`nn.Module`, layers, and architectural building blocks

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_dim, hidden, out_dim, dropout=0.1):
        super().__init__()
        self.fc1 = nn.Linear(in_dim, hidden)
        self.act = nn.GELU()
        self.drop = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden, out_dim)

    def forward(self, x):
        return self.fc2(self.drop(self.act(self.fc1(x))))

model = MLP(784, 512, 10)
logits = model(torch.randn(32, 784))   # __call__ → forward

nn.Module is the base class for every layer, sub-network, and full model. Its job is to track parameters. When you assign a nn.Linear to self, the module registers its weights and biases; later model.parameters() walks the tree and gives the optimizer every learnable tensor. That's the contract — get it right and everything works; subvert it (e.g. store params in a plain list) and the optimizer silently misses them.

For lists of layers, use nn.ModuleList or nn.Sequential, not a plain list. Use nn.ParameterList / nn.ParameterDict for collections of parameters. For non-learnable tensors you want to ship with the model (e.g. positional encodings), use self.register_buffer('pos', ...) — they move with .to(device) and save with state_dict() but aren't optimized.

• `nn.Linear(in, out)` — fully connected. y = xW^T + b. The workhorse.
• `nn.Conv1d/2d/3d(in, out, kernel)` — convolution. 2d for images, 1d for audio/sequences, 3d for video/volumes.
• `nn.ConvTranspose2d` — 'deconvolution' for upsampling in generators and U-Nets.
• `nn.MaxPool2d(k)` / `nn.AvgPool2d(k)` / `nn.AdaptiveAvgPool2d((1,1))` — downsampling and global pooling.
• `nn.BatchNorm2d(c)` — normalize per-channel across the batch. Standard for vision CNNs.
• `nn.LayerNorm(d)` — normalize per-sample across the feature axis. Standard for transformers.
• `nn.GroupNorm(g, c)` / `nn.InstanceNorm2d` — batch-size-independent alternatives.
• `nn.RMSNorm(d)` (PyTorch 2.4+) — simplified LayerNorm used by LLaMA and friends.
• `nn.Embedding(vocab, d)` — learnable lookup table from token id → vector. First layer of every text model.
• `nn.Dropout(p)` / `nn.Dropout2d(p)` — randomly zeros activations during training. Cheap regularization, expensive if you forget eval() at inference.
• `nn.MultiheadAttention(d, h)` — built-in attention block. Often you write your own for clarity.
• `nn.LSTM(d, h, num_layers)` / `nn.GRU(...)` — recurrent layers. Less common in 2026 but still relevant for streaming / very long sequences.
• Activations: nn.ReLU (vision baseline), nn.GELU (transformers), nn.SiLU/nn.Swish (LLaMA, many modern), nn.LeakyReLU (GANs), nn.Tanh (rarely now), nn.Sigmoid (binary output).

Every interesting architecture is mostly these primitives, composed cleverly.

Stack 50 linear layers naively with default init and your gradient is either zero or NaN by the time it reaches layer 1. Three ideas saved deep learning:

1. Careful initialization. Xavier / Glorot for tanh/sigmoid (var = 1/n_in or 2/(n_in + n_out)). He / Kaiming for ReLU/GELU (var = 2/n_in). nn.Linear and nn.Conv2d use sensible defaults. For custom layers, nn.init.kaiming_normal_(self.weight, nonlinearity='relu').

2. Residual connections out = x + F(x). The gradient gets a 'highway' to early layers because d(out)/dx = I + dF/dx. The I keeps signal flowing no matter how deep. ResNets, transformers, U-Nets — every modern deep architecture is residual at heart. Pre-norm variant (out = x + F(LayerNorm(x))) is the transformer default and trains more stably than post-norm.

3. Normalization layers keep activation statistics under control. BatchNorm uses the current minibatch's mean/var during training and a running average at eval; great for vision but breaks when batch sizes are tiny or sequences are long. LayerNorm normalizes per-sample, batch-size-independent — the choice for transformers and any model where the batch isn't statistically meaningful.

Misuse case: vision people sometimes drop BatchNorm and the model stops converging. Transformer people sometimes add BatchNorm and the model trains slower and worse. Use the norm that matches your data shape.

Some layers behave differently during training and inference: Dropout zeros activations during training and is a no-op at inference; BatchNorm uses batch statistics during training and a running mean/var at inference. The model needs to know which mode it's in:

model.train()        # before the training loop
... # train ...
model.eval()         # before validation / inference
with torch.no_grad():
    preds = model(batch)

Forgetting eval() is one of the top three production deep-learning bugs. Symptom: the model is great on the training metric but flaky on val/prod inputs. Cause: dropout still on, BN using mini-batch stats from a batch of size 1, etc. Make .eval() + with torch.no_grad(): a reflex around inference.

There's a subtler bug: if you do hyperparameter search and forget to set .eval() between trials, BatchNorm running statistics get re-trained on val data — a slow leak that inflates your reported numbers. Always set the mode explicitly at the boundaries.

Knowing 'how big is my model' is a daily skill. The standard one-liner:

n_total = sum(p.numel() for p in model.parameters())
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'{n_trainable/1e6:.1f}M trainable / {n_total/1e6:.1f}M total')

For each layer's contribution:

• nn.Linear(in, out) → in*out + out params.
• nn.Conv2d(c_in, c_out, k) → c_in * c_out * k * k + c_out params (way fewer than equivalent Linear).
• nn.Embedding(V, d) → V * d. Often the biggest layer in a language model.

For a parameter budget check on a 7B LLM: ~7 × 10⁹ params, 7 × 10⁹ × 2 bytes (bf16) = 14 GB just for weights, before grads/optim. That's why a 24 GB consumer GPU can't fine-tune 7B with vanilla AdamW. We'll fix this with LoRA in chapter 12.

torchinfo.summary(model, input_size=(1, 3, 224, 224)) (from the torchinfo package) prints a Keras-style table: layer name, output shape, param count, MACs. Run it on any model you're about to train; it catches misshapen heads and silent param leaks.