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.

-

  • Weight Initialization
  • [Parameter Sharing & Weight Tying](/06-pytorch/02-module-and-layer-engineering/(03-parameter-sharing-weight-tying/)
  • Hooks