Checkpointing means periodically saving a model's weights (and optimizer state) to disk during training — protecting hours or days of compute from being lost to a crash, and preserving the model's best-performing version even if later epochs make things worse.
What Actually Gets Saved
| Component | Why Save It |
|---|---|
Model weights (state_dict) | The core thing you need — the learned parameters |
| Optimizer state | Adam and similar optimizers maintain per-parameter running statistics (\(m_t\), \(v_t\) from Adam Optimizer) — without saving these, resuming training "cold" can temporarily destabilize training |
| Current epoch number | So training can resume exactly where it left off, not restart from epoch 0 |
| Learning rate scheduler state | So the schedule continues correctly rather than restarting |
Code — Saving a Full Checkpoint
import torch
checkpoint = {
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'val_loss': avg_val_loss,
}
torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pt')Code — Resuming From a Checkpoint
checkpoint = torch.load('checkpoint_epoch_10.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1
for epoch in range(start_epoch, num_epochs):
# ... training continues exactly where it left off ...
passTwo Common Checkpointing Strategies
| Strategy | What It Saves | Use Case |
|---|---|---|
| Periodic checkpointing | Every N epochs, regardless of performance | Crash recovery for long training runs |
| Best-model checkpointing | Only when validation performance improves on the best seen so far | Ensures you always have the best-performing version saved, even if later epochs overfit and get worse |
Best-model checkpointing directly complements Early Stopping — even if you keep training past the point where validation performance peaks (to see if it improves further), the checkpoint from the best epoch is preserved and can be reloaded as the final model.
Code — Best-Model Checkpointing
best_val_loss = float('inf')
for epoch in range(num_epochs):
# ... training and validation for this epoch ...
if avg_val_loss < best_val_loss:
best_val_loss = avg_val_loss
torch.save(model.state_dict(), 'best_model.pt')
print(f"New best model saved at epoch {epoch}, val_loss={avg_val_loss:.4f}")Common Mistakes
- Saving only the model weights and forgetting the optimizer state when planning to resume a long training run — resuming without optimizer state effectively restarts Adam's momentum/variance estimates from scratch, which can cause a temporary instability right after resuming.
- Overwriting a single checkpoint file every epoch without keeping the best-performing one separately — if training later degrades (e.g. due to overfitting or a bad hyperparameter change), the best version has already been lost.
Interview Relevance
Q: "Why save the optimizer's state, not just the model's weights, when checkpointing for a long training run?" Optimizers like Adam maintain per-parameter running statistics (moving averages of gradients and their squares) that took many steps to build up — restarting training with fresh, zero-initialized optimizer state (even with the correct model weights) means the optimizer briefly behaves as if training just began, which can cause a temporary but real disruption to training stability right after resuming.
Practice Question
A training run needs to survive an unreliable server that occasionally reboots. Would periodic checkpointing or best-model-only checkpointing better serve this specific goal, and why might you want both simultaneously?
