Skip to main content

Checkpoints and a reusable fit function

Save a run’s whole state with torch.save, keep the weights with the best validation loss, resume a stopped run so that it continues exactly as if it had never stopped, and package it all in a fit function that the next two modules use unchanged.
DL01.8.4L2 Intermediate25 min
In review

The maths and code on this page are checked by automated tests on every change, but an independent reader has not signed it off yet. Spot a mistake or a confusing step? Tell us.

Before this lesson: DL01.8.3 First checks when training fails

What you will be able to do
  • List what a checkpoint must hold to resume a run exactly, and save and load it with torch.save and torch.load.
  • Keep the weights of the epoch with the lowest validation loss and restore them at the end of training.
  • Prove that restoring the full state of a run makes the resumed run identical to an uninterrupted one, and check it by comparing losses.
  • Write a fit(model, loss_fn, optimizer, train, val, epochs) function that returns a history and can resume from a checkpoint.

The question

A run is to last 150 epochs. At epoch 100 the laptop goes to sleep, or the shared machine ends your session. And the validation loss was lowest at epoch 73, long before the end. Can you carry on from epoch 100 so that the result is exactly what an uninterrupted run would have given, and come away with the weights of epoch 73 rather than those of the last epoch?

From the last lesson

1. What starting loss should a two-class and a ten-class classifier show, and why?

Show answer

About \({\ln 2 \approx 0.693}\) and \({\ln 10 \approx 2.303}\): with small logits every class gets probability about \(1/K\), and \({-\ln(1/K) = \ln K}\).

2. Training loss falling, validation loss rising: what is happening, and which weights do you want?

Show answer

Overfitting. You want the weights from the epoch with the lowest validation loss.

Intuition

A checkpoint is a file that holds everything a run needs to carry on: not just the weights, but the optimizer’s state, the epoch reached, the history so far, and the state of every random-number generator the run draws from. Write one after every epoch, and a stop costs at most one epoch of work.

A chart of loss against epoch from 0 to 150 for the wider network on 30 training points. The solid purple training loss falls to nearly zero, apart from a one-epoch spike near epoch 54 that leaves the top of the chart, labelled off the chart, 1.0 at epoch 54; the dashed purple validation loss falls to a black star at epoch 73, labelled best.pt, weights of epoch 73, lowest validation loss, then rises to about 0.24. A dotted vertical line at epoch 100 is labelled stopped after epoch 100, last.pt resumes it exactly. A row of small grey ticks along the bottom marks a checkpoint written every epoch.
DL01.8.4-F1

One run, two files. last.pt is rewritten every epoch and lets a stopped run carry on exactly; best.pt holds the weights of the epoch with the lowest validation loss.

Keeping the best weights is a separate job from resuming. On the machine that built this page, the validation loss of DL01.8.3’s overfitting run is lowest at epoch 73 and higher at every epoch after it, so the weights you want are not the last ones. The loop only has to remember the best so far: whenever an epoch beats it, copy the weights.

Build it up

Step 1. Weights are a dictionary. model.state_dict() maps each parameter’s name to its tensor: 0.weight of shape \({(16, 2)}\), 0.bias of shape \({(16)}\) and so on, 354 numbers for the 2-16-16-2 network. torch.save(model.state_dict(), path) writes it; a new model of the same architecture takes it back with load_state_dict(torch.load(path, weights_only=True)). The architecture itself is not saved: it lives in the code (DL01.7.3).

Step 2. The optimizer has a state too. optimizer.state_dict() holds its settings (the learning rate and the rest) and any buffers it keeps per parameter. Plain SGD keeps none, but SGD with momentum and Adam (DL02.1) do, and a resumed run without them takes different steps.

Step 3. Randomness has a state. Shuffling uses the loader’s generator, dropout uses PyTorch’s global generator. Each has a state, a tensor that torch.get_rng_state() (for the global generator) or the loader generator’s get_state() returns, and that torch.set_rng_state() or set_state() restores.

Step 4. Put it all in one dictionary, and save it after every epoch.

A grey panel representing one file, last.pt, written by torch.save after every epoch, holding five entries. In orange: model, the model's state_dict, weights and biases, 354 numbers in six tensors; optimizer, its state_dict, the learning rate and settings, plus buffers for optimizers that keep them. In purple: epoch and history, where to restart and the curves so far. In orange: best, with its epoch, loss and weights, the weights to keep at the end. In grey: torch_rng and loader_rng, the random states for dropout and the order of the next batches.
DL01.8.4-F2

What a checkpoint holds. Leave out any entry and the resumed run can differ from the original.

Step 5. Seeds make the start repeatable. Setting torch.manual_seed, NumPy’s seed, Python’s random.seed and the loader’s generator seed before building the model fixes the starting weights and the order of every batch. Two runs started this way give the same numbers on the same machine. The checkpoint does the same job for a run that is already under way.

The formal statement

NoteExact resumption

Let the state of a run after \(t\) epochs be \({S_t = (\theta_t,\, O_t,\, \rho_t)}\): the parameters, the optimizer’s state and the states of the random-number generators. With the code, the data and the machine fixed, one epoch is a fixed function \({S_{t+1} = F(S_t)}\). If a checkpoint stores \(S_k\) and loading restores it exactly, the resumed run reaches the same states \({S_{k+1}, \dots, S_E}\) as the uninterrupted one.

Proof, by induction on \(t \ge k\). Write \(S'_t\) for the states of the resumed run. At \({t = k}\), \({S'_k = S_k}\) because loading restores the stored state exactly. Suppose \({S'_t = S_t}\) for some \({t \ge k}\). The resumed run applies the same epoch to it, so

\[ S'_{t+1} = F(S'_t) = F(S_t) = S_{t+1}. \]

So \({S'_t = S_t}\) for every \(t\) from \(k\) to \(E\), and with them every loss and accuracy computed from them. ∎

Two chains of orange circles labelled S0, S1, dots, Sk, Sk plus 1, dots, SE, joined by green arrows. The top chain is the uninterrupted run, with the label one epoch: S t plus 1 equals F of S t. The bottom chain, stopped after epoch k and resumed in a new process, has no arrow between Sk and Sk plus 1: instead Sk is saved into a box labelled last.pt, which is loaded before applying F. A green two-headed arrow joins the two final states SE, labelled identical, bit for bit.
DL01.8.4-F4

Saving and loading a state in the middle of the chain does not change where the chain ends.

The proof shows exactly what can go wrong. If \(S_k\) is stored only in part, the resumed run applies \(F\) to a different state: the same weights but another batch order, or the same order without the momentum buffers. And if \(F\) itself changes, through another PyTorch version, another processor or another number of threads, the floating-point sums can round differently. The PyTorch notes say so plainly: completely reproducible results are not guaranteed across releases or platforms. Reproducibility in full is DL02.11.

Keeping the best. Let \(V_t\) be the validation loss after epoch \(t\). The kept weights are \({\theta_{t^\ast}}\) with \({t^\ast = \argmin_t V_t}\). The loop finds \(t^\ast\) without storing every epoch’s weights: it keeps the lowest \(V\) so far and replaces the stored copy only when an epoch is strictly lower.

Worked example

A run’s validation losses over eight epochs are 0.52, 0.40, 0.33, 0.31, 0.32, 0.30, 0.31 and 0.33.

  • Epochs 1 to 4 each beat the best so far, so each writes best.pt: the best is then 0.31 from epoch 4.
  • Epoch 5 scores 0.32, which is not lower: nothing is written, and best.pt still holds epoch 4.
  • Epoch 6 scores 0.30, a new best: best.pt now holds epoch 6.
  • Epochs 7 and 8 score 0.31 and 0.33, neither lower: nothing is written.

At the end the weights of epoch 6 are loaded, with validation loss 0.30. The last weights, from epoch 8, are worse on validation (0.33).

A table with columns epoch 1 to epoch 8 and four rows. Validation loss: 0.52, 0.40, 0.33, 0.31, 0.32, 0.30 highlighted, 0.31, 0.33. Best so far: 0.52, 0.40, 0.33, 0.31, 0.31, 0.30, 0.30, 0.30. Best epoch: 1, 2, 3, 4, 4, 6, 6, 6. Write best.pt: yes for epochs 1 to 4 and 6, no for 5, 7 and 8. A sentence below: at the end, fit loads the weights of epoch 6, loss 0.30, not those of epoch 8.
DL01.8.4-F3

The best-so-far rule traced over eight epochs.

Resuming, by the numbers. Suppose the run stopped after epoch 5. last.pt holds the weights of epoch 5, the optimizer’s state, the history of five epochs, the best so far (0.31 from epoch 4, with its weights) and the random states. The resumed run starts at epoch 6, finds 0.30 against the stored 0.31, and writes best.pt exactly as the uninterrupted run did.

Code

The module’s data and network, a seed function, and the two functions of DL01.8.2 (now returning loss and accuracy).

import copy, math, os, random, tempfile
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split

torch.set_num_threads(1)
X, labels = make_moons(n_samples=200, noise=0.2, random_state=0)
X_train, X_val, y_train, y_val = train_test_split(X, labels, test_size=0.25,
                                                  stratify=labels, random_state=0)
X_train, X_val = torch.tensor(X_train, dtype=torch.float32), torch.tensor(X_val, dtype=torch.float32)
y_train, y_val = torch.tensor(y_train), torch.tensor(y_val)
loss_fn = nn.CrossEntropyLoss()

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

def make_model(width=16, seed=0):
    set_seed(seed)
    return nn.Sequential(nn.Linear(2, width), nn.ReLU(), nn.Linear(width, width), nn.ReLU(),
                         nn.Linear(width, 2))

def loaders(n=150, batch_size=32, seed=0):
    train = DataLoader(TensorDataset(X_train[:n], y_train[:n]), batch_size=batch_size, shuffle=True,
                       generator=torch.Generator().manual_seed(seed))
    val = DataLoader(TensorDataset(X_val, y_val), batch_size=64)
    return train, val

def train_one_epoch(model, loss_fn, optimizer, data):
    model.train()
    total, correct, count = 0.0, 0, 0
    for X, y in data:
        out = model(X)
        loss = loss_fn(out, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total += loss.item() * len(y)
        correct += (out.argmax(dim=1) == y).sum().item()
        count += len(y)
    return total / count, correct / count

def evaluate(model, loss_fn, data):
    model.eval()
    total, correct, count = 0.0, 0, 0
    with torch.no_grad():
        for X, y in data:
            out = model(X)
            total += loss_fn(out, y).item() * len(y)
            correct += (out.argmax(dim=1) == y).sum().item()
            count += len(y)
    return total / count, correct / count

The fit function. Given a checkpoint path, it saves the whole state after every epoch, and if the file already exists it resumes from it. At the end it loads the best weights into the model.

def fit(model, loss_fn, optimizer, train, val, epochs, checkpoint=None):
    """Train until `epochs` epochs are done in all; return the history of the run."""
    history = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}
    best = {"epoch": 0, "loss": math.inf, "weights": copy.deepcopy(model.state_dict())}
    done = 0
    if checkpoint and os.path.exists(checkpoint):            # resume
        state = torch.load(checkpoint, weights_only=True)
        model.load_state_dict(state["model"])
        optimizer.load_state_dict(state["optimizer"])
        done, history, best = state["epoch"], state["history"], state["best"]
        torch.set_rng_state(state["torch_rng"])
        if train.generator is not None and state["loader_rng"] is not None:
            train.generator.set_state(state["loader_rng"])
    for epoch in range(done + 1, epochs + 1):
        results = train_one_epoch(model, loss_fn, optimizer, train) + evaluate(model, loss_fn, val)
        for key, value in zip(history, results):
            history[key].append(value)
        if history["val_loss"][-1] < best["loss"]:               # strictly better: keep these weights
            best = {"epoch": epoch, "loss": history["val_loss"][-1],
                    "weights": copy.deepcopy(model.state_dict())}
        if checkpoint:
            torch.save({"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(),
                        "history": history, "best": best, "torch_rng": torch.get_rng_state(),
                        "loader_rng": train.generator.get_state() if train.generator is not None else None},
                       checkpoint)
    model.load_state_dict(best["weights"])
    history["best_epoch"] = best["epoch"]
    return history

Keeping the best. The overfitting run of DL01.8.3: the wider network on 30 training points.

folder = tempfile.mkdtemp()
model = make_model(width=64)
optimizer = torch.optim.SGD(model.parameters(), lr=0.5)
train, val = loaders(n=30, batch_size=10)
history = fit(model, loss_fn, optimizer, train, val, epochs=150,
              checkpoint=os.path.join(folder, "small.pt"))
best = history["best_epoch"]
print(f"best epoch {best}: validation loss {history['val_loss'][best - 1]:.3f}, "
      f"accuracy {history['val_acc'][best - 1]:.2f}")
print(f"last epoch: validation loss {history['val_loss'][-1]:.3f}, accuracy {history['val_acc'][-1]:.2f}")
print("the returned model scores", evaluate(model, loss_fn, val))
best epoch 73: validation loss 0.079, accuracy 0.98
last epoch: validation loss 0.241, accuracy 0.88
the returned model scores (0.07908251136541367, 0.98)

The model that fit returns scores the best epoch’s validation loss, not the last epoch’s.

Resuming exactly. Run A trains 20 epochs without a break. Run B trains 10 epochs and stops; then new objects, built from a different seed so that nothing carries over by accident, resume from the checkpoint up to epoch 20. Run C does the same with a checkpoint whose loader state was removed.

def run(epochs, path=None, seed=0):
    model = make_model(seed=seed)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.3)
    train, val = loaders(seed=seed)
    return fit(model, loss_fn, optimizer, train, val, epochs, checkpoint=path)

a = run(20)
again = run(20)                                        # the same seed: the same run
path_b, path_c = os.path.join(folder, "b.pt"), os.path.join(folder, "c.pt")
run(10, path_b)
b = run(20, path_b, seed=123)                          # resumed in new objects
run(10, path_c)
state = torch.load(path_c, weights_only=True)
state["loader_rng"] = None                             # forget the order of the batches
torch.save(state, path_c)
c = run(20, path_c, seed=123)

print("same seed, same history:", a == again)
print("resumed with the full state, identical:", a == b)
gap = max(abs(x - y) for x, y in zip(a["train_loss"], c["train_loss"]))
print(f"resumed without the loader state: largest loss difference {gap:.3f}")
print("checkpoint keys:", list(torch.load(path_b, weights_only=True)))
same seed, same history: True
resumed with the full state, identical: True
resumed without the loader state: largest loss difference 0.067
checkpoint keys: ['epoch', 'model', 'optimizer', 'history', 'best', 'torch_rng', 'loader_rng']

The comparison is with ==: the resumed history equals the original to the last bit, as the proof says, and the run without the loader state drifts away from epoch 11 on. The optimizer entry of the checkpoint shows "state": {}, because plain SGD keeps no buffers.

Weights only, for later use. To use a trained network you need only its state_dict.

weights = os.path.join(folder, "weights.pt")
torch.save(model.state_dict(), weights)
copy_of_model = make_model(width=64, seed=7)           # same architecture, other start
copy_of_model.load_state_dict(torch.load(weights, weights_only=True))
copy_of_model.eval()
print("same outputs:", torch.equal(copy_of_model(X_val), model(X_val)))
same outputs: True

fit is used unchanged in DL01.9 and DL01.10, with the image loaders in place of these.

Explore

The two chains of states again: the uninterrupted run from S0 to SE, and the run saved to last.pt after Sk and resumed, which reaches the same SE.
DL01.8.4-F4

The induction of the formal statement as a picture.

Try the chain argument against the code: in the cell above, delete the line that restores torch_rng from fit and rerun. Run B still matches run A, because this network has no dropout and nothing else draws from the global generator after the weights are built. Then delete the line that restores loader_rng instead, and the match breaks from epoch 11 on, as for run C. The loop stepper of DL01.8.1 shows the five lines that each epoch repeats.

Pitfalls

Two charts of training loss over 20 epochs, each with a solid purple line for a run without a break and a dotted vertical line after epoch 10. Hollow circles show a run resumed after epoch 10. Left, crossed, model and optimizer only: from epoch 11 the circles wander off the line, with a largest difference of a few hundredths. Right, ticked, every random state too: the circles sit exactly on the line, difference exactly 0.
DL01.8.4-F5

Wrong beside right. A checkpoint without the loader’s random state resumes with another batch order, and the run drifts away from the original. With every state saved the two runs are identical.

“To resume, the weights are enough.”

The weights are \(\theta_k\), one part of the state \(S_k\). Without the loader’s state the batches come in a different order (the figure); without the optimizer’s buffers, momentum restarts from zero (exercise 2). The run still trains, but it is no longer the run you started, and a result you report cannot be reproduced.

“The final epoch’s weights are the trained model.”

On the 30-point run the last epoch has about three times the validation loss of the best epoch. Keep the best epoch’s weights, and say which epoch they came from.

“The best validation loss is an honest estimate of how good the model is.”

It was chosen because it was the lowest of 150 numbers, so it is optimistic. Picking an epoch is a choice made with the validation set, like any other hyperparameter; the honest estimate needs data used for nothing else: the test set of ML01.9.4, used once in DL01.10.3.

Check yourself

Recall

1. Name the entries a checkpoint needs to resume a run exactly.

Show answer

The model’s state_dict, the optimizer’s state_dict, the epoch reached (and the history so far and the best so far, for the curves and the kept weights), and the states of the random-number generators: the loader’s and PyTorch’s global one.

2. What does model.state_dict() contain, and what does it not contain?

Show answer

A dictionary from parameter names to tensors. It does not contain the architecture, which must be built in code before load_state_dict.

3. State the induction step that proves exact resumption.

Show answer

If \({S'_t = S_t}\) then \({S'_{t+1} = F(S'_t) = F(S_t) = S_{t+1}}\), where \(F\) is one epoch of training.

4. When does the best-so-far rule write best.pt?

Show answer

Only when an epoch’s validation loss is strictly lower than the lowest seen so far.

5. Why is the validation loss of the kept epoch an optimistic estimate?

Show answer

Because that epoch was chosen for having the lowest validation loss. An honest estimate needs a test set that played no part in any choice.

Exercises

1. ★ The validation losses of six epochs are 0.70, 0.45, 0.47, 0.38, 0.38 and 0.41. Which epochs write best.pt, and which weights does fit return?

Show solution

Epoch 1 (0.70, the first), epoch 2 (0.45) and epoch 4 (0.38) write it. Epoch 3 (0.47) is higher; epoch 5 (0.38) only ties, and the rule needs strictly lower; epoch 6 (0.41) is higher. fit returns the weights of epoch 4.

2. ★★ SGD with momentum (DL02.1) keeps a buffer \(v\) per parameter and updates \({v \leftarrow \beta v + g}\), then \({\theta \leftarrow \theta - \eta v}\). One parameter has buffer \({v = 2.0}\) when a run is saved; the next gradient is \({g = 1.0}\), with \({\beta = 0.9}\) and \({\eta = 0.1}\). Compare the first step after resuming with and without the optimizer’s state.

Show solution

With the state: \({v = 0.9 \cdot 2.0 + 1.0 = 2.8}\), so \({\theta}\) moves by \({-0.1 \cdot 2.8 = -0.28}\). Without it, PyTorch starts a new buffer at the gradient, \({v = 1.0}\), and \(\theta\) moves by \({-0.1 \cdot 1.0 = -0.1}\): barely more than a third of the step the original run takes. Every later step differs too, because the buffers differ.

3. ★★★ For the 2-16-16-2 network trained with plain SGD, which entries of the checkpoint could be dropped without breaking exact resumption, and why does fit save them anyway?

Show solution

PyTorch’s global random state: the network has no dropout, and after the weights are built nothing else draws from the global generator (the loader draws from its own). And the optimizer’s state: plain SGD keeps no buffers, and the new optimizer is built with the same learning rate. In the induction, neither part changes what \(F\) does. fit saves them because it is meant for every network and every optimizer: with dropout (DL02.5) the global state decides which units are dropped, and with momentum or Adam (DL02.1) the buffers decide every step, as exercise 2 shows.

Summary card

Save the whole state every epoch; keep the best weights; resume exactly.

The anchor figure again: training and validation loss over 150 epochs, a star at the lowest validation loss at epoch 73 marking best.pt, and a line at epoch 100 where last.pt lets the run resume exactly.
DL01.8.4-F1

last.pt to resume, best.pt to keep.

  • State \({S_t = (\theta_t, O_t, \rho_t)}\); one epoch \({S_{t+1} = F(S_t)}\); restore \(S_k\) and every later state matches.
  • Checkpoint: model and optimizer state_dict, epoch, history, best, random states.
  • Best weights: \({\theta_{t^\ast}}\) with \({t^\ast = \argmin_t V_t}\), replaced only on a strict improvement.
  • fit(model, loss_fn, optimizer, train, val, epochs) returns the history and the best model.

Sources and next steps

Sources. Each was opened and checked against this lesson on 15 November 2026.

  • PyTorch tutorials, Saving and loading models: a state_dict maps each layer to its parameter tensors and optimizers have one too; a general checkpoint for resuming training must hold more than the model’s state_dict, and the example saves the epoch, both state_dicts and the loss; call model.eval() or model.train() after loading; the example loads with weights_only=True.
  • PyTorch 2.15 documentation, Reproducibility: completely reproducible results are not guaranteed across releases, commits or platforms; torch.manual_seed seeds the generator for all devices; NumPy needs its own seed; a DataLoader is made repeatable with a seeded generator.
  • Aston Zhang, Zachary C. Lipton, Mu Li and Alexander J. Smola, Dive into Deep Learning, section 6.6, File I/O: saving a model saves its parameters, not the architecture, which must be generated in code before the parameters are loaded.

Next. DL01.9, Datasets, loaders and batching, opens up the DataLoader that this module used as a black box.

This lesson unlocks

DL01.9 Datasets, loaders and batching

Your progress is saved in this browser only. No account, no tracking.
Back to top