Day 31: Checkpointing & inference mode
Saving, loading, and serving a trained model
A trained model is worthless if it evaporates when the process exits. Checkpointing saves the model's learned parameters (the state_dict) to disk so you can reload it for more training or for inference. The convention is to save the state_dict (just the weights) rather than the whole model object — it's portable and version-robust.
# save
torch.save(model.state_dict(), "garment_classifier.pt")
# load for inference
model = GarmentClassifier(num_classes=8) # recreate the architecture
model.load_state_dict(torch.load("garment_classifier.pt"))
model.eval() # inference mode
# a full checkpoint also saves optimizer state for resuming training
torch.save({
"epoch": epoch,
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
}, "checkpoint.pt")Two kinds of checkpoint, two purposes
For *deployment*, save just the model weights. For *resuming training*, also save the optimizer state (Adam's momentum buffers) and the epoch — otherwise resuming restarts the optimizer cold and the loss jumps. Knowing the difference signals you've actually run long training jobs, not just tutorials.
Key terms
- state_dict
- A dictionary mapping each layer to its learned parameter tensors — PyTorch's portable representation of a model's weights.
- Checkpoint
- A saved snapshot of training state (weights, optimizer state, epoch) allowing inference or resumed training.
- Inference mode
- Running a trained model to make predictions, with eval() set and gradients disabled.
You want to pause a long training run and resume it later exactly where it left off. Saving only model.state_dict() is insufficient. What else must you save?