Checkpoints — Module 6.2

How This Lesson Fits the Module

Training runs for hours or days. Checkpoints snapshot model weights, optimizer state, and epoch counters so you can resume after crashes, compare runs, and load the validation-best model for testing.

Without checkpoints, a GPU preemption at epoch 89 of 100 means starting over—unacceptable in production training jobs.

Prior Lesson The test loop loads a saved checkpoint. This lesson covers what to save and how to restore.

Learning Objectives

By the end of this lesson, students should be able to:

Save and load model.state_dict() and optimizer.state_dict().
Build a checkpoint dict with epoch, metrics, and hyperparameters.
Resume training from a checkpoint without losing optimizer momentum.
Distinguish saving full model vs weights-only checkpoints.
Implement “save best validation” and periodic checkpoint policies.
Use map_location when loading on a different device.

What Belongs in a Checkpoint

At minimum: model weights. For resume training: optimizer state, epoch, scheduler state, random seeds, and the validation metric that triggered the save.

Field	Needed for inference	Needed to resume training
`model_state_dict`	Yes	Yes
`optimizer_state_dict`	No	Yes
`scheduler_state_dict`	No	Yes (if using scheduler)
`epoch`, `best_val_loss`	No	Recommended
Hyperparameters / config	Helpful	Yes

Saving and Loading Checkpoints

Prefer state_dict over pickling entire model objects—portable across code versions if architecture class is unchanged.

def save_checkpoint(path, model, optimizer, epoch, val_loss, scheduler=None):
    checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "val_loss": val_loss,
        "config": {"lr": 1e-3, "batch_size": 64},
    }
    if scheduler is not None:
        checkpoint["scheduler_state_dict"] = scheduler.state_dict()
    torch.save(checkpoint, path)

def load_checkpoint(path, model, optimizer=None, scheduler=None, device="cpu"):
    ckpt = torch.load(path, map_location=device)
    model.load_state_dict(ckpt["model_state_dict"])
    if optimizer and "optimizer_state_dict" in ckpt:
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    if scheduler and "scheduler_state_dict" in ckpt:
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
    return ckpt.get("epoch", 0), ckpt.get("val_loss", float("inf"))

Best-vs-Last Checkpoint Policy

Maintain two files: last.pt every epoch (crash recovery) and best.pt when validation improves (deployment and test evaluation). Delete or rotate old checkpoints on disk-constrained clusters.

best_val = float("inf")
for epoch in range(start_epoch, num_epochs + 1):
    train_one_epoch(...)
    val_loss, _ = validate(...)

    save_checkpoint(f"checkpoints/last.pt", model, optimizer, epoch, val_loss)

    if val_loss < best_val:
        best_val = val_loss
        save_checkpoint(f"checkpoints/best.pt", model, optimizer, epoch, val_loss)
        print(f"New best val_loss={best_val:.4f} at epoch {epoch}")

Critical Mistake — Saving Only state_dict to Wrong Architecture

load_state_dict requires identical layer names and shapes. Changing num_classes or depth after saving yields size-mismatch errors. Version your model config alongside the checkpoint.

Resuming After Interruption

Load last.pt, restore optimizer and epoch, continue the loop from start_epoch = ckpt_epoch + 1. Adam’s momentum buffers live in optimizer state—without them, resumed training behaves like a cold restart.

Engineering Habit — Atomic Writes

Write to best.pt.tmp then rename to best.pt so a crash mid-write does not corrupt the only good checkpoint.

Knowledge Check

Short Answer: What method exports model weights? Answer: model.state_dict().
True/False: Optimizer state is needed for inference deployment. Answer: False—only weights matter at inference.
Multiple Choice: map_location is used when: (a) loading on CPU from GPU save, (b) speeding training, (c) clipping gradients. Answer: (a).
Short Answer: Why save both best.pt and last.pt? Answer: Best for quality; last for exact resume point.
Short Answer: What is lost if you save model but not optimizer state? Answer: Momentum/adaptive estimates when resuming.
True/False: Pickling the entire nn.Module is more portable than state_dict. Answer: False—state_dict is preferred.
Multiple Choice: Test evaluation should load: (a) best val checkpoint, (b) random weights, (c) optimizer only. Answer: (a).
Short Answer: What metadata helps reproduce a run? Answer: Config, epoch, metric, library versions.
Short Answer: When does load_state_dict fail? Answer: Architecture mismatch with saved keys/shapes.
Multiple Choice: Checkpoint every epoch on a 1 TB model is often: (a) always fine, (b) impractical without rotation, (c) impossible. Answer: (b).

Key Takeaways

Checkpoints bundle weights and, for resume, optimizer/scheduler state plus metadata.
Save best for deployment/test and last for crash recovery.
Use state_dict + config versioning for portable saves.
Next: Early Stopping—stop training when validation stops improving.

Trainer’s Guide

Hands-on idea: Train 5 epochs, kill the job, resume from last.pt, confirm loss curve continuity.

Discussion prompt: How do cloud platforms (SageMaker, Vertex) handle checkpoint paths and sync to object storage?

What’s Next Automate “save best” with stopping logic in Early Stopping.