Model Checkpointing, Serialization & Security

1. The Golden Rule: Always Save state_dict

PyTorch provides two methods to serialize models using torch.save(), but one of them is an anti-pattern:

torch.save(model, "model_complete.pt")

torch.save(model.state_dict(), "weights.pt")

2. Saving & Loading for Inference

When deploying to production for inference, serialize the model's weights and re-initialize the architecture prior to loading:

import torch
import torchvision.models as models

# --- SAVING WEIGHTS ---
model = models.resnet18(weights=None)
# ... after training finishes ...
torch.save(model.state_dict(), "resnet18_production.pt")

# --- LOADING WEIGHTS (Inference Service) ---
# 1. Instantiate identical architecture
inference_model = models.resnet18(weights=None)

# 2. Load the state dictionary safely (PyTorch 2.4+ default)
state_dict = torch.load("resnet18_production.pt", weights_only=True)

# 3. Inject weights into module
inference_model.load_state_dict(state_dict)

# 4. CRITICAL: Switch to evaluation mode!
inference_model.eval()

# Move to target hardware (e.g. CPU or GPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
inference_model.to(device)

3. Resuming Preempted Training: Comprehensive Checkpoints

On cloud clusters with spot instances or long training runs, node preemption can kill your process. To resume without loss of progress, your checkpoint must preserve optimizer momentum buffers and scheduler states:

# --- SAVING FULL CHECKPOINT ---
checkpoint = {
    'epoch': current_epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'best_val_loss': best_val_loss,
    'rng_state': torch.get_rng_state()  # Guarantees exact reproducibility
}
torch.save(checkpoint, f"checkpoint_epoch_{current_epoch}.tar")

# --- RESUMING FROM CHECKPOINT ---
def resume_training(checkpoint_path, model, optimizer, scheduler):
    print(f"Resuming training from {checkpoint_path}...")
    checkpoint = torch.load(checkpoint_path, weights_only=False)
    
    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']
    torch.set_rng_state(checkpoint['rng_state'])
    
    print(f"Resumed successfully! Restarting at Epoch {start_epoch}")
    return start_epoch, best_loss

4. Cross-Device Hardware Loading with map_location

If you train a model on an NVIDIA GPU (e.g., cuda:0) and attempt to load it on a CPU-only server or Apple Silicon Mac without flags, PyTorch errors with RuntimeError: Attempting to deserialize object on a CUDA device.... Use map_location:

5. The Pickle Vulnerability & weights_only=True

To prevent this vulnerability:

• Always use weights_only=True: Starting in PyTorch 2.4, weights_only=True is the default. It constrains the deserializer strictly to pure numerical tensors, primitive datatypes, and state dictionaries.
• Migrate to Safetensors: For production Hugging Face and open-weights models, the safetensors format replaces pickle entirely with a zero-copy memory-mapped binary format.