A checkpoint is a dictionary you can load back. It must include: model state dict, optimizer state dict, scheduler state dict, GradScaler state dict (if using mixed precision), step number, epoch number, and your config (so you know what trained this checkpoint).
Saving only the model state dict is fine for inference but useless for resuming training. The optimizer and scheduler carry running averages and momentum that you can't reconstruct from the model alone.
Tip: Always save the optimizer too. The number of times someone resumed from 'just the model' and watched their carefully tuned learning rate restart from zero is uncountable.
Best-checkpoint vs latest-checkpoint
Best — the checkpoint at the highest validation metric so far. This is what you ship.
Latest — the most recent checkpoint, regardless of metric. This is what you resume from after a crash.
Both should exist. Best is for production; latest is for engineering. Save them to different filenames so you don't overwrite one with the other.
Where to save
Local disk is fine for short training. For multi-day runs, save to S3 / GCS / a fileserver so the checkpoint survives a node death. Some training frameworks (PyTorch Lightning, Composer, HF Trainer) handle this for you.
Principle: Checkpointing is insurance. It is the single cheapest thing you can do to make a long training run not waste a week of GPU time when something crashes at 80%.
Code
Save and resume a full training state·python
import torch
def save_ckpt(path, model, opt, sch, scaler, step, epoch, cfg, metric):
torch.save({
"model": model.state_dict(),
"optimizer": opt.state_dict(),
"scheduler": sch.state_dict() if sch else None,
"scaler": scaler.state_dict() if scaler else None,
"step": step,
"epoch": epoch,
"config": cfg,
"best_metric": metric,
}, path)
def load_ckpt(path, model, opt=None, sch=None, scaler=None):
ckpt = torch.load(path, map_location="cpu")
model.load_state_dict(ckpt["model"])
if opt and ckpt.get("optimizer"):
opt.load_state_dict(ckpt["optimizer"])
if sch and ckpt.get("scheduler"):
sch.load_state_dict(ckpt["scheduler"])
if scaler and ckpt.get("scaler"):
scaler.load_state_dict(ckpt["scaler"])
return ckpt["step"], ckpt["epoch"], ckpt.get("best_metric", 0.0)
Train for 5 epochs, save a checkpoint, kill the process, then resume and train 5 more epochs. The validation curve should be continuous across the resume — no jump in either direction. If there is a jump, something in your save/load is missing.
Progress
Progress is local-only — sign in to sync across devices.