Complete PyTorch Series — Module 11 of 12

Deep neural networks often train for hours, days, or weeks across hundreds of epochs. Without robust serialization and checkpointing, hardware failures, spot-instance terminations, or unexpected power cuts can erase days of compute progress. In this module, you will master PyTorch serialization best practices and implement an Early Stopping engine.

state_dict vs Full Model Pickling

PyTorch offers two ways to save a model, but only one is recommended for production:

Why You Should NEVER use torch.save(model, path)

# AVOID THIS PATTERN!
torch.save(model, "entire_model.pth")

Saving the entire model uses Python’s pickle module, which serializes the exact directory layout and Python class definitions. If you rename your class, move the file, or load the weights in another project or C++ environment, deserialization will fail with ModuleNotFoundError or unpickling exceptions.

The Recommended Standard: state_dict

A state_dict is a standard Python dictionary that maps each layer name to its learnable parameter tensors. It contains pure numerical weights with zero dependency on Python code structures:

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 2)
)

# 1. Inspect the state_dict
print("Model keys in state_dict:")
for param_tensor in model.state_dict():
    print(f" - {param_tensor:20s}: shape {model.state_dict()[param_tensor].shape}")

# 2. Save only the weights dictionary
torch.save(model.state_dict(), "model_weights.pth")
print("\nModel weights successfully saved to model_weights.pth")

# 3. Load weights into a newly instantiated model
new_model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 2)
)

new_model.load_state_dict(torch.load("model_weights.pth"))
new_model.eval() # Always set to eval() after loading!
print("Weights loaded into new model successfully.")

Production Checkpoint Bundles for Seamless Resuming

When training multi-epoch models, saving only model weights is insufficient to resume training faithfully. You must also preserve the optimizer state (such as Adam’s momentum buffers), the learning rate scheduler state, the current epoch number, and the best validation score:

import torch

# Saving a full training checkpoint bundle
checkpoint = {
    'epoch': 45,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'best_val_loss': 0.284,
}
torch.save(checkpoint, "checkpoint_epoch_45.tar")
print("Full checkpoint bundle saved.")

# Resuming training from checkpoint
def resume_training(checkpoint_path, model, optimizer, scheduler):
    checkpoint = torch.load(checkpoint_path, map_location='cpu')

    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    start_epoch = checkpoint['epoch'] + 1
    best_loss = checkpoint['best_val_loss']

    print(f"Resumed successfully at epoch {start_epoch} with best loss {best_loss:.4f}")
    return start_epoch, best_loss

Loading Weights Across Different Devices (GPU to CPU)

If you trained a model on a multi-GPU cloud cluster and need to load it on a CPU server or MacBook for inference, attempting to load directly will crash because PyTorch tries to allocate weights to GPU devices that do not exist.

Use the map_location parameter to redirect tensors during deserialization:

# Load a GPU-trained model safely on a CPU machine
device = torch.device("cpu")
state_dict = torch.load("model_weights.pth", map_location=device)
model.load_state_dict(state_dict)

Production Early Stopping Callback Implementation

Early stopping halts training when validation loss stops improving for a specified number of consecutive epochs (patience), preventing overfitting and saving compute credits:

import numpy as np
import torch

class EarlyStopping:
    def __init__(self, patience=5, min_delta=1e-4, checkpoint_path="best_model.pth"):
        self.patience = patience
        self.min_delta = min_delta
        self.checkpoint_path = checkpoint_path
        self.counter = 0
        self.best_loss = np.inf
        self.early_stop = False

    def __call__(self, val_loss: float, model: torch.nn.Module):
        # Check if loss improved by at least min_delta
        if val_loss < (self.best_loss - self.min_delta):
            print(f"Validation loss improved ({self.best_loss:.4f} --> {val_loss:.4f}). Saving model...")
            self.best_loss = val_loss
            torch.save(model.state_dict(), self.checkpoint_path)
            self.counter = 0 # Reset patience counter
        else:
            self.counter += 1
            print(f"EarlyStopping patience: {self.counter} of {self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True
                print("Early stopping triggered! Training halted.")

# Usage inside training loop:
early_stopping = EarlyStopping(patience=3)
# for epoch in range(epochs):
#     ... validation pass ...
#     early_stopping(val_loss, model)
#     if early_stopping.early_stop:
#         break

Module 11 Recap

• Always save models using model.state_dict(), never pickle the entire model object.
• Always call model.eval() immediately after loading weights to lock Dropout and BatchNorm.
• Save complete checkpoint bundles (model, optimizer, scheduler, epoch, metric) to allow flawless training resumption.
• Use torch.load(..., map_location=device) to load weights seamlessly across mismatched hardware.
• Implement EarlyStopping to avoid overfitting and save compute costs.

Next up — Module 12: Production deployment: export models to TorchScript and ONNX, and serve low-latency predictions with FastAPI.