Visual walkthrough — Building models with nn.Module

Before we begin: everything here builds on the parent's Lego-brick idea — each nn.Module is a self-contained brick. We are going to zoom into a single brick, see what lives inside it, then snap bricks together and watch a collector walk through them.

---

Step 1 — What is a parameter, and what is a plain tensor?

WHAT. A tensor is just a grid of numbers — a list, a table, or a stack of tables. Nothing more. In PyTorch you make one with torch.randn(3) (three random numbers in a row).

A parameter is a special tensor: a plain tensor wearing a name tag that says "train me." In code it is nn.Parameter(torch.randn(3)). Same numbers inside — one extra flag on the outside.

WHY this distinction and not just "all tensors." A neural network holds two kinds of numbers:

• the numbers it is trying to learn (weights, biases) — these must be updated during training,
• the numbers that are just passing through (your input image, an intermediate result) — these must not be updated.

If PyTorch trained everything, it would try to "improve" your input photo instead of the network. So we need a tag that separates the two.

PICTURE. On the left, a plain grid of numbers. On the right, the same grid with a bright "train me" tag pinned to it. The tag is the only difference.

---

Step 2 — What happens the instant you write self.weight = ...

WHAT. Inside __init__ you write a line like:

$\underset{\text{this brick}}{\underbrace{\text{self}}}.\underset{\text{the name}}{\underbrace{\text{weight}}}=\underset{\text{a train-me tensor}}{\underbrace{\text{nn.Parameter}(\dots )}}$

Reading it term by term: self is this particular module object; weight is the label you chose; the right-hand side is a parameter (the train-me tensor from Step 1).

In an ordinary Python object this line would just glue the value onto the object and forget about it. But nn.Module does something extra.

WHY the extra step is needed. Later, a collector will need to find this parameter. If assignment did nothing special, the parameter would be an anonymous attribute lost among dozens of others. So nn.Module intercepts every assignment and checks the type of what you're storing.

PICTURE. An assignment arrow enters a checkpoint (a gate). The gate inspects the package: "Is this a Parameter? a sub-Module? or ordinary?" and drops it into the matching drawer inside the module.

---

Step 3 — A single brick's own drawer

WHAT. Take the simplest possible module: one linear layer that computes $y=xW+b$. Its __init__ stores two parameters, weight (that is $W$) and bias (that is $b$). After construction, its _parameters drawer literally holds:

$\text{\_parameters}=\{\underset{\text{the name}}{\underbrace{\text{"weight"}}}:W,\underset{\text{the name}}{\underbrace{\text{"bias"}}}:b\}$

Each entry is a name pointing at a train-me tensor.

WHY store names, not just a list. Names let you save and reload a model reliably — when you load a file, PyTorch matches numbers back to slots by name. A bare list would lose track of which grid is which.

PICTURE. One brick, its lid open, showing a small drawer with two labelled slots: weight → W, bias → b. Nothing else is in the drawer.

For this lonely brick, asking "give me all your parameters" is trivial: just read out the drawer, $[W,b]$. This is our base case — a brick with no sub-bricks returns only its own parameters. Remember that phrase; the next step needs it.

---

Step 4 — Bricks inside bricks (the tree)

WHAT. Real networks nest. Your SimpleClassifier stores three linear layers:

self.fc1 = nn.Linear(...)   # a whole Module, goes in _modules
self.fc2 = nn.Linear(...)   # a whole Module, goes in _modules
self.fc3 = nn.Linear(...)   # a whole Module, goes in _modules

Each nn.Linear is itself a module (with its own weight/bias drawer). So by the gate rule from Step 2, these go into the parent's _modules drawer, not its _parameters drawer. The parent's own _parameters drawer is empty — it owns no loose parameters, only children.

WHY this makes a tree. A module contains modules, which contain modules… This is exactly the parent note's hierarchical function $f(x)=f_3(f_2(f_1(x)))$ turned into a data structure. Drawing the containment gives a tree: the classifier at the root, three linear bricks as branches, each branch holding two leaf tensors.

PICTURE. A tree. Root node SimpleClassifier (empty own-parameter drawer). Three child nodes fc1, fc2, fc3. Under each child, two leaves: its weight and bias.

---

Step 5 — The collector walks the tree (recursion)

WHAT. .parameters() is a walker. Standing at any module it does exactly two things:

$\text{collect}(M)=\underset{\text{my own train-me tensors}}{\underbrace{M.\text{\_parameters}}}\cup {\bigcup }_{c\in M.\text{\_modules}}\underset{\text{ask each child the same question}}{\underbrace{\text{collect}(c)}}$

Term by term: $M.\text{\_parameters}$ is this node's own drawer (Step 3); the big union symbol $\bigcup$ means "gather together"; $\text{collect}(c)$ is the walker calling itself on every child $c$. That self-call is recursion — solving a problem by asking a smaller copy of yourself the same question.

WHY recursion is the right tool here. The tree can be any depth — you decided that at runtime with num_layers. A fixed loop of "layer1 then layer2" cannot handle an unknown depth. Recursion does not care how deep the tree goes: it just says "collect my own, then tell every child to do the same," and the children handle their children, forever, automatically. It is the natural match for a tree of unknown size.

• Base case (Step 3): a leaf module with no children returns only its own drawer. This is where the recursion stops — without it the walk would never end.
• Recursive case: any module returns its own drawer plus whatever its children return.

PICTURE. The same tree, now with a glowing path showing the visit order: root → fc1 → its weight, bias → back up → fc2 → weight, bias → fc3 → weight, bias. Collected tensors drop into an output list at the bottom in the order visited.

Run it on SimpleClassifier(784, 128, 1). The root's own drawer is empty, so the whole list is just the three children's drawers, in order:

$[W_1,b_1,W_2,b_2,W_3,b_3]$

Counting the numbers inside them (each $W$ has $\text{in}\times \text{out}$ entries, each $b$ has $\text{out}$):

$\underset{100480}{\underbrace{(784\cdot 128+128)}}+\underset{16512}{\underbrace{(128\cdot 128+128)}}+\underset{129}{\underbrace{(128\cdot 1+1)}}=117121$

which matches the parent note's parameter count exactly. The walker found every one and skipped nothing.

---

Step 6 — The edge cases that break naïve versions

A good walker must survive the weird trees, not just the tidy one. Three traps:

Case A — the Python-list trap. If you write self.layers = [nn.Linear(...), nn.Linear(...)], the gate in Step 2 sees a plain Python list, not an nn.Module. So it lands in the "ignored" column: the bricks inside are invisible to the walker, never trained. The fix is nn.ModuleList([...]), which is a module and registers each element into _modules. Same-looking code, completely different tree.

Case B — the empty module. A module with no parameters and no submodules (say, a pure nn.ReLU, which just computes $\max (0,x)$ and learns nothing). Its own drawer is empty and it has no children, so $\text{collect}=\{\}$ — an empty list. The recursion handles this gracefully because the base case allows an empty drawer; nothing crashes.

Case C — the shared parameter. Suppose two slots point at the same tensor object (weight tying — common in language models where input and output embeddings share weights). The walker would meet it twice. PyTorch's real walk therefore keeps a "seen" set and yields each tensor once, so an optimizer does not update the same numbers twice.

PICTURE. Three mini-trees side by side: (A) a broken tree where a plain-list branch is greyed-out and unreachable; (B) a tiny tree whose only node is empty, returning []; (C) a tree where two arrows converge on one shared leaf, marked "counted once."

---

The one-picture summary

The whole story in a single frame: an assignment goes through the gate (Step 2) into one of two drawers (Steps 3–4), the drawers form a tree, and the recursive collector (Step 5) walks that tree — surviving the empty, the list-trap, and the shared-weight edge cases (Step 6) — to hand the optimizer exactly the set of train-me tensors.

Recall Feynman retelling — say it back in plain words

Every layer in PyTorch is a little box. Inside __init__, whenever you set self.something = ..., a gatekeeper looks at what you're storing. If it's a "train-me" tensor, it drops it in the parameters drawer. If it's another box, it drops it in the submodules drawer. Everything else it ignores. So each box ends up holding its own train-me tensors and a list of child boxes — which makes a family tree of boxes. When you finally ask the top box "give me everything trainable," it hands over its own drawer and politely asks every child box the exact same question; each child asks its children, all the way down to boxes with no children (that's where the asking stops). Collect all the answers and you have every weight in the network — none missed, none doubled — ready for the optimizer. The one thing you must never forget: call super().__init__() first, or the gatekeeper was never hired, and everything falls apart. And never hide layers in a plain Python list — the gatekeeper won't recognise it, and those layers silently never learn.

Recall Quick self-test

What decides which drawer a value goes into? ::: Its type — nn.Parameter → _parameters, nn.Module → _modules, anything else is ignored. Why does .parameters() use recursion instead of a fixed loop? ::: The module tree can be any depth (chosen at runtime), and recursion naturally walks a tree of unknown size. What is the base case that stops the recursion? ::: A module with no submodules returns only its own parameter drawer. Why does self.layers = [nn.Linear(...)] silently fail? ::: A plain list is not an nn.Module, so the gate ignores it and those layers never get registered or trained — use nn.ModuleList. How many parameters does SimpleClassifier(784,128,1) have? ::: $117121$.