Skip to main content

Module 3 — nn.Module: structuring a model

Modules 1 and 2 handled raw tensors and their gradients. Deep learning is more than that: a real model contains dozens of learnable tensors and a fixed order in which to apply operations to them. nn.Module is the object that ties those together — parameters, sub-modules and the forward function — and it is where the Fashion-MNIST classifier finally takes shape. Every abstraction in the rest of the course (DataLoader, optimisers, checkpoints, TorchScript) reads and writes an nn.Module; investing ten minutes to understand its rules pays for itself many times over.

The contract: __init__ and forward​

Every model subclasses nn.Module and implements two methods.

import torch
import torch.nn as nn
import torch.nn.functional as F

class FashionMLP(nn.Module):
    def __init__(self, hidden=256, num_classes=10):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.fc3 = nn.Linear(hidden, num_classes)
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        x = x.view(x.size(0), -1)      # flatten 28x28 to 784
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        return self.fc3(x)             # logits, no softmax

__init__ registers sub-modules and parameters. forward describes how a batch flows through them. That is the whole contract. Everything else — parameters(), state_dict(), .to(device), .train(), .eval() — is derived automatically from those two methods, and specifically from the objects you assign to self.

The registration mechanism, explained once​

This is the point most tutorials rush over, and it explains half of the surprises that follow.

When you write self.fc1 = nn.Linear(...), nn.Module intercepts the assignment via __setattr__. It detects that the right-hand side is another nn.Module, stores it in an internal dictionary, and thereafter recurses into it whenever it needs parameters, buffers or state. The same trick works for nn.Parameter objects, which are simply tensors marked as trainable.

The direct consequence: putting learnable weights in a plain list breaks everything.

# BROKEN: the layers are not registered
self.layers = [nn.Linear(10, 10) for _ in range(3)]

# CORRECT: nn.ModuleList delegates to __setattr__
self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(3)])

A Python list holding modules is invisible to parameters(). Training will run, gradients will be computed on those tensors, but the optimiser will never see them and their weights will stay at their initialisation. The symptom looks like a broken initialisation: accuracy hovers around chance and never moves. nn.ModuleList and nn.ModuleDict are the two collections that solve this, and they exist for exactly this reason.

Naked lists and dicts kill training silently

The same rule applies to nn.Parameter: a list of parameters must live in nn.ParameterList, not a Python list. Any time you find yourself iterating to create modules dynamically, reach for nn.ModuleList or nn.Sequential, never [].

parameters() and state_dict(): what the model exposes​

Once registration is correct, two methods matter for the rest of the course.

parameters() returns an iterator over every learnable tensor in the module and its descendants, in declaration order. It is what the optimiser consumes:

model = FashionMLP()
optimiser = torch.optim.Adam(model.parameters(), lr=1e-3)
print(sum(p.numel() for p in model.parameters()))   # total parameter count

state_dict() returns an ordered dictionary mapping each parameter and buffer to its tensor value. It is the object that serialises the model:

torch.save(model.state_dict(), "fashion_mlp.pt")
# ... later, in another script
model = FashionMLP()
model.load_state_dict(torch.load("fashion_mlp.pt", weights_only=True))

Loading requires that the target model has the exact same structure as the source. load_state_dict compares keys and raises on any mismatch, which is exactly what you want — a silent partial load would leave half the network at its random initialisation. strict=False exists for transfer learning and is the topic of module 9.

nn.Sequential versus a dedicated class​

For strictly linear networks, nn.Sequential avoids writing a class:

model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.3),
    nn.Linear(256, 256), nn.ReLU(),
    nn.Linear(256, 10),
)

The result is a real nn.Module: parameters, state_dict, .to(device) all work. The forward pass is generated automatically as a chain. This is the fastest way to sketch a first architecture and is what module 4 will feed with the Fashion-MNIST DataLoader.

The dedicated class becomes necessary as soon as one of three things happens: two paths diverge and reunite, a tensor from an earlier layer is reused (skip connections), or the forward code contains an if. nn.Sequential cannot express any of those, and forcing it to leads to unreadable nested constructs.

Buffers: state that is not a parameter​

Some tensors belong to a module but must not receive gradients: the running mean and variance of a BatchNorm2d, a fixed positional encoding, a mask. Those are buffers, registered with register_buffer and included in state_dict — so they survive save and load — but excluded from parameters().

class Normalised(nn.Module):
    def __init__(self, mean, std):
        super().__init__()
        self.register_buffer("mean", torch.tensor(mean))
        self.register_buffer("std", torch.tensor(std))

    def forward(self, x):
        return (x - self.mean) / self.std

A buffer follows the module through .to(device) and appears in checkpoints, which is precisely what a NumPy attribute would not do. When in doubt between a plain attribute and a buffer: if the tensor needs to move to the GPU with the model and appear on load, it is a buffer.

Initialisation: default is usually fine, until it isn't​

Every layer in torch.nn initialises its weights with a reasonable scheme — Kaiming uniform for Linear and Conv2d, ones and zeros for BatchNorm. For a Fashion-MNIST MLP with ReLU activations, the default is fine, and a custom initialisation is a distraction.

Custom schemes matter in three situations: very deep networks without normalisation, unusual activations (SELU, GELU with residual paths), or reproducing a paper. The pattern is a walk over sub-modules:

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
        nn.init.zeros_(m.bias)

model.apply(init_weights)

apply recurses over every sub-module, which is why it needs the isinstance guard: it visits the model itself, all nn.Linear layers, all activations, everything.

Read the sub-module tree once

Call print(model) on any fresh model. It prints an indented tree of every registered sub-module and its shapes. Doing this once at construction time catches mismatched layer sizes, forgotten dropouts and any layer that ended up in a plain list.

In summary​

  • An nn.Module registers sub-modules and parameters through __setattr__; a plain Python list of layers is invisible to parameters() and to state_dict(), and training silently ignores those weights.
  • parameters() feeds the optimiser, state_dict() feeds save and load; both walk the module tree automatically once registration is done properly.
  • nn.Sequential is enough for strictly linear networks; a dedicated class becomes necessary as soon as forward code has branches, skip connections or conditions.
  • Buffers carry non-trainable state — running statistics, masks, constants — through .to(device) and checkpoints, unlike plain attributes.

Next module: Dataset and DataLoader, the pair that turns 60 000 Fashion-MNIST files on disk into the batches this network is going to consume.