What Data Does PyTorch state_dict Serialize?

In PyTorch, saving model checkpoints typically relies on the state_dict object. This article explores the precise data serialized when saving a PyTorch model's state_dict, explaining the exact components included—specifically learnable parameters and persistent buffers—as well as the critical elements omitted during serialization that developers must handle separately.

The Structure of a state_dict

A PyTorch state_dict is a standard Python dictionary (collections.OrderedDict) that maps each layer or submodule to its corresponding parameter and buffer tensors. Only layers with internal, persistent states are represented in this dictionary.

When you serialize a model's state_dict using torch.save(model.state_dict(), PATH), the file stores:

1. Learnable Parameters

The primary data stored in a state_dict consists of all objects instantiated as torch.nn.Parameter within the model and its child submodules. These parameters update during backpropagation:

• Weights: Tensors representing connection strengths across layers (e.g., conv1.weight, linear.weight, embedding.weight).
• Biases: Tensors containing additive bias terms for those operations (e.g., conv1.bias, linear.bias).

Each tensor is saved along with its current values, data type (dtype), and shape, maintaining the state of the model at the exact moment of saving.

2. Registered Persistent Buffers

Not all stateful values are updated via gradient descent. Models frequently require non-trainable tensors that persist across training and inference iterations. Any tensor registered using model.register_buffer() is serialized into the state_dict by default (unless registered with persistent=False).

Common examples of registered buffers include:

• Batch Normalization Statistics: Running statistics such as running_mean, running_var, and num_batches_tracked.
• Positional Encodings: Static coordinate or sequence grids in Transformer architectures.
• Custom Masks: Fixed binary attention masks or causal masks.

What is NOT Serialized

A common misconception is that a state_dict contains everything required to re-instantiate a model from scratch. The state_dict excludes:

• Model Architecture: The underlying Python class definitions, the computation graph, and the sequence of forward operations are omitted. To reload a state_dict, the exact model class must first be instantiated in Python code before calling model.load_state_dict().
• Optimizer States: Momentum values, gradient histories, and step counts are stored in the optimizer's own state_dict (optimizer.state_dict()), not the model's.
• Standard Python Attributes: Any standard class variable not explicitly registered as a nn.Parameter or buffer (e.g., self.learning_rate = 0.01 or standard Python lists) is ignored.
• Hooks and External State: Forward and backward hooks, training/evaluation modes (self.training), and device assignments (CPU or CUDA) are not preserved.

Standard Complete Checkpoint Practice

Because model.state_dict() captures only weights and persistent buffers, standard PyTorch practice packages the model state alongside operational metadata when saving checkpoints for later resumption:

checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, 'checkpoint.pth')

This comprehensive dictionary ensures that the model parameters, buffer states, optimizer parameters, and training progress markers are fully preserved in a single serialized file.