Course outline

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.

About 4 min readContent reviewed
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.

python Illustrative CPU PyTorch 2.14 API program; torch unavailable, expected output derived
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()))
Expected output
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.

python Executed deterministic standard-library companion
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"]))
Expected output
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.

Official references