Saving and loading models (checkpoints)

The Core Problem: What Do We Actually Need to Save?

When training a neural network, we're updating millions (or billions) of parameters via gradient descent. The training state consists of:

1. Model Parameters: The weights $\mathbf{W}$ and biases $\mathbf{b}$ for every layer
2. Optimizer State: Momentum bufers (SGD with momentum), running averages (Adam's $m_t$, $v_t$), learning rate schedules
3. Epoch/Step Counter: Where we are in training
4. Random Seeds: For reproducibility
5. Model Architecture (optional): The blueprint to reconstruct the network

Why all of this? Parameters alone define the model's predictions, but optimizer state is needed to continue training smoothly. If you only save weights and later resume with a fresh optimizer, you lose accumulated momentum—like a car suddenly starting from zero velocity instead of cruising speed.

Framework-Specific Approaches

PyTorch: State Dictionaries

Why this design? Separation of architecture (code) and parameters (data). Your model class can evolve, but old checkpoints remain loadable if layer names match.

import torch
import torch.nn as nn
 
# Define model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
 
# After training for N epochs...
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': train_loss,
    'accuracy': val_acc
}
torch.save(checkpoint, 'model_epoch_10.pt')

Why this step? We bundle everything needed to resume training into a single dict, then serialize with torch.save (uses Python's pickle under the hood).

Loading back:

checkpoint = torch.load('model_epoch_10.pt')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
start_epoch = checkpoint['epoch'] + 1

Why load_state_dict and not direct assignment? It validates tensor shapes and missing/unexpected keys, catching architecture mismatches.

TensorFlow/Keras: Two Paradigms

1. TensorFlow Checkpoints (.ckpt)

Saves only weights as binary files + an index:

model = tf.keras.Sequential([...])
checkpoint_path = "training_1/cp-{epoch:04d}.ckpt"
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    save_weights_only=True,
    save_freq='epoch'
)
model.fit(x_train, y_train, epochs=10, callbacks=[checkpoint_callback])
 
# Load
model.load_weights('training_1/cp-010.ckpt')

Why save_weights_only=True? Faster saves, smaller files. Requires you to define the model architecture separately.

2. SavedModel Format (Full Model)

model.save('my_model')  # Creates a directory with architecture + weights
loaded_model = tf.keras.models.load_model('my_model')

Why this format? Bundles architecture, weights, and even the computational graph—no need for the original code. Ideal for production deployment with TensorFlow Serving.

Checkpoint Strategies: When and What to Save

Strategy 1: Save Only the Best Model

# PyTorch example
best_acc = 0.0
for epoch in range(num_epochs):
    val_acc = evaluate(model, val_loader)
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'accuracy': val_acc
        }, 'best_model.pt')

Why this works? Most experiments only need the peak performance model. Intermediate checkpoints are wasted storage.

Strategy 2: Save Every N Epochs + Keep Last K

# Keras callback
checkpoint_callback = ModelCheckpoint(
    filepath='model_{epoch:02d}_{val_loss:.2f}.h5',
    save_best_only=False,
    save_freq='epoch',
    period=5  # Every 5 epochs
)

Then manually prune old checkpoints or use a rolling window (keep only the last 3).

Why? Long training runs benefit from fallback points. If validation loss spikes at epoch 47, you can reload from epoch 45.

Strategy 3: Save on Plateau or Early Stopping

from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
 
early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)
checkpoint = ModelCheckpoint('checkpoint.h5', monitor='val_loss', save_best_only=True)

Why patience=10? Validation metrics are noisy. A single bad epoch shouldn't trigger a stop. Patience gives the model time to escape local minima.

Derivation: Why Optimizer State Matters

Consider SGD with momentum:

Where:

• $\mathbf{v}_t$ is the velocity (moving average of gradients)
• $\beta$ is the momentum coefficient (typically 0.9)
• $\alpha$ is the learning rate

What happens if you only save $\theta_t$ and restart?

At resume, $\mathbf{v}_t$ resets to $\mathbf{0}$. The first step becomes vanilla SGD:

$\theta_{t+1} = \theta_t - \alpha \nabla_\theta \mathcal{L}$

Why is this bad? Momentum accumulates gradient information over many steps. Resetting it means:

1. Loss of directional inertia: The model "forgets" which direction was promising
2. Slower convergence: Takes ~$1/(1-\beta)$ steps to rebuild momentum (11 steps for $\beta=0.9$)
3. Unstable training: If you're in a sharp valley, the first few steps might diverge

The fix: Save optimizer.state_dict() which includes $\mathbf{v}_t$ for every parameter.

Example 1: Resuming Training After Crash

import torch
import torch.nn as nn
 
# Training loop with checkpointing
def train_with_checkpoints(model, optimizer, train_loader, num_epochs, checkpoint_dir):
    start_epoch = 0
    # Try to resume from latest checkpoint
    checkpoint_files = sorted(Path(checkpoint_dir).glob('checkpoint_*.pt'))
    if checkpoint_files:
        latest = checkpoint_files[-1]
        checkpoint = torch.load(latest)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        start_epoch = checkpoint['epoch'] + 1
        print(f"Resumed from epoch {start_epoch}")
    for epoch in range(start_epoch, num_epochs):
        for batch in train_loader:
            loss = train_step(model, batch, optimizer)
        # Save checkpoint every epoch
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss.item()
        }, f'{checkpoint_dir}/checkpoint_{epoch:04d}.pt')

Why start_epoch adjustment? Prevents repeating work. If we crashed at epoch 7, we pick up at epoch 8.

Why sort and take [-1]? Files might be listed out of order. Sorting by name (with zero-paded numbers) ensures we load the truly latest checkpoint.

Example 2: Exporting for Inference (No Optimizer)

# Training complete, save inference-only checkpoint
torch.save({
    'model_state_dict': model.state_dict(),
    'accuracy': final_test_acc,
    'class_names': ['cat', 'dog', 'bird']
}, 'production_model.pt')
 
# Load in production
checkpoint = torch.load('production_model.pt', map_location='cpu')
model = MyModel()  # Architecture must match
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()  # Disable dropout/batchnorm training mode
 
# Inference
with torch.no_grad():
    predictions = model(new_data)

Why map_location='cpu'? The model was trained on GPU (CUDA). Loading directly on a CPU server would fail without this flag.

Why model.eval()? Dropout and BatchNorm behave differently during training (stochastic) vs. inference (deterministic). Forgetting this is a common deployment bug.

Advanced: Distributed Training Checkpoints

When training across multiple GPUs, each device holds a shard of the model. Checkpointing requires coordination:

# PyTorch DistributedDataParallel
if torch.distributed.get_rank() == 0:  # Only master process saves
    torch.save({
        'model_state_dict': model.module.state_dict(),  # .module unwraps DP
        'optimizer_state_dict': optimizer.state_dict()
    }, 'checkpoint.pt')

Why only rank 0? All GPUs have identical parameters (synchronized via AllReduce). Saving redundantly wastes I/O.

Why .module? DP wraps your model in a DistributedDataParallel object. The actual model is model.module.

Connections

• Overfitting and Regularization: Checkpoints enable early stopping by saving the model before validation loss diverges.
• Learning Rate Schedules: LR schedulers depend on epoch counters; checkpoints must preserve this state.
• Distributed Training: Multi-GPU setups require synchronized checkpointing to avoid data corruption.
• Model Deployment: Production models are loaded from checkpoints; architecture mismatches cause runtime failures.
• Transfer Learning: Pretrained models (BERT, ResNet) are just checkpoints from someone else's training run.

#flashcards/ai-ml

What is a checkpoint in deep learning? :: A saved snapshot of a model's parameters (weights, biases), optimizer state, and training metadata (epoch, loss) to disk, allowing training resumption or deployment.

Concept Map

Hinglish (regional understanding)