Module

Base class for layers: named parameters, buffers and children.

Written in Python rather than C++ on purpose. A Module holds references and iterates dicts, none of which is hot; the hot part is the tensor operators it calls, and those are already in C++.

The structure — named children plus named parameters, recursive with a dotted prefix — exists so parameter names line up with a state_dict for loading and comparison. It mirrors torch.nn closely enough that PyTorch code reads across unchanged.

Assigning a Tensor attribute makes a parameter, and assigning a Module makes a child. That is done in __setattr__, so self.weight = ... registers without any explicit call. register_buffer is the explicit opt-out for state that must travel with the module but must not be trained.

Construction¶

__init__¶

register_buffer¶

Records persistent state that is not trained. A buffer appears in state_dict, so it saves, loads and interoperates with a torch checkpoint — but never in parameters(), so an optimiser cannot see it. Batch normalisation's running statistics are the motivating case: they carry no gradient, and letting an optimiser "train" them would destroy the estimate. Raises if the tensor has requires_grad.

Traversal¶

named_parameters¶

parameters¶

named_buffers¶

named_modules¶

Yields self first, then children depth-first, so a caller can match on the root as well as the leaves.

State¶

state_dict¶

Parameters and buffers, by dotted name. Buffers are included because that is what makes a checkpoint complete: a batch-normalised model restored without its running statistics evaluates against the wrong distribution while looking perfectly healthy.

load_state_dict¶

Copies values in by name, in place, preserving each entry's device, dtype and requires_grad. Every key must match — a missing or unexpected one raises rather than being ignored.

Placement and mode¶

to¶

Moves every parameter and buffer to a device, in place.

Call this before constructing an optimiser. The optimiser captures the parameter list when it is built, and to replaces each parameter with a new tensor — so an optimiser made first would keep updating the old ones while the model used the new. torch has the same ordering constraint for the same reason.

Transfer goes through the host, because that is what a transfer to a discrete device is. Gradients move with their parameters; dropping them would leave a subsequent optimiser step silently updating nothing.

train¶

Sets training mode recursively. Only Dropout and BatchNorm2d behave differently between modes.

eval¶

train(False).

zero_grad¶

Clears every parameter's gradient by assigning an undefined tensor. Required between steps, because backward accumulates.

Interface¶

forward¶

__call__¶

__repr__¶

Attribute plumbing¶

__setattr__¶

Routes a Tensor into _parameters and a Module into _modules, which is what makes self.weight = ... register without a call.

__getattr__¶

Only invoked when normal lookup fails, so parameters, buffers and children resolve without shadowing real attributes.

See also backward, save_module, load_module