PyTorch Model Checkpointing

Checkpointing saves the state of your model and training progress to disk at regular intervals. If training crashes, you restart from the last checkpoint instead of the beginning. Checkpointing also lets you save the best-performing model during training so you can always recover peak performance even if later epochs cause overfitting.

What to Save in a Checkpoint

Minimum (inference only):
  model.state_dict()   ← all weights and biases

Full training checkpoint (resumable):
  model.state_dict()
  optimizer.state_dict()   ← optimizer momentum, learning rate state
  scheduler.state_dict()   ← scheduler progress
  epoch                    ← where to resume from
  best_val_loss            ← for early stopping logic

state_dict — What It Is

A state dict is an ordered dictionary that maps each layer name to its parameters (weights and biases). It does not store the model architecture — only the learned values. You must recreate the model class before loading a state dict into it.

import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 2)
)

sd = model.state_dict()
for k, v in sd.items():
    print(f"{k}: {v.shape}")
# 0.weight: torch.Size([8, 4])
# 0.bias:   torch.Size([8])
# 2.weight: torch.Size([2, 8])
# 2.bias:   torch.Size([2])

Saving a Model

import torch

# Save weights only (for inference)
torch.save(model.state_dict(), 'model_weights.pth')

# Save full checkpoint (for resuming training)
checkpoint = {
    'epoch': 25,
    'model_state': model.state_dict(),
    'optimizer_state': optimizer.state_dict(),
    'scheduler_state': scheduler.state_dict(),
    'best_val_loss': 0.142
}
torch.save(checkpoint, 'checkpoint_epoch25.pth')

Loading a Model

# Load weights only
model = MyModel()   # must recreate architecture first
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

# Load full checkpoint to resume training
checkpoint = torch.load('checkpoint_epoch25.pth')
model.load_state_dict(checkpoint['model_state'])
optimizer.load_state_dict(checkpoint['optimizer_state'])
scheduler.load_state_dict(checkpoint['scheduler_state'])
start_epoch = checkpoint['epoch'] + 1
best_val_loss = checkpoint['best_val_loss']

print(f"Resuming from epoch {start_epoch}")

Save Best Model During Training

best_val_loss = float('inf')

for epoch in range(100):
    train_one_epoch(model, train_loader, optimizer, loss_fn)

    val_loss = evaluate(model, val_loader, loss_fn)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best_model.pth')
        print(f"  ✓ Saved best model at epoch {epoch+1} (val_loss={val_loss:.4f})")

# After training: load the best-performing checkpoint
model.load_state_dict(torch.load('best_model.pth'))

Periodic Checkpointing

for epoch in range(100):
    train_one_epoch(model, train_loader, optimizer, loss_fn)

    # Save every 10 epochs
    if (epoch + 1) % 10 == 0:
        ckpt = {
            'epoch': epoch,
            'model_state': model.state_dict(),
            'optimizer_state': optimizer.state_dict()
        }
        torch.save(ckpt, f'checkpoint_epoch{epoch+1}.pth')
        print(f"  Saved checkpoint at epoch {epoch+1}")

Loading on a Different Device

If a checkpoint was saved on a GPU and you want to load it on a CPU (or a different GPU), use the map_location argument.

# Saved on GPU, loading on CPU
model.load_state_dict(
    torch.load('model_weights.pth', map_location=torch.device('cpu'))
)

# Load to current device automatically
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(
    torch.load('model_weights.pth', map_location=device)
)

Checkpoint Naming Strategy

Good naming conventions:
  best_model.pth               ← always the best so far
  checkpoint_epoch_025.pth     ← periodic saves with zero-padded epoch
  checkpoint_final.pth         ← saved at end of training

Folder approach for long experiments:
  checkpoints/
  ├── best.pth
  ├── epoch_010.pth
  ├── epoch_020.pth
  └── epoch_030.pth

Full Checkpoint Pattern — End to End

import torch
import os

def save_checkpoint(model, optimizer, epoch, val_loss, path):
    torch.save({
        'epoch': epoch,
        'model_state': model.state_dict(),
        'optimizer_state': optimizer.state_dict(),
        'val_loss': val_loss
    }, path)

def load_checkpoint(path, model, optimizer):
    ckpt = torch.load(path)
    model.load_state_dict(ckpt['model_state'])
    optimizer.load_state_dict(ckpt['optimizer_state'])
    return ckpt['epoch'], ckpt['val_loss']

# Resume if checkpoint exists
start_epoch = 0
ckpt_path   = 'latest.pth'

if os.path.exists(ckpt_path):
    start_epoch, _ = load_checkpoint(ckpt_path, model, optimizer)
    print(f"Resumed from epoch {start_epoch}")

for epoch in range(start_epoch, 100):
    # ... training ...
    save_checkpoint(model, optimizer, epoch, val_loss, ckpt_path)

Summary

Checkpointing saves model weights and training state to disk. Use torch.save(model.state_dict(), path) for inference-only saves. Save the full checkpoint dictionary — model, optimizer, scheduler states, epoch, and best metric — when you need to resume training. Save the best model separately whenever validation loss improves. Load checkpoints with model.load_state_dict(torch.load(path)) after recreating the model architecture. Use map_location when loading across different devices. Checkpointing protects weeks of training from hardware failures and gives you the best-performing model at the end regardless of when overfitting starts.