Checkpoint: weights are only part of a training run

A checkpoint saves model state at a point in training. For inference, weights plus the matching architecture and preprocessing may be enough. Resuming training also needs the optimizer's state: momentum and adaptive statistics are learned history too.

Save the state you will restore

These PyTorch helpers save and restore a model, optimizer and update counter. They expect already-created matching model and optimizer objects.

import torch

def save_checkpoint(path, model, optimizer, step):
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "step": step,
    }, path)

def restore_checkpoint(path, model, optimizer):
    state = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(state["model"])
    optimizer.load_state_dict(state["optimizer"])
    return state["step"]

The architecture is still supplied by code; saving a state dictionary does not define it. Exact continuation can additionally require the scheduler, mixed-precision scaler, random-number states and data position.

A checkpoint is also different from an exported deployment artifact. Quantization, conversion and runtime packaging may change the representation, while the training checkpoint remains useful for further learning.

Keep a bounded set of useful checkpoints rather than every update. A recent recovery point and a validation-selected model serve different purposes.