Python PyTorch: Save and Load state_dict and Checkpoints

When working with PyTorch, saving and loading model state is a routine task that appears in almost every training pipeline. The phrase python pytorch save load state_dict and checkpoints covers two related but distinct operations: saving the model's parameter tensors via state_dict, and saving a full checkpoint that includes optimizer state, epoch, and other training metadata. Understanding the difference and knowing when to use each is essential for reproducible experiments and resumable training.

What a state_dict Contains

Every PyTorch nn.Module has a state_dict() method that returns a Python dictionary mapping each learnable parameter and persistent buffer to its current tensor value. For a convolutional network, this includes the convolution weights and biases, batch norm running statistics, and any other registered parameters. The state_dict is a plain dictionary, so it can be inspected, modified, and saved directly.

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
print(model.state_dict().keys())
# odict_keys(['weight', 'bias'])

The keys are the exact names of the parameters as they appear in the module's named_parameters(). For nested modules, the keys are prefixed with the module path, such as features.0.weight. This naming is what allows the same model architecture to load a state_dict later.

Saving and Loading a state_dict

The simplest way to persist a model is to save only the state_dict. This is sufficient when you need to reuse the model for inference or transfer learning, and when you do not need to resume training from the exact point where it stopped.

# Save
torch.save(model.state_dict(), 'model_weights.pth')

# Load
model = nn.Linear(10, 5)  # must match the original architecture
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

torch.save uses Python's pickle to serialize the dictionary, and torch.load deserializes it. The file extension is arbitrary; .pt, .pth, and .bin are all common. When loading, you must instantiate the model first and then call load_state_dict on it. The model architecture must match the saved parameters exactly, otherwise the load will fail.

Saving and Loading a Full Checkpoint

A full checkpoint goes beyond the model parameters. It also saves the optimizer's state_dict, the current epoch, learning rate scheduler state, and any other training variables you need to resume training exactly where you left off. This is the standard approach for long-running training jobs where interruptions are expected.

# During training
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, 'checkpoint.pt')

# To resume
def load_checkpoint(path, model, optimizer):
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch'] + 1
    return model, optimizer, start_epoch

The checkpoint dictionary is a regular Python dict, so you can add any metadata you need, such as validation metrics, random number generator states, or configuration arguments. This makes the checkpoint self-contained and allows you to reproduce the exact training state.

Handling Device Placement When Loading

PyTorch tensors are device-specific. When you save a model trained on a GPU, the state_dict tensors are stored on the GPU. If you later load that state_dict on a machine without a GPU, or on a different GPU, you must explicitly move the tensors to the target device. The torch.load function accepts a map_location argument to handle this.

# Load on CPU from a GPU-trained checkpoint
device = torch.device('cpu')
checkpoint = torch.load('checkpoint.pt', map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])

Without map_location, PyTorch will attempt to load tensors to their original device, which fails if that device is unavailable. For mixed-device environments, a common pattern is to load with map_location='cpu' first, then move the model to the target device after loading. This avoids partial-load errors and keeps the code portable.

Dealing with Missing or Unexpected Keys

When loading a state_dict into a model, PyTorch validates that the keys match exactly. If you try to load a checkpoint from a different architecture, or if the model definition changed, you will get an error listing missing and unexpected keys. The load_state_dict method has a strict parameter, set to True by default.

# Non-strict loading: ignore missing or unexpected keys
model.load_state_dict(torch.load('weights.pth'), strict=False)

Using strict=False is useful when you are fine-tuning a model with a modified head, or when you want to load only a subset of layers. However, it silently ignores mismatches, which can lead to subtle bugs if you are not careful. A safer approach is to load with strict=True and explicitly handle the missing keys by inspecting the error message.

For transfer learning, you often need to load only the backbone weights. You can filter the state_dict before loading:

pretrained_dict = torch.load('backbone.pth')
model_dict = model.state_dict()
# Filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

This pattern lets you reuse pretrained layers while keeping the new layers randomly initialized.

Checkpointing Strategies for Long Training Runs

Saving checkpoints at regular intervals is critical for long training runs. A common practice is to save every N epochs, and also keep the best model based on a validation metric. This protects against data center failures and lets you revert to a known-good state.

best_loss = float('inf')
for epoch in range(num_epochs):
    train_one_epoch()
    val_loss = validate()
    
    # Save latest checkpoint
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'val_loss': val_loss,
    }, 'checkpoint_last.pt')
    
    # Save best model separately
    if val_loss < best_loss:
        best_loss = val_loss
        torch.save(model.state_dict(), 'best_model.pth')

Storing the best model as a plain state_dict is convenient for inference because you do not need the optimizer or epoch. The full checkpoint is for resuming training. Separating the two keeps the inference artifact small and avoids loading unnecessary metadata.

When disk space is a concern, you can overwrite the latest checkpoint and keep only the best model. Some training frameworks also use cyclic checkpoints, saving every K epochs and deleting older ones. The tradeoff is between recovery granularity and storage cost.

Another consideration is the file format. torch.save uses Python's pickle, which is not secure against maliciously crafted files. Only load checkpoints from trusted sources. For large models, you can use torch.save with _use_new_zipfile_serialization=True (the default in recent versions) to produce a zip archive that is more portable and easier to inspect.

Finally, when saving checkpoints on distributed training, ensure that only one process writes to disk to avoid file corruption. The common pattern is to save only on the main process, or to use a temporary file and rename it atomically after writing.

By understanding the difference between state_dict and full checkpoints, and by following these loading and saving patterns, you can build training pipelines that are robust, reproducible, and easy to resume.