Model Saving and Loading

This note covers the practical mechanics of saving and loading PyTorch models correctly — the difference between saving just the weights versus the whole model object, and the specific mistakes that most commonly break a saved model when loading it back.

Two Ways to Save a Model

Saving the state_dict is the standard, recommended practice — it decouples the saved weights from the exact code structure, which matters enormously for long-term reproducibility and sharing models across different codebases or versions.

Code — Saving and Loading a state_dict

import torch

# Saving
torch.save(model.state_dict(), 'model_weights.pt')

# Loading -- you must recreate the SAME model architecture first
model = MyModelClass()               # the architecture must match exactly
model.load_state_dict(torch.load('model_weights.pt'))
model.eval()                           # switch to evaluation mode before inference

Code — Saving and Loading the Whole Model (Less Recommended)

torch.save(model, 'full_model.pt')     # saves the class definition + weights together

model = torch.load('full_model.pt')    # no need to redefine the class... but fragile
model.eval()

The Architecture-Mismatch Trap

When loading a state_dict, PyTorch matches saved parameter names to the current model's parameter names exactly. If you change the model's architecture (add a layer, rename a variable, change a layer's size) between saving and loading, load_state_dict will raise an error about mismatched keys or shapes — this is a genuinely common source of confusion, especially after refactoring model code.

Code — Loading Partial Weights (e.g. for Transfer Learning)

pretrained_dict = torch.load('pretrained_weights.pt')
model_dict = model.state_dict()

# Only keep weights whose keys AND shapes match the current model
matched_dict = {k: v for k, v in pretrained_dict.items()
                 if k in model_dict and v.shape == model_dict[k].shape}
model_dict.update(matched_dict)
model.load_state_dict(model_dict)   # loads what matches, keeps the rest at its current initialization

This partial-loading pattern is exactly what makes transfer learning practical — loading a pretrained backbone's weights while leaving a newly added, differently-sized output layer at its fresh initialization (covered fully in the Transfer Learning category).

Common Mistakes

• Forgetting model.eval() after loading a model intended for inference — dropout and batch norm will still behave in training mode, producing inconsistent, incorrect predictions.
• Loading a state_dict into a model with a different architecture without using the partial-loading pattern above — this crashes with a key/shape mismatch error rather than silently doing the wrong thing, which is at least easy to catch (but confusing the first time you see it).
• Saving the whole model object via pickle for long-term storage or sharing across teams — this can break if the class definition changes even slightly, unlike the more portable state_dict approach.

Interview Relevance

Q: "Why is saving a model's state_dict generally preferred over saving the entire model object?" A state_dict only contains the learned parameter values, decoupled from the exact class definition and code structure — it's portable across minor code refactors and different environments. Saving the whole object (via Python's pickle, which is what torch.save(model, ...) uses) ties the saved file to the exact class definition at save time, which can break if that code changes at all before the file is loaded again.

Practice Question

You want to fine-tune a pretrained image classifier on a new task with a different number of output classes. What loading strategy would you use to keep the pretrained backbone's weights while letting the new output layer train from scratch?

Related DL Notes

• Checkpointing
• PyTorch Transfer Learning
• Model Serialization