Unit 03.05: Saving metrics and checkpoints

A model you cannot reload is a model you cannot report. Checkpointing takes three lines and turns a run into something you can defend.

Save on validation, not on training

Training loss falls almost monotonically, so checkpointing on it saves the most overfitted model you produced. Validation loss is the signal that reflects generalisation, so that is what should trigger a save.

Keep the history too — the per-epoch record is what lets you draw a learning curve later:

import json
import tempfile
from pathlib import Path

import torch
from torch import nn

torch.manual_seed(0)
X = torch.randn(200, 2)
y = (2 * X[:, 0] - 3 * X[:, 1] + 1).unsqueeze(1)
split = 160
Xtr, ytr, Xva, yva = X[:split], y[:split], X[split:], y[split:]

model = nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
loss_fn = nn.MSELoss()

outdir = Path(tempfile.mkdtemp())
history, best = [], float("inf")

for epoch in range(100):
    optimizer.zero_grad()
    loss_fn(model(Xtr), ytr).backward()
    optimizer.step()

    with torch.no_grad():          # validation never needs gradients
        train_loss = loss_fn(model(Xtr), ytr).item()
        val_loss = loss_fn(model(Xva), yva).item()
    history.append({"epoch": epoch, "train": train_loss, "val": val_loss})

    # Checkpoint on the validation score, not the training score.
    if val_loss < best:
        best = val_loss
        torch.save({"epoch": epoch, "state_dict": model.state_dict()},
                   outdir / "best.pt")

(outdir / "history.json").write_text(json.dumps(history, indent=1))

restored = nn.Linear(2, 1)
checkpoint = torch.load(outdir / "best.pt", weights_only=True)
restored.load_state_dict(checkpoint["state_dict"])

print("epochs recorded :", len(history))
print("best val loss   :", round(best, 6), "at epoch", checkpoint["epoch"])
with torch.no_grad():
    print("restored matches:", round(loss_fn(restored(Xva), yva).item(), 6))
# A run you cannot reload is a run you cannot report.

The restored model reproduces the validation loss exactly, which confirms the checkpoint captured everything needed. Note that torch.load uses weights_only=True: loading a checkpoint executes whatever is inside it, and a file from an untrusted source can run arbitrary code.

The mistake this prevents

Saving only the final model. If validation loss bottomed out at epoch 40 and you trained to 100, the file on disk is the worse model — and the good one is gone.

Takeaway

Checkpoint on validation, keep the epoch it came from, and record the history alongside it.

Mark complete