nn.Module and Model Definition

torch.nn.Module is the base class for all neural network components. A module encapsulates:

1. Parameters — learnable tensors registered via nn.Parameter
2. Sub-modules — child nn.Module instances (set as attributes)
3. Forward logic — the forward() method defining the computation

import torch
import torch.nn as nn

class LinearModel(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)

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

Calling model(x) invokes model.forward(x) through __call__, which also runs registered hooks.

nn.Parameter is a Tensor subclass that, when assigned as a module attribute, is automatically added to the list of parameters:

$$\theta = \{\mathbf{W}_1, \mathbf{b}_1, \mathbf{W}_2, \mathbf{b}_2, \ldots\}$$

class ManualLinear(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.W = nn.Parameter(torch.randn(d_out, d_in))
        self.b = nn.Parameter(torch.zeros(d_out))

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

nn.Sequential chains modules so each output feeds the next:

$$\hat{\mathbf{y}} = f_L \circ f_{L-1} \circ \cdots \circ f_1(\mathbf{x})$$

mlp = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
)

For named layers, use OrderedDict:

from collections import OrderedDict
mlp = nn.Sequential(OrderedDict([
    ("fc1", nn.Linear(784, 256)),
    ("act", nn.ReLU()),
    ("fc2", nn.Linear(256, 10)),
]))

Activation functions introduce non-linearity:

$$\text{ReLU}(x) = \max(0, x), \qquad \text{Sigmoid}(x) = \frac{1}{1 + e^{-x}}, \qquad \text{Tanh}(x) = \frac{e^x - e^{-x}}{e^x + e^{-x}}$$

GELU (used in transformers): $$\text{GELU}(x) = x \cdot \Phi(x)$$

where $\Phi$ is the standard Gaussian CDF.

In PyTorch: nn.ReLU(), nn.Sigmoid(), nn.GELU(), or their functional forms F.relu(x), torch.sigmoid(x), etc.

Proper initialisation prevents vanishing/exploding activations. For a layer with fan-in $n_{\text{in}}$ and fan-out $n_{\text{out}}$:

Kaiming (He) for ReLU networks: $$W_{ij} \sim \mathcal{N}\!\left(0, \frac{2}{n_{\text{in}}}\right)$$

Xavier (Glorot) for tanh/sigmoid: $$W_{ij} \sim \mathcal{N}\!\left(0, \frac{2}{n_{\text{in}} + n_{\text{out}}}\right)$$

nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')
nn.init.xavier_uniform_(layer.weight)
nn.init.zeros_(layer.bias)

Hooks let you inspect or modify data flowing through modules:

• Forward hook: module.register_forward_hook(fn) — called after forward(), receives (module, input, output)
• Backward hook: module.register_full_backward_hook(fn) — called during backward pass, receives (module, grad_input, grad_output)

activations = {}
def save_activation(name):
    def hook(module, input, output):
        activations[name] = output.detach()
    return hook

model.fc1.register_forward_hook(save_activation('fc1'))