PyTorch / Tensor, gradient, and module foundations / Lesson 3
Modules, parameters, and buffers
Use nn.Module registration to make trainable values, persistent state, and composition explicit.
On this page
Lesson goals
- Register child modules, Parameters, and buffers under clear names.
- Separate state persistence from gradient participation.
- Call the module interface and inspect named state deliberately.
A module is a registered computation boundary
Subclass nn.Module, call super().__init__, register the components that belong to the model, and define forward for its input-to-output computation. Calling model(x) uses the module call path, including supported hook behavior, rather than merely invoking a plain mathematical function. Calling forward directly can bypass that surrounding machinery. A module can contain other modules, giving names and traversal to a hierarchy. That registration is useful for parameter discovery, state dictionaries, dtype or device movement, and mode changes; it does not itself validate feature meaning or establish a model’s statistical usefulness.
Parameters and buffers have different jobs
Assigning nn.Parameter as a module attribute registers a parameter, normally eligible for gradients. register_buffer records a tensor that belongs to model state without making it an optimizer parameter. Persistent buffers enter state_dict by default; a deliberately nonpersistent buffer does not. Running normalization statistics are a common example of nonparameter state, while a fixed centering value makes the distinction easy to inspect here. Freezing a parameter’s gradients does not automatically freeze mutable buffers or mode-dependent behavior. Define separately which named values may change during training, evaluation, adaptation, or inference.
Inspect a complete tiny registered module
The illustrative module has one scalar weight and one stored center. Its forward maps three and four to three and six. Named parameter and buffer traversals expose their roles, and sorted state_dict keys include both. There is no random initialization or external model to obscure the expected calculation. An ordinary tensor attribute can participate numerically without being a registered buffer, so seeing the correct output is weaker evidence than checking registration too. Likewise, a plain Python list of layers does not provide the same child-module registration as ModuleList or another documented container.
import torch
from torch import nn
class CenteredScale(nn.Module):
def __init__(self):
super().__init__()
self.weight = nn.Parameter(torch.tensor(3.0, dtype=torch.float64, device="cpu"))
self.register_buffer("center", torch.tensor(2.0, dtype=torch.float64, device="cpu"))
def forward(self, x):
return self.weight*(x-self.center)
model = CenteredScale()
x = torch.tensor([3.0,4.0], dtype=torch.float64, device="cpu")
print("output", model(x).tolist())
print("parameters", sorted(dict(model.named_parameters())))
print("buffers", sorted(dict(model.named_buffers())))
print("state keys", sorted(model.state_dict()))
output [3.0, 6.0]
parameters ['weight']
buffers ['center']
state keys ['center', 'weight']
A Parameter and a persistent buffer are registered under different roles but both appear in state_dict. Expected registration and tensor results are reasoned, not executed.
Count trainable state separately from complete state
The standard-library companion calculates the same affine result with explicit parameter and buffer dictionaries. It has one trainable scalar but two state entries. This is role-aware bookkeeping only; Python dictionaries do not reproduce Module behavior. In a real module, inspect named_parameters, named_buffers, and state_dict according to their documented scope rather than assuming every stored tensor is optimized. Parameter sharing and duplicate-removal behavior can affect enumeration, so do not infer independent parameter counts from attribute names alone. Preserve a clear mapping from state names to their mathematical and operational meaning.
state = {"parameters":{"weight":3.0}, "buffers":{"center":2.0}}
values = [3.0,4.0]
outputs = [state["parameters"]["weight"]*(x-state["buffers"]["center"]) for x in values]
parameter_count = len(state["parameters"])
print("output", outputs)
print("trainable scalar count", parameter_count)
print("state entries", parameter_count+len(state["buffers"]))
output [3.0, 6.0]
trainable scalar count 1
state entries 2
The companion checks role-aware scalar arithmetic; ordinary dictionaries do not implement module registration, state migration, hooks, or optimization.
Make composition and state changes reviewable
Use registered containers when creating repeated blocks, and choose stable meaningful names for components that will be saved or inspected. Module train and eval states affect only implementations that define mode-specific behavior; they do not turn gradients on or off. Moving a module to another dtype or device should include its registered state, while unrelated Python objects need their own handling. All examples remain CPU-only and no move to an accelerator is requested. The framework snippet was not executed, and its companion proves only the supplied arithmetic and counts, not registration, serialization, hooks, or runtime compatibility.
Common errors and fixes
An optimizer misses a layer stored in a plain list
- Why it happens
- The intended child modules were not registered through a supported container.
- How to fix it
- Use ModuleList, Sequential, or explicit registered attributes as appropriate.
- Prevent it next time
- Inspect expected named parameters before constructing the optimizer.
A fixed tensor is missing from saved state
- Why it happens
- It was an ordinary attribute or intentionally nonpersistent buffer.
- How to fix it
- Register persistent model state with the documented buffer policy.
- Prevent it next time
- Test state keys and known-input behavior together.
A frozen model still changes behavior
- Why it happens
- Only parameter gradients were disabled while buffers or training mode changed.
- How to fix it
- Specify and verify parameter, buffer, and mode policies separately.
- Prevent it next time
- Snapshot relevant named state before and after the controlled operation.
Check your understanding
1. Why does center appear in state_dict without appearing in named_parameters?
Show worked solution
It is registered as a persistent buffer. Its value is needed to reproduce the function, but it is not declared as an optimizer parameter. The weight is a Parameter and belongs to the trainable role. State persistence and gradient participation answer different questions.
2. Would obtaining output [3,6] prove that center was registered correctly?
Show worked solution
No. An unregistered tensor attribute could produce the same arithmetic while being absent from expected state handling. Inspect buffer names and state_dict keys as well as the output. In this course those are derived framework expectations; only the separate dictionary arithmetic was executed.
Key takeaways
- Registration controls discoverable model structure and state.
- Persistent state is broader than trainable parameters.
- Inspect names, behavior, and mode policies together.