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
| Situation | Risk Level |
|---|---|
| Loading a model file you trained and saved yourself | Low โ you control the source |
| Loading a model file from a trusted internal team, verified source | Low to moderate โ depends on your organization's supply chain trust |
| Loading a model file downloaded from an arbitrary public URL or untrusted third party | High โ treat as a genuine security risk, prefer weights_only=True or the safetensors format instead |
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?