Skip to content

Custom Layers & Advanced Containers

Overview

PyTorch lets you write layers as plain Python classes — but the framework discovers parameters and submodules through attribute inspection, not through your code. Learning the discovery rules lets you build clean, correct, composable layers instead of fighting the framework.

  • nn.Module subclasses get params/submodules from attributes assigned in __init__.
  • Plain Python containers (list, dict) are invisible to the module tree → use ModuleList/ModuleDict/ParameterList.
  • forward is just Python; any logic is allowed, but the graph comes from tensor ops.
  • Reuse primitive ops (F.conv2d, F.linear) — you rarely need to implement raw kernels in Python.

💡 If a param "disappears" from model.parameters() after serialization, it was never registered — check how you store it.

Anatomy of a Custom Layer

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

class DenseBlock(nn.Module):
    def __init__(self, in_ch, out_ch, kernel=3, stride=1):
        super().__init__()
        # REGISTERED as submodule -> device moves & state_dict follow automatically
        self.conv = nn.Conv2d(in_ch, out_ch, kernel, stride=stride, padding=kernel // 2)
        self.bn = nn.BatchNorm2d(out_ch)
        # Plain tensor would NOT be a param:
        self.alpha = nn.Parameter(torch.tensor(1.0))  # <- registered leaf param

    def forward(self, x):
        return self.alpha * F.relu(self.bn(self.conv(x)))

model = DenseBlock(3, 8)
print("params:", [n for n, _ in model.named_parameters()])

The Discovery Rules (the "module tree")

class Naive(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(4, 4) for _ in range(2)]   # 🚫 invisible!
        self.also_lost = {"a": nn.Linear(4, 4)}             # 🚫 invisible!

class Fixed(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
        self.also_kept = nn.ModuleDict({"a": nn.Linear(4, 4)})

naive = Naive(); fixed = Fixed()
print("Naive params   :", sum(p.numel() for p in naive.parameters()))  # 0
print("Fixed params   :", sum(p.numel() for p in fixed.parameters()))  # 2*20

Parameter vs plain tensor

Storage Found by .parameters() Tracked by state_dict Grad by default
nn.Parameter(t) ✅ ✅ ✅
plain tensor attr ❌ ❌ (unless register_buffer) ❌ 
class BufferExample(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("running_mean", torch.zeros(3))  # moved with .to(), saved, no grads

b = BufferExample()
b.to('meta')  # buffers follow device moves too
print(b.running_mean.device)

Containers: ModuleList, ModuleDict, Sequential

m = nn.ModuleList([
    nn.Linear(4, 8),
    nn.Linear(8, 2),
])
# index access, iterate, keep tree intact
out = m[1](m[0](torch.randn(3, 4)))

md = nn.ModuleDict({"up": nn.Linear(4, 8), "down": nn.Linear(8, 4)})
print(md["up"])

When Sequential isn't enough

nn.Sequential hardwires order; build forward manually for branching/skipping/attention masks:

class Branchy(nn.Module):
    def __init__(self):
        super().__init__()
        self.trunk = nn.Sequential(nn.Linear(4, 8), nn.ReLU())
        self.side = nn.Linear(4, 8)
        self.head = nn.Linear(8, 2)

    def forward(self, x):
        return self.head(self.trunk(x) + self.side(x))  # residual branch

Init Discipline — Where & When

class W(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.w = nn.Parameter(torch.empty(*dims))
        self.reset_parameters()          # call in __init__ right after alloc

    def reset_parameters(self):
        nn.init.xavier_uniform_(self.w)
        # mirror torch's own pattern (Linear.reset_parameters)

⚠️ Avoid lazy init inside forward — it breaks export, device moves, and state_dict round-trips. (Except nn.LazyLinear etc., which manage the bookkeeping.)

Building Blocks Best Practices

  1. Compose via small modules, not giant forward methods, to keep named_modules() usable.
  2. Use F.* ops for pointwise/activation logic — no params, no state.
  3. Keep non-tensor config as plain attributes (self.kernel = kernel) — they don't need registration.
  4. Override extra_repr() for readable print(model), and __repr__ stays informative.
  5. Never del or reassign params inside forward; assign once in __init__.
class ConvBlock(nn.Module):
    def __init__(self, c, k=3):
        super().__init__()
        self.conv = nn.Conv2d(c, c, k, padding=k // 2)
        self.gn = nn.GroupNorm(min(8, c), c)

    def extra_repr(self):
        return f"groupnorm_groups={self.gn.num_groups}"

    def forward(self, x):
        return F.relu(self.gn(self.conv(x)))

Key Takeaways

  • The module tree is attribute-based: registration is what makes params trainable, movable, serializable.
  • Reaching for nn.ModuleList/Dict/ParameterList instead of Python lists fixes 90% of "missing param" bugs.
  • forward freedom is your friend, but correctness (params, buffers, init) lives in __init__.
  • Prefer F.* primitives; keep classes small and extra_repr informative.
  • Init immediately after allocation, not lazily during forward.