Module 8 — Checkpoints and resuming training
The Fashion-MNIST classifier from module 7 trains in a few minutes on a GPU. A real deep network trains for hours, and often days. Anything that runs longer than a coffee break must survive a lost SSH session, a preempted cloud instance, a driver crash or a colleague who kicks the power strip. Checkpointing is that survival mechanism. Getting it right is more subtle than saving the weights: the weights are only one piece of what a training run needs to resume exactly where it left off.
What actually needs to be saved
Naive checkpointing saves model.state_dict() and calls it done. That is enough to reload a trained model for inference. It is not enough to resume training.
To resume identically, five pieces of state matter:
| State | Held by | Why it matters |
|---|---|---|
| Model weights | model.state_dict() | the learned parameters and buffers |
| Optimiser state | optimizer.state_dict() | momentum, Adam moments, decoupled state per parameter |
| Scheduler state | scheduler.state_dict() | how far along the learning-rate curve we are |
| Epoch counter | your loop | so we know where to resume |
| RNG state | torch.get_rng_state(), and numpy/python if used | so shuffling and dropout continue coherently |
The optimiser state is the most frequently forgotten. Adam maintains first- and second-moment estimates per parameter; discarding them restarts the adaptive learning rate from scratch, and the first few post-resume batches produce updates that look nothing like a natural continuation. Training does not diverge — Adam recovers within a few dozen batches — but the loss curve has a visible bump exactly where the resume happened, which is a reliable tell in code reviews.
The canonical save
def save_checkpoint(path, model, optimizer, scheduler, epoch, best_val):
torch.save({
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"epoch": epoch,
"best_val": best_val,
"torch_rng": torch.get_rng_state(),
"cuda_rng": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
}, path)
torch.save serialises whatever object you give it — dict, list, model, tensor — with pickle. A checkpoint is just a dictionary, which is exactly the shape that survives well when the model definition changes later. Adding a new key next month does not break old checkpoints; renaming an existing key does.
The canonical load
def load_checkpoint(path, model, optimizer, scheduler, device):
ckpt = torch.load(path, map_location=device, weights_only=False)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
scheduler.load_state_dict(ckpt["scheduler"])
torch.set_rng_state(ckpt["torch_rng"])
if ckpt["cuda_rng"] is not None and torch.cuda.is_available():
torch.cuda.set_rng_state_all(ckpt["cuda_rng"])
return ckpt["epoch"], ckpt["best_val"]
map_location=device is essential. Without it, a checkpoint saved on a GPU tries to reload on that same GPU by default. If you saved on cuda:0 and reload on a machine with only a CPU — the typical laptop after cloud training — the load fails. map_location="cpu" reloads onto CPU regardless of where the checkpoint originated; map_location=torch.device("cuda") reloads onto whichever GPU is current. This one argument is what makes checkpoints portable.
weights_only=False is required when the checkpoint contains a plain dict with tensors, integers and RNG state — anything beyond pure weights. Since PyTorch 2.6, weights_only=True is the default for security, and non-tensor payloads must be opted in explicitly.
The training loop, extended
start_epoch = 0
best_val = float("inf")
if Path("checkpoint.pt").exists():
start_epoch, best_val = load_checkpoint("checkpoint.pt", model, optimizer, scheduler, device)
start_epoch += 1
for epoch in range(start_epoch, total_epochs):
train_one_epoch(...)
val_loss = evaluate(...)
save_checkpoint("checkpoint.pt", model, optimizer, scheduler, epoch, best_val)
if val_loss < best_val:
best_val = val_loss
save_checkpoint("best.pt", model, optimizer, scheduler, epoch, best_val)
Two files: checkpoint.pt for the latest epoch, best.pt for the best validation. The distinction between "last" and "best" matters and is the subject of the next section.
Best epoch versus last epoch: two artifacts, two purposes
The best-validation checkpoint is what you deploy. The last-epoch checkpoint is what you resume from. Conflating them causes two symmetric bugs.
If you resume from best.pt, subsequent epochs restart from a snapshot that already generalised well, but the optimiser state is that of the best epoch, not the current one. If several epochs of overfitting happened after it, resuming loses that recent history, and the loss trace shows a discontinuity where the resume happened.
If you deploy checkpoint.pt, you deploy the last epoch — which may be the worst-generalising one, since training tends to overfit near the end. The safe production artifact is best.pt, chosen by an explicit metric on a validation set, and never overwritten by the next epoch.
Unlike some Keras callbacks, PyTorch does not restore best weights at the end of training on its own. If you want the deployed model to be the best one on validation, save it explicitly during training and reload it before serving. Assuming that model at the end of the loop is the best model is a specific class of production bug.
Determinism, or the illusion of it
Full bit-for-bit determinism across resumes is possible but expensive. It requires:
torch.use_deterministic_algorithms(True)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
plus setting every seed in every library that generates random numbers — Python's random, NumPy, torch, CUDA — and saving all their RNG states in the checkpoint. It also slows training because some fast non-deterministic algorithms get disabled. For most projects, "resumes cleanly with a small numerical difference" is a better trade than "runs identically but takes 25 % longer". Reserve strict determinism for scientific reproducibility, and accept a small drift in exchange for speed on production runs.
Version compatibility
A checkpoint saved with PyTorch 2.6 is not guaranteed to load with PyTorch 2.4. The change that most often breaks compatibility is a rename or restructure of a state_dict key — a layer added or removed in the model definition, a submodule renamed. load_state_dict(..., strict=False) skips missing keys and reports them, which is what you want for fine-tuning but never for resume.
The safest habit: pin the PyTorch version alongside the training code, and save that pin in the checkpoint metadata:
"pytorch_version": torch.__version__,
"python_version": sys.version,
On load, comparing versions and warning on mismatch prevents debugging the wrong problem for an hour.
Saving every epoch is standard for training of a few hours. Saving every batch is wasteful — the disk I/O eats into training time, and 99 % of those checkpoints will be deleted. A reasonable heuristic is one save per epoch plus one save after any "best validation so far", which is exactly the two-file scheme above.
In summary
- A checkpoint contains five things, not one: model weights, optimiser state, scheduler state, epoch counter and RNG state; skipping the optimiser produces a visible bump in the loss trace at resume.
map_locationat load makes checkpoints portable across devices; saving on GPU and loading on CPU without it is a common blocking failure.- Maintain two files, one for the latest epoch and one for the best validation; do not deploy the last-epoch file unless you have verified it is also the best.
- Strict determinism is expensive; near-determinism with a saved RNG is the right default, and pinning the PyTorch version alongside the checkpoint prevents mysterious load failures.
Next module: transfer learning with torchvision, where a pretrained ResNet18 replaces our MLP and unlocks another twenty points of accuracy.