# Saving and loading with state_dict — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-saving

> Save and restore model and optimizer state with state_dict, write resumable checkpoints, and load GPU-trained weights on a CPU.

## Save the weights, not the object

A module's `state_dict()` is a dictionary mapping names such as `fc1.weight` to tensors: every parameter plus buffers like BatchNorm's running averages. The recommended way to save a model is `torch.save(model.state_dict(), "model.pt")`. To load it, first construct the model class in code, then call `model.load_state_dict(torch.load("model.pt"))`.

You *can* `torch.save(model)` the whole object, but that pickles a reference to your class and module path; renaming a file or class later breaks loading. Saving only the state dict keeps weights independent of code layout.

To **resume** training, save a checkpoint dictionary: the model's state dict, the optimizer's state dict (Adam's running averages are part of training state), the scheduler's, and the epoch number. Pass `map_location="cpu"` (or your device) to `torch.load` so weights saved on a GPU load on a machine without one. After loading for inference, call `model.eval()`. `load_state_dict` is strict by default and reports missing or unexpected keys, which usually means the architecture changed.

## From trained model to service

A trained model is saved as a state_dict checkpoint, loaded and verified on the target machine, optionally exported to a portable format such as torch.export or ONNX, and finally served behind an application.

![Diagram of 5 connected stages from left to right: trained model, state_dict checkpoint, load and verify, export (torch.export / ONNX), serve.](assets/figures/pytorch/section-9-map.svg) — Figure 9.1: trained model -> state_dict checkpoint -> load and verify -> export (torch.export / ONNX) -> serve

## A resumable checkpoint

`SmallCNN`, `model`, `optimizer` and `epoch` come from the training loop. `weights_only=True` restricts unpickling to tensors and plain containers, which is all a checkpoint like this needs.

```python
import torch

checkpoint = {
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, "checkpoint.pt")

# Later, possibly on another machine:
model = SmallCNN()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
ckpt = torch.load("checkpoint.pt", map_location="cpu", weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1
```

## torch.load runs pickle: treat files as code

`torch.load` uses Python's pickle format, and unpickling an untrusted file can execute arbitrary code. Recent PyTorch versions default to `weights_only=True`, which only allows tensors and simple types; never switch it to `False` for files you did not create. For sharing weights publicly, the `safetensors` format stores tensors with no executable content.

**Quiz:** What is the recommended way to save a trained model for later use?

- [ ] `torch.save(model)` to pickle the whole object
- [x] `torch.save(model.state_dict(), path)`, then rebuild the class and call `load_state_dict`
- [ ] Print the weights and copy them into a text file
- [ ] Save only the optimizer's state dict

*Answer:* `torch.save(model.state_dict(), path)`, then rebuild the class and call `load_state_dict`. A state dict holds only tensors keyed by name, so it survives refactoring. Pickling the whole model ties the file to your exact class and module path.