02 · Checkpoints & Resuming ⏸️

4 min read

Saving weights lets you reload a finished model. A checkpoint lets you pause a half-trained one and pick up exactly where you left off.

1. The 'Why'

Real training runs are long — hours or days — and the world is unreliable. A machine reboots, a cloud instance is preempted, a run crashes at epoch 47 of 100. If all you saved was the model's state_dict, you'd lose more than the weights: you'd lose the optimizer's internal state (Adam's per-parameter momentum estimates, for instance), the current epoch number, and any learning-rate schedule progress. Resuming from a bare state_dict restarts the optimizer cold, which can jolt training and waste the momentum it had built up. A checkpoint solves this by bundling everything needed to continue into one saved dictionary.

The idea is simple and powerful: because torch.save can serialize any dictionary, you save a dict containing the model's state_dict, the optimizer's state_dict, the epoch you stopped at, and whatever else you need (best validation score, loss history, RNG state). To resume, you load that dict and restore each piece into place, then continue your training loop from the saved epoch. This is also the mechanism behind saving the best model during training (sub-module 04) and is standard practice in every serious project. Master checkpoints and long training runs stop being fragile.

2. Core Concepts

A checkpoint is just a dictionary. You assemble a plain Python dict of everything you want to persist and hand it to torch.save. There's no special checkpoint object.

What to include. At minimum: model.state_dict() and optimizer.state_dict(). Usually also the epoch, the best metric so far, and often the loss/metric history. Optionally the LR scheduler's state_dict.

Why the optimizer state matters. Optimizers like Adam maintain running statistics per parameter. Dropping them on resume effectively restarts the optimizer, causing a visible bump in the loss. Restoring them makes resuming seamless.

Restoring on resume. Re-create the model and optimizer objects, then call load_state_dict on each from the saved checkpoint, and read back the scalar values (epoch, best score) to continue the loop.

3. Code in Action

Saving a checkpoint

import torch

def save_checkpoint(path, model, optimizer, epoch, best_val):
    checkpoint = {
        "epoch": epoch,                          # where we stopped
        "model_state": model.state_dict(),       # learned weights
        "optimizer_state": optimizer.state_dict(),# momentum, step counts, etc.
        "best_val": best_val,                    # best validation score so far
    }
    torch.save(checkpoint, path)                 # one file holds it all

# ...inside training, periodically:
# save_checkpoint("ckpt.pt", model, optimizer, epoch, best_val)

Loading a checkpoint to resume

import torch
import torch.nn as nn

# 1. Re-create the SAME model and optimizer objects
model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 3))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# 2. Load the checkpoint dict and restore each component
ckpt = torch.load("ckpt.pt")
model.load_state_dict(ckpt["model_state"])          # restore weights
optimizer.load_state_dict(ckpt["optimizer_state"])  # restore optimizer state
start_epoch = ckpt["epoch"] + 1                      # continue from the NEXT epoch
best_val = ckpt["best_val"]
print(f"Resuming from epoch {start_epoch}")

Resuming the training loop

for epoch in range(start_epoch, 100):     # note: starts at the resumed epoch
    model.train()
    for X, y in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(X), y)
        loss.backward()
        optimizer.step()
    # periodically checkpoint again so the next crash is cheap
    save_checkpoint("ckpt.pt", model, optimizer, epoch, best_val)

4. Common Pitfalls

Pitfall 1 — Saving only the model, then wondering why resume is bumpy. Without optimizer.state_dict(), Adam/SGD-momentum restart from scratch and the loss jumps on resume. Always checkpoint the optimizer alongside the model.

Pitfall 2 — Off-by-one on the resume epoch. If you saved at the end of epoch N, resume at N + 1, not N, or you'll silently repeat an epoch. Store the completed epoch and add one when restoring.

Pitfall 3 — Restoring into mismatched objects. optimizer.load_state_dict must go into an optimizer built over the same parameters with the same type; loading Adam state into an SGD optimizer (or after changing the model) errors or corrupts state. Re-create model and optimizer identically before restoring.

5. Further Reading & Watch List

⬅️ Prev: 01 · state_dict Basics · ➡️ Next: 03 · Loading for Inference