Building on the general concepts from Model Saving and Loading, this note gives the exact, practical PyTorch syntax for saving models correctly, including full training checkpoints.
Saving Just the Weights (Recommended)
import torch
torch.save(model.state_dict(), "model_weights.pt")Saving a Full Training Checkpoint
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'best_val_loss': best_val_loss,
}
torch.save(checkpoint, f"checkpoint_epoch_{epoch}.pt")This is the exact pattern from Checkpointing — saving not just the weights, but everything needed to resume training exactly where it left off.
Saving the Best Model During Training
best_val_loss = float('inf')
for epoch in range(num_epochs):
# ... training and validation for this epoch ...
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save(model.state_dict(), "best_model.pt")
print(f"New best model saved at epoch {epoch}")Common Mistakes
- Overwriting the same checkpoint file every epoch without a separate "best model" save — if performance degrades in later epochs (e.g. due to overfitting), the best-performing version has already been lost.
- Saving to a path that doesn't get backed up or persisted (e.g. inside a temporary/ephemeral cloud instance's local disk) — for long training runs, saving checkpoints to persistent storage is essential to avoid losing progress.
Interview Relevance
Q: "What should a complete training checkpoint include, beyond just the model's weights?" The optimizer's state (essential for optimizers like Adam that maintain per-parameter running statistics — resuming without it effectively restarts those statistics from scratch), the learning rate scheduler's state (so the schedule continues correctly), the current epoch number, and often the best validation metric seen so far — everything needed to resume training exactly as if it had never stopped.
Practice Question
Why is it good practice to save a separate "best model" checkpoint, distinct from a periodic checkpoint saved every N epochs?
