Building networks: layers, containers and initialisation

Sequential, ModuleList and ModuleDict, wiring branches and shared layers, and why default initialisation is not always what you want.

The three containers

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, sizes):
        super().__init__()
        # ModuleList, not a list: a plain Python list is invisible to .parameters()
        self.blocks = nn.ModuleList([
            nn.Linear(a, b) for a, b in zip(sizes[:-1], sizes[1:])
        ])
        self.norms = nn.ModuleList([nn.LayerNorm(b) for b in sizes[1:]])
        self.act = nn.GELU()

    def forward(self, x):
        for linear, norm in zip(self.blocks, self.norms):
            x = self.act(norm(linear(x)))
        return x

model = MLP([32, 64, 64, 10])
print(sum(p.numel() for p in model.parameters()))   # 9280

# a plain list silently loses everything
broken = nn.Module()
broken.layers = [nn.Linear(4, 4) for _ in range(3)]
print(len(list(broken.parameters())))               # 0

Branches, sharing and weight tying

class TwoTower(nn.Module):
    """Shared trunk, two heads, plus a tied embedding matrix."""

    def __init__(self, vocab=1000, dim=128, n_classes=5):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim)
        self.trunk = nn.Sequential(nn.Linear(dim, dim), nn.ReLU())

        # the SAME layer instance used twice: parameters are shared, not copied
        self.shared = nn.Linear(dim, dim, bias=False)

        self.head = nn.Linear(dim, n_classes)
        self.proj = nn.Linear(dim, dim, bias=False)   # tied to the embedding below
        self.proj.weight = self.embed.weight          # weight tying

    def forward(self, ids):
        x = self.embed(ids).mean(dim=1)
        x = self.trunk(x)
        x = x + self.shared(x)                        # a residual branch
        return self.head(x)

m = TwoTower()
names = dict(m.named_parameters())
print(names["proj.weight"].data_ptr() == names["embed.weight"].data_ptr())  # True: one tensor
print(len(list(m.parameters())))     # tied matrix counted once

• Assigning a module twice shares parameters by reference. That is how weight tying works, and it is also how you accidentally share weights when you meant to copy.
• Use nn.ModuleList for repeated blocks, nn.ModuleDict when the forward pass must choose by a name (an expert mixture, or a task-conditioned head).
• Residual connections need matching shapes. If a downsample changes the channel count, add a 1x1 convolution on the skip path rather than padding with zeros.
• Parameter count is a check, not a goal: print sum(p.numel()) and compare it with your hand calculation to catch a container that is not registered.

Initialisation

import math

def init_weights(module):
    if isinstance(module, nn.Linear):
        # Kaiming for the activation actually used in the network
        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(module, nn.Embedding):
        nn.init.normal_(module.weight, mean=0.0, std=0.02)
    elif isinstance(module, nn.LayerNorm):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model.apply(init_weights)

# a residual block should scale its output projection, or training diverges at depth
class ResidualBlock(nn.Module):
    def __init__(self, dim, scale=0.1):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fc1 = nn.Linear(dim, 4 * dim)
        self.fc2 = nn.Linear(4 * dim, dim)
        self.act = nn.GELU()
        nn.init.zeros_(self.fc2.bias)
        self.fc2.weight.data.mul_(scale)      # shrink the residual branch at init

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

FAQ

Related

Modules, data and saving Debugging PyTorch models

Last refreshed 2026-09-18.