Checkpointing — Save Everything You'd Need to Resume

state_dict over full model — every time

Two ways to save a model:

state_dict (recommended) — saves only the parameters and buffers. To restore, you instantiate the model class and call load_state_dict. Robust to code changes, framework version differences, and class renames.
torch.save(model) — pickles the whole Python object. Fragile: it requires the exact class definition on load, and any rename / refactor breaks it.

Always use state_dict.

What goes in a "training" checkpoint

For inference, just save the model state_dict. For resuming training, you need:

Model state_dict
Optimizer state_dict (Adam's running moments are not trivial to recreate)
Scheduler state_dict (the current step / last_lr / etc.)
Current epoch / step number
Best validation metric so far (for early stopping)
RNG state if you care about exact reproducibility

weights_only=True

When loading, pass weights_only=True. PyTorch 2.x will eventually default to it, but explicit is good. It refuses to deserialize arbitrary Python — only tensor data. Defends against the (real) class of attacks where a malicious .pt file executes code on load.

Early stopping

Save a "best so far" checkpoint whenever val loss improves; track patience and stop when N epochs pass without improvement. Five lines of class state, prevents wasted compute.

Code

Save and load — the simple inference case·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

# Train
model = TinyMLP()
torch.save(model.state_dict(), 'tiny_mlp.pt')

# Load — instantiate first, then restore weights
model2 = TinyMLP()
model2.load_state_dict(torch.load('tiny_mlp.pt', weights_only=True))
model2.eval()

Full training-resume checkpoint·python

import torch

def save_ckpt(path, model, optimizer, scheduler, epoch, best_val):
    torch.save({
        'epoch': epoch,
        'best_val': best_val,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'rng_state': torch.get_rng_state(),
    }, path)

def load_ckpt(path, model, optimizer, scheduler):
    ckpt = torch.load(path, weights_only=False)   # need False for non-tensor objects
    model.load_state_dict(ckpt['model_state_dict'])
    optimizer.load_state_dict(ckpt['optimizer_state_dict'])
    scheduler.load_state_dict(ckpt['scheduler_state_dict'])
    torch.set_rng_state(ckpt['rng_state'])
    return ckpt['epoch'], ckpt['best_val']

Early stopping — five lines of state·python

class EarlyStopping:
    def __init__(self, patience=5, min_delta=0.0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best = None
        self.should_stop = False

    def __call__(self, val_loss):
        if self.best is None or val_loss < self.best - self.min_delta:
            self.best = val_loss
            self.counter = 0
        else:
            self.counter += 1
            if self.counter >= self.patience:
                self.should_stop = True

# Usage
early = EarlyStopping(patience=5)
for epoch in range(100):
    train_one_epoch(...)
    val_loss, _ = evaluate(...)
    early(val_loss)
    if early.should_stop:
        print(f"Early stopping at epoch {epoch}")
        break

Checkpointing — Save Everything You'd Need to Resume

state_dict over full model — every time

What goes in a "training" checkpoint

weights_only=True

Early stopping

Code

External links

Exercise

Progress

Comments 0