Checkpoints: saving training to resume it

A checkpoint holds more than the weights

The weights alone are enough to use a finished model. To continue training, you need everything that the next step depends on:

• The model's weights.
• The optimizer state: the numbers the optimizer keeps between steps. For example, AdamW keeps two running averages for every weight. Without them, the first steps after a resume behave differently.
• The step count, so that the learning-rate schedule continues from the right place and the run knows when to stop.
• Any extra copies, such as the averaged weights.
• The state of the random-number generators, so that the batches and random choices that follow are the ones the run would have made (see random draws).

A full resume continues exactly where training stopped

The snippet runs the same training in four parts. Run A takes 20 steps without stopping. Run B takes 10 steps and saves a checkpoint. Run C loads that checkpoint into new objects and takes the other 10 steps. Run D loads only the weights and starts a fresh optimizer.

resume.py

import torch

def make():  # a linear layer (3 inputs, 1 output), with AdamW for its weights
    model = torch.nn.Linear(3, 1)
    return model, torch.optim.AdamW(model.parameters(), lr=0.01)

def train(model, opt, steps):
    for _ in range(steps):                        # _ is a name the loop never uses
        x = torch.randn(8, 3)                     # a random minibatch
        # The loss is the mean squared gap between the output and each row's sum.
        # keepdim=True keeps the sums in a column, the shape of model(x).
        loss = (model(x) - x.sum(1, keepdim=True)).pow(2).mean()
        # Clear the old gradients, compute new ones, and take one step.
        opt.zero_grad(); loss.backward(); opt.step()

torch.manual_seed(0)                              # run A: 20 steps without stopping
a, opt_a = make()
train(a, opt_a, 20)

torch.manual_seed(0)                              # run B: 10 steps, save, stop
b, opt_b = make()
train(b, opt_b, 10)
torch.save({"model": b.state_dict(), "optimizer": opt_b.state_dict(),
            "step": 10, "rng": torch.get_rng_state()}, "ckpt.pt")

c, opt_c = make()                                 # later: new objects, then resume
ck = torch.load("ckpt.pt")
c.load_state_dict(ck["model"]); opt_c.load_state_dict(ck["optimizer"])
torch.set_rng_state(ck["rng"])
train(c, opt_c, 20 - ck["step"])

d, opt_d = make()                                 # resume, but forget the optimizer
d.load_state_dict(ck["model"])
torch.set_rng_state(ck["rng"])
train(d, opt_d, 20 - ck["step"])

# torch.equal is True only when two tensors hold exactly the same numbers.
print("full resume matches run A:   ", torch.equal(a.weight, c.weight))
print("without optimizer state:     ", torch.equal(a.weight, d.weight))

full resume matches run A:    True
without optimizer state:      False

With everything saved, the resumed run ends with exactly the same weights as the run that never stopped. Without the optimizer state, it ends somewhere else. A state_dict() is PyTorch's name for all of an object's numbers, gathered into a dictionary that torch.save can write. torch.manual_seed(0) starts PyTorch's random numbers from the seed 0, so runs A and B start from the same weights and draw the same minibatches. torch.get_rng_state() reads where PyTorch's generator is in its sequence, and torch.set_rng_state puts it back there (see random draws).

Save often enough to lose little, rarely enough to stay fast

Saving takes time, during which training pauses. Saving rarely risks losing more work when the job stops. The interval between saves trades one against the other.

Try it. The figure repeats the worked example with an interval you choose. Find an interval that keeps the worst-case loss under 15 minutes and the time spent saving under 3%.

Write each checkpoint safely

• Write to a temporary file, then rename it to the real name. A rename happens all at once, so a stop during the save leaves the previous checkpoint whole.
• Keep the last few checkpoints, not only the latest, in case one is damaged.
• Keep the checkpoint with the lowest validation loss as well; that is often the model you want in the end.