Checkpoints and Early Stopping

Three days into a training run, at ninety percent complete, the power goes out. Every deep learning engineer has a story like this one. And it only happens once, because after that, you never train anything without saving as you go.

• Day three of training
• The power goes out
• It only happens once

Knowing when to stop

Look at these two curves. Training loss keeps falling, forever. But validation loss falls, and then it turns and starts to rise. That's overfitting, in action. The best model is at the bottom of the orange curve, not at the end of training. Early stopping watches this curve, and it stops the run when things stop improving.

An early stopper

Here's a small class that does this. Each epoch, it receives the validation loss. If that's a new best, it saves a copy of the model, and resets its counter. If not, it adds to the counter. And after too many bad epochs in a row, it says stop. We feed it a pretend history. Run it. It stops at epoch six, and keeps the model from epoch four. We get the best version, not just the last one.

import copy

class EarlyStopper:
    def __init__(self, patience=3):
        self.patience, self.best, self.bad = patience, float("inf"), 0
        self.state = None

    def step(self, val_loss, model_state):
        if val_loss < self.best:
            self.best, self.bad = val_loss, 0
            self.state = copy.deepcopy(model_state)
            return False
        self.bad += 1
        return self.bad >= self.patience

stopper = EarlyStopper(patience=2)
history = [0.9, 0.7, 0.55, 0.48, 0.5, 0.53]

for epoch, loss in enumerate(history, 1):
    stop = stopper.step(loss, {"epoch": epoch})
    print(f"epoch {epoch}  val loss {loss}  stop? {stop}")
    if stop:
        break
print("Best model came from:", stopper.state)

epoch 1  val loss 0.9  stop? False
epoch 2  val loss 0.7  stop? False
epoch 3  val loss 0.55  stop? False
epoch 4  val loss 0.48  stop? False
epoch 5  val loss 0.5  stop? False
epoch 6  val loss 0.53  stop? True
Best model came from: {'epoch': 4}

Saving a checkpoint

A checkpoint needs more than the weights. To resume training properly, you also need the optimiser's state, and the epoch number. State dict gives you a dictionary of all weights. Torch save writes it all to one file. To resume, we load it, and pour the weights into a fresh model. Run it. The weights are identical, and we can carry on from epoch seven. Save a checkpoint at each epoch, and keep the best one.

import torch
from torch import nn

model = nn.Linear(4, 2)
opt = torch.optim.AdamW(model.parameters())

torch.save({"epoch": 7, "model": model.state_dict(),
            "optimiser": opt.state_dict()}, "checkpoint.pt")

saved = torch.load("checkpoint.pt", weights_only=True)
fresh = nn.Linear(4, 2)
fresh.load_state_dict(saved["model"])

same = torch.equal(model.weight, fresh.weight)
print("Resuming from epoch:", saved["epoch"])
print("Weights identical:", same)

Resuming from epoch: 7
Weights identical: True

Making runs repeatable

One last habit. Neural networks use randomness everywhere, in initial weights, in shuffling, and in dropout. If two runs differ, you can't tell whether your change helped. Setting the seeds at the start makes results repeatable. Run it. The same seed gives the same numbers. Even then, runs on a G P U can differ slightly, so for serious comparisons, average over a few seeds.

import random
import numpy as np
import torch

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

set_seed(42)
a = torch.rand(3)
set_seed(42)
b = torch.rand(3)
print("Same numbers:", torch.equal(a, b))

Same numbers: True

Summary