Building Models with nn.Module

By the end of this lesson you will be able to build any PyTorch model as a class, from a simple stack of layers to a network with its own custom logic. You will know how PyTorch finds every weight inside your model, and you will avoid the classic bug where some layers are silently never trained.

In NumPy, our network was a dictionary of arrays plus a few functions. PyTorch organises the same things with nn.Module, the base class of every layer and every model. Learn its few rules and you can read almost any PyTorch project.

Step 1: your first model class

A model is a class that inherits from nn.Module and does two things:

1. In __init__, it creates its layers and stores them as attributes (self.something = ...).
2. In forward, it says how an input flows through those layers.

Think of a tiffin box. __init__ decides which compartments it has: rice, dal, sabzi. forward is the order in which you eat from them. The box itself keeps track of everything inside it.

import torch
from torch import nn

class DigitNet(nn.Module):
    def __init__(self, n_in=64, n_hidden=128, n_classes=10):
        super().__init__()                      # always call this first
        self.hidden = nn.Linear(n_in, n_hidden)
        self.act = nn.ReLU()
        self.out = nn.Linear(n_hidden, n_classes)

    def forward(self, x):
        x = self.act(self.hidden(x))
        return self.out(x)                       # raw logits, no softmax

torch.manual_seed(0)
model = DigitNet()
print(model)
x = torch.randn(5, 64)                           # a batch of 5 fake digits
logits = model(x)                                # calls forward for you
print("logits shape:", logits.shape)
print("parameters:", sum(p.numel() for p in model.parameters()))

DigitNet(
  (hidden): Linear(in_features=64, out_features=128, bias=True)
  (act): ReLU()
  (out): Linear(in_features=128, out_features=10, bias=True)
)
logits shape: torch.Size([5, 10])
parameters: 9610

Three things to notice:

• super().__init__() sets up the machinery that tracks layers. Forget it and PyTorch raises an error when you assign the first layer.
• You call model(x), not model.forward(x). Calling the model runs some extra PyTorch bookkeeping (such as hooks) and then your forward.
• forward returns raw logits, because nn.CrossEntropyLoss applies softmax itself, as you learned in Loss Functions.

The parameter count is 64×128 + 128 + 128×10 + 10 = 9,610, exactly the formula from Layers & the Forward Pass.

nn.Sequential for simple stacks

When the data simply flows through layers one after another, nn.Sequential saves you writing a class:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(128, 10),
)
print(model)
print("first layer:", model[0])
print("output:", model(torch.randn(2, 64)).shape)

Sequential(
  (0): Linear(in_features=64, out_features=128, bias=True)
  (1): ReLU()
  (2): Dropout(p=0.2, inplace=False)
  (3): Linear(in_features=128, out_features=10, bias=True)
)
first layer: Linear(in_features=64, out_features=128, bias=True)
output: torch.Size([2, 10])

You can index a Sequential like a list: model[0] is the first layer. Use Sequential for straight stacks, and write your own class when forward needs more: two inputs, a skip connection, a loop, or an if.

Step 2: how PyTorch finds your weights

model.parameters() returns every learnable tensor in the model. The optimiser gets this list, so a weight that is missing from it is never trained. PyTorch finds parameters by looking at the attributes you set in __init__:

• An attribute that is an nn.Module (like nn.Linear) is registered, along with everything inside it.
• An attribute that is an nn.Parameter (a tensor marked as learnable) is registered.
• A plain Python list of layers is not registered. This is the classic bug. Use nn.ModuleList instead.

import torch
from torch import nn

class Broken(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(16, 16) for _ in range(3)]     # plain list: invisible!

class Fixed(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(16, 16) for _ in range(3)])
        self.scale = nn.Parameter(torch.ones(16))                # a learnable vector

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x * self.scale

print("Broken parameters:", sum(p.numel() for p in Broken().parameters()))
print("Fixed parameters: ", sum(p.numel() for p in Fixed().parameters()))
for name, p in Fixed().named_parameters():
    print(f"  {name:<16} {tuple(p.shape)}")

Broken parameters: 0
Fixed parameters:  832
  scale            (16,)
  layers.0.weight  (16, 16)
  layers.0.bias    (16,)
  layers.1.weight  (16, 16)
  layers.1.bias    (16,)
  layers.2.weight  (16, 16)
  layers.2.bias    (16,)

The broken model reports zero parameters. Its layers would still run in forward, but the optimiser would never update them, and you would see a model that does not learn without any error message. named_parameters() is the quickest way to check what your model really contains. The same rule applies to dictionaries: use nn.ModuleDict, not a plain dict.

Step 3: modes, devices and the state dict

Every module carries a few more things you will use daily:

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 8), nn.BatchNorm1d(8), nn.ReLU(), nn.Dropout(0.5), nn.Linear(8, 2))
print("training mode?", model.training)
model.eval()
print("after eval():", model.training, "| dropout layer:", model[3].training)
for key, value in model.state_dict().items():
    print(f"  {key:<28} {tuple(value.shape)}")

training mode? True
after eval(): False | dropout layer: False
  0.weight                     (8, 4)
  0.bias                       (8,)
  1.weight                     (8,)
  1.bias                       (8,)
  1.running_mean               (8,)
  1.running_var                (8,)
  1.num_batches_tracked        ()
  4.weight                     (2, 8)
  4.bias                       (2,)

model.eval() reaches every layer inside, so you call it once on the whole model. The state dict includes the batch norm's running statistics (buffers) as well as the weights. You will save it in Saving & Loading Models.

Modules versus functions

PyTorch offers many operations in two forms: a module such as nn.ReLU() and a function such as torch.relu or torch.nn.functional.relu. Both compute the same thing. A rule of thumb: anything with learnable weights or state (Linear, Conv2d, BatchNorm, Dropout's train/eval switch) should be a module created in __init__. Stateless maths (ReLU, softmax, reshaping) can be a function call in forward.

Step 4: the idea in NumPy

There is nothing magic about parameters(). It walks through the model's attributes and collects every weight it finds, going inside sub-modules. Here is a tiny NumPy version you can run, so the idea is clear.

import numpy as np

class Module:
    def parameters(self):
        found = []
        for value in vars(self).values():
            if isinstance(value, np.ndarray):
                found.append(value)                        # a weight array
            elif isinstance(value, Module):
                found.extend(value.parameters())           # look inside sub-modules
        return found

    def __call__(self, x):
        return self.forward(x)

class Linear(Module):
    def __init__(self, n_in, n_out, rng):
        self.W = rng.normal(0, np.sqrt(2 / n_in), (n_in, n_out))
        self.b = np.zeros(n_out)

    def forward(self, x):
        return x @ self.W + self.b

class DigitNet(Module):
    def __init__(self, rng):
        self.hidden = Linear(64, 128, rng)
        self.out = Linear(128, 10, rng)
        self.notes = "not a weight, ignored"

    def forward(self, x):
        return self.out(np.maximum(0, self.hidden(x)))

rng = np.random.default_rng(0)
model = DigitNet(rng)
print("weight arrays found:", len(model.parameters()))
print("total parameters:", sum(p.size for p in model.parameters()))
print("output shape:", model(rng.normal(size=(5, 64))).shape)

weight arrays found: 4
total parameters: 9610
output shape: (5, 10)

Our toy finds the same 9,610 parameters as PyTorch. Real nn.Module does much more (devices, modes, hooks, saving), but the core idea is this recursive search.

Common mistakes

• Forgetting super().__init__(). You get "cannot assign module before Module.__init__() call".
• Storing layers in a plain list or dict. They are not registered and never train. Use nn.ModuleList or nn.ModuleDict.
• Applying softmax in forward and then using CrossEntropyLoss. Return logits; apply softmax only when you need probabilities for display.
• Creating new layers inside forward. A layer created there gets fresh random weights on every call and is never trained. Create layers in __init__.
• Calling model.forward(x) directly. It works, but skips hooks. Call model(x).

Key points

• A PyTorch model is a class inheriting nn.Module, with layers created in __init__ and the data flow written in forward.
• nn.Sequential builds a straight stack of layers without a class.
• parameters() collects every registered module and nn.Parameter; plain lists are invisible, so use nn.ModuleList.
• model.train(), model.eval(), model.to(device) and model.state_dict() act on the whole model at once.
• Layers with weights or state are modules; stateless operations can be plain functions.

import torch
from torch import nn

class ResidualBlock(nn.Module):
    def __init__(self, n=32):
        super().__init__()
        self.linear1 = nn.Linear(n, n)
        self.linear2 = nn.Linear(n, n)

    def forward(self, x):
        return torch.relu(x + self.linear2(torch.relu(self.linear1(x))))

class ResNetMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.inp = nn.Linear(64, 32)
        self.blocks = nn.ModuleList([ResidualBlock(32) for _ in range(3)])
        self.out = nn.Linear(32, 10)

    def forward(self, x):
        x = torch.relu(self.inp(x))
        for block in self.blocks:
            x = block(x)
        return self.out(x)

torch.manual_seed(0)
model = ResNetMLP()
print("output:", model(torch.randn(4, 64)).shape)
print("parameters:", sum(p.numel() for p in model.parameters()))

output: torch.Size([4, 10])
parameters: 8746