Pickle Models

Pickle is Python's built-in object serialization format, and it underlies torch.save() internally — understanding what it actually does, and its real security risks, matters for safe deployment.

What Pickle Does

Pickle converts (nearly) any Python object into a byte stream that can be written to disk and later reconstructed back into a live Python object — including custom classes, provided their exact class definition is importable at load time. torch.save()/torch.load() use pickle under the hood by default for the surrounding container structure, even when the tensors themselves are stored more efficiently.

Code — Direct Pickle Usage

import pickle

# Saving any Python object, e.g. a scikit-learn model or a config dictionary
with open("model.pkl", "wb") as f:
    pickle.dump(sklearn_model, f)

# Loading it back
with open("model.pkl", "rb") as f:
    loaded_model = pickle.load(f)

The Real Security Risk — Never Unpickle Untrusted Data

This is worth stating plainly and directly, since it's a genuine, serious production risk: unpickling data from an untrusted or unverified source can execute arbitrary code. Pickle's format allows encoding instructions that run during deserialization, not just passive data — a maliciously crafted pickle file can compromise the system that loads it. This isn't a theoretical concern; it's a well-documented, real-world attack vector.

# NEVER do this with a model file from an untrusted source
with open("model_downloaded_from_random_url.pkl", "rb") as f:
    model = pickle.load(f)   # could execute arbitrary malicious code

# Safer alternatives for untrusted or cross-team model sharing:
# - torch.load(..., weights_only=True)  -- restricts loading to tensor data only
# - safetensors format -- a format designed specifically to store only tensor
#   data, with no ability to execute arbitrary code during loading

When Pickle Is Fine vs When It Isn't

Common Mistakes

• Loading pickled model files from untrusted or unverified sources without using safer, restricted-loading alternatives — a genuine, exploitable security vulnerability, not just a theoretical concern.
• Assuming pickle files are just "data" with no execution risk — pickle's format inherently supports arbitrary code execution during deserialization, which is precisely what makes untrusted pickle files dangerous.

Interview Relevance

Q: "Why is loading a pickled model file from an untrusted source considered a security risk, not just a data format concern?" Python's pickle format can encode instructions that execute during deserialization, not just passive data values — a maliciously crafted pickle file can run arbitrary code on the machine that loads it. This makes unpickling untrusted data a genuine attack vector, which is why safer alternatives exist for handling models from unverified sources: torch.load(..., weights_only=True), which restricts loading to tensor data only, or dedicated formats like safetensors designed specifically to exclude any code-execution capability.

Practice Question

A colleague suggests downloading a pretrained model checkpoint from a random file-sharing link and loading it directly with pickle.load(). What would you recommend instead, and why?

Related DL Notes

• Model Serialization
• Saving PyTorch Models
• DL Model Saving