🎓 A1Learner
🔍

PyTorch Checkpoint and Resume

A checkpoint is a file that holds enough state to continue a training run later. Saving the model, the optimizer, the epoch, and the random state lets a resumed run match an unbroken one.

Random state round trip

Saving the random state freezes the sequence of random draws. Restoring it replays the same draws.

torch.manual_seed(7)
state = torch.get_rng_state()
a = torch.randn(2)
torch.set_rng_state(state)
b = torch.randn(2)
print(torch.equal(a, b))

What the file holds

One dict carries all four pieces. The keys print in sorted order.

checkpoint = {"model": "weights", "optimizer": "state", "epoch": 3, "rng": "state"}
print(sorted(checkpoint.keys()))

Example

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

def make_data():
    g = torch.Generator().manual_seed(3)
    x = torch.randn(64, 2, generator=g)
    y = (x[:, 0] * 2 - x[:, 1]).reshape(-1, 1)
    return x, y

def train_epochs(model, opt, loader, x, y, loss_fn, epochs, trace):
    for _ in range(epochs):
        model.train()
        for batch_x, batch_y in loader:
            loss = loss_fn(model(batch_x), batch_y)
            opt.zero_grad()
            loss.backward()
            opt.step()
        model.eval()
        with torch.no_grad():
            trace.append(round(loss_fn(model(x), y).item(), 4))

x, y = make_data()
loss_fn = nn.MSELoss()

# One unbroken run of 6 epochs
torch.manual_seed(5)
model_a = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
opt_a = torch.optim.Adam(model_a.parameters(), lr=0.02)
loader_a = DataLoader(TensorDataset(x, y), batch_size=16, shuffle=True)
trace_a = []
train_epochs(model_a, opt_a, loader_a, x, y, loss_fn, 6, trace_a)

# A run split at epoch 3 by a checkpoint
torch.manual_seed(5)
model_b = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
opt_b = torch.optim.Adam(model_b.parameters(), lr=0.02)
loader_b = DataLoader(TensorDataset(x, y), batch_size=16, shuffle=True)
trace_b = []
train_epochs(model_b, opt_b, loader_b, x, y, loss_fn, 3, trace_b)
torch.save({
    "model": model_b.state_dict(),
    "optimizer": opt_b.state_dict(),
    "epoch": 3,
    "rng": torch.get_rng_state(),
}, "checkpoint.pt")

# "Restart": fresh objects, everything restored from the file
fresh = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
fresh_opt = torch.optim.Adam(fresh.parameters(), lr=0.02)
saved = torch.load("checkpoint.pt")
fresh.load_state_dict(saved["model"])
fresh_opt.load_state_dict(saved["optimizer"])
torch.set_rng_state(saved["rng"])
loader_c = DataLoader(TensorDataset(x, y), batch_size=16, shuffle=True)
train_epochs(fresh, fresh_opt, loader_c, x, y, loss_fn, 3, trace_b)

print("Unbroken run:", trace_a)
print("Resumed run: ", trace_b)
print("Traces match:", trace_a == trace_b)

Key points

  • Checkpoint: one saved file holding the pieces a run needs to continue
  • Use torch.save on a dict to store model, optimizer, epoch, and random state together
  • Optimizer state matters, because Adam remembers per-weight step sizes
  • Use torch.get_rng_state to capture the random state and torch.set_rng_state to restore it
  • A shuffled DataLoader draws from the global random state, so its order depends on that state

In this example, the data is 64 rows built with its own generator and seed. The first run trains for six unbroken epochs and records the loss after each one. The second run starts from the same seed and trains three epochs. It then saves a checkpoint dict with the model state, the optimizer state, the epoch, and the random state. Fresh model and optimizer objects load that state, and the random state is restored. A new loader shuffles with the restored state, so its order continues the original sequence. Three more epochs complete the second trace. The two traces print side by side, and the final line confirms they match exactly.

Note

Saving only the model weights restarts the optimizer and the shuffle, so the resumed run drifts from the original

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the checkpoint include the random state?