Parameters, Buffers, and the State Dict

Two kinds of tensors live inside a Module

Parameters (nn.Parameter) — learnable. Show up in model.parameters(). The optimizer updates them.
Buffers (registered via self.register_buffer(name, tensor)) — non-learnable but stateful. Examples: BatchNorm's running mean/var, position encodings, attention masks, EMA weights of a teacher model.

Both are saved in state_dict() and both move with .to(device). The difference is just whether the optimizer touches them.

state_dict — the universal serialization format

model.state_dict() returns an OrderedDict mapping parameter / buffer names to their tensors. load_state_dict(d) restores them. This is the recommended format for saving models because it's framework-version-tolerant: even if the model class evolves, you can load the weights into a new architecture as long as the names and shapes match (or use strict=False to allow mismatches).

Freezing

To freeze a parameter (so the optimizer doesn't update it), set p.requires_grad = False. This is how transfer learning starts — freeze the pretrained backbone, train only the new head. The optimizer still sees the parameter (it's in model.parameters()) but with no gradient there's nothing to update.

Code

Counting parameters and buffers·python

import torch
import torch.nn as nn

class DemoModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 4)
        self.bn = nn.BatchNorm1d(4)
        # A custom non-learnable buffer
        self.register_buffer('class_centroids', torch.zeros(4, 4))

    def forward(self, x):
        return self.bn(self.fc(x))

m = DemoModel()

# Parameters — learnable
print('PARAMETERS')
for n, p in m.named_parameters():
    print(f"  {n}: shape={tuple(p.shape)}, requires_grad={p.requires_grad}")

# Buffers — non-learnable state
print('BUFFERS')
for n, b in m.named_buffers():
    print(f"  {n}: shape={tuple(b.shape)}")

state_dict — save and load·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP()
torch.save(model.state_dict(), '/tmp/model_weights.pt')

# Load
model2 = TinyMLP()
sd = torch.load('/tmp/model_weights.pt', weights_only=True)
model2.load_state_dict(sd)

# Confirm
for k in sd:
    print(k, torch.equal(sd[k], model2.state_dict()[k]))

Freezing for transfer learning·python

import torch.nn as nn
from torchvision.models import resnet18, ResNet18_Weights

model = resnet18(weights=ResNet18_Weights.DEFAULT)

# Freeze the whole backbone
for p in model.parameters():
    p.requires_grad = False

# Replace the head — new params default to requires_grad=True
model.fc = nn.Linear(model.fc.in_features, 5)

# Optimizer should only see the trainable params
import torch.optim as optim
optimizer = optim.AdamW(
    [p for p in model.parameters() if p.requires_grad],
    lr=1e-3,
)
print(f"Trainable params under optimizer: {sum(p.numel() for g in optimizer.param_groups for p in g['params']):,}")

Parameters, Buffers, and the State Dict

Two kinds of tensors live inside a Module

state_dict — the universal serialization format

Freezing

Code

External links

Exercise

Progress

Comments 0