PyTorch: Zero to Hero — Step 04

In the previous post we built a neural network using raw tensors and autograd.

We had weights.

We had biases.

We had a forward calculation.

We had a loss.

We called backward().

We updated the parameters.

Nothing essential was missing from the mathematics of training.

Now we are going to add nn.Module.

But I do not want to introduce it as another PyTorch class whose methods you need to memorize.

I want to use it to show you what I think is the most important organizing idea in PyTorch:

Complicated PyTorch systems are built by putting simple PyTorch objects inside other PyTorch objects, and PyTorch knows how to walk the resulting structure.

Understand that one idea and a surprising amount of the framework becomes predictable.

Why does model.parameters() find parameters buried several layers deep?

Why does model.to("cuda") move tensors you never mentioned explicitly?

Why does state_dict() produce names such as encoder.blocks.3.attention.q_proj.weight?

Why does model.eval() affect dropout layers hidden deep inside a transformer?

Why does nn.ModuleList work when a normal Python list can silently break registration?

Why is nn.Parameter different from a tensor with requires_grad=True?

These can look like unrelated PyTorch rules.

They are not.

They are consequences of the same structure.

The idea in one picture

Suppose we eventually build a transformer.

It may look intimidating from the outside:

Transformer
    ↓
TransformerBlock
    ↓
Attention + MLP
    ↓
Linear layers
    ↓
Parameters
    ↓
Tensors

The transformer is not made from a fundamentally new kind of object.

It is smaller objects composed into larger objects.

A block contains modules.

Those modules contain other modules.

Eventually we reach parameters.

Those parameters are tensors.

The same pattern repeats at every scale.

That is the first recursion we are going to understand.

There is also a second structure.

When the model runs, operations depend on earlier operations:

parameters
    ↓
linear operation
    ↓
activation
    ↓
next layer
    ↓
prediction
    ↓
loss

Autograd records the dependency graph needed to propagate gradients backward through that computation.

So PyTorch gives us two related ways to build complexity from simple pieces:

MODEL STRUCTURE                  COMPUTATION STRUCTURE

Model                            loss
  ↓                                ↑
Block                            operation
  ↓                                ↑
Layer                            operation
  ↓                                ↑
Parameter                        operation
  ↓                                ↑
Tensor                           parameters

The first structure answers:

What belongs to this model?

The second answers:

How did this result depend on those tensors?

A great deal of PyTorch follows from those two questions.

1. Start with one layer

Let us begin with the smallest useful module.

import torch
import torch.nn as nn

layer = nn.Linear(4, 3)

print(layer)

You should see something like:

Linear(in_features=4, out_features=3, bias=True)

A linear layer contains learnable tensors.

Inspect them:

for name, parameter in layer.named_parameters():
    print(name, parameter.shape)

Output:

weight torch.Size([3, 4])
bias torch.Size([3])

Already we have a hierarchy:

Linear
├── weight: Parameter
└── bias: Parameter

The layer is an nn.Module.

Its weight and bias are nn.Parameter objects.

And each parameter is a tensor.

We can go all the way down:

print(type(layer))
print(type(layer.weight))
print(isinstance(layer.weight, torch.Tensor))

The important point is not the inheritance hierarchy itself.

The important point is that larger PyTorch objects contain smaller PyTorch objects in a form that the framework understands.

2. Put modules inside another module

Now wrap two linear layers in a model.

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()

        self.fc1 = nn.Linear(4, 8)
        self.fc2 = nn.Linear(8, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.relu(x)
        return self.fc2(x)


model = TinyNet()

The model now looks like:

TinyNet
├── fc1: Linear
│   ├── weight: Parameter
│   └── bias: Parameter
│
└── fc2: Linear
    ├── weight: Parameter
    └── bias: Parameter

We did not create a special parameter registry manually.

We wrote:

self.fc1 = nn.Linear(4, 8)
self.fc2 = nn.Linear(8, 2)

Because fc1 and fc2 are modules assigned as attributes of another module, PyTorch registers them as child modules.

Now ask the parent model what it contains:

for name, module in model.named_modules():
    print(f"{name!r:10s} -> {type(module).__name__}")

Typical output:

''         -> TinyNet
'fc1'      -> Linear
'fc2'      -> Linear

Then ask for all parameters:

for name, parameter in model.named_parameters():
    print(name, parameter.shape)

Output:

fc1.weight torch.Size([8, 4])
fc1.bias torch.Size([8])
fc2.weight torch.Size([2, 8])
fc2.bias torch.Size([2])

Notice what happened.

We asked the top-level model for its parameters.

PyTorch walked into fc1 and fc2 and found the parameters inside them.

That is the pattern to remember.

The parent does not need to know the details of the child

This is more powerful than it first appears.

Suppose fc1 were replaced by a much larger component:

class Block(nn.Module):
    def __init__(self, dim):
        super().__init__()

        self.up = nn.Linear(dim, dim * 4)
        self.down = nn.Linear(dim * 4, dim)

    def forward(self, x):
        return self.down(torch.relu(self.up(x)))

Then:

class Model(nn.Module):
    def __init__(self, dim=32):
        super().__init__()

        self.block = Block(dim)
        self.head = nn.Linear(dim, 2)

    def forward(self, x):
        return self.head(self.block(x))

The top-level model does not need special code saying:

look inside block
then find up
then find its weight
then find its bias
then find down
then find its weight
...

It only needs registered structure.

Ask:

model = Model()

for name, _ in model.named_parameters():
    print(name)

and PyTorch can expose names such as:

block.up.weight
block.up.bias
block.down.weight
block.down.bias
head.weight
head.bias

Each dot is effectively another step through the hierarchy.

3. Build deeper without changing the idea

Let us make several blocks.

class TinyBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()

        self.up = nn.Linear(dim, dim * 2)
        self.down = nn.Linear(dim * 2, dim)

    def forward(self, x):
        return x + self.down(torch.relu(self.up(x)))

Now put several of them inside a model:

class TinyDeepModel(nn.Module):
    def __init__(self, dim=16, depth=3):
        super().__init__()

        self.blocks = nn.ModuleList([
            TinyBlock(dim)
            for _ in range(depth)
        ])

        self.head = nn.Linear(dim, 1)

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return self.head(x)

The tree is now:

    graph TD
    M[TinyDeepModel] --> L[blocks: ModuleList]
    M --> H[head: Linear]
    L --> B0[0: TinyBlock]
    L --> B1[1: TinyBlock]
    L --> B2[2: TinyBlock]
    B0 --> U0[up: Linear]
    B0 --> D0[down: Linear]
    B1 --> U1[up: Linear]
    B1 --> D1[down: Linear]
    B2 --> U2[up: Linear]
    B2 --> D2[down: Linear]
  

Ask for parameter names:

model = TinyDeepModel()

for name, parameter in model.named_parameters():
    print(name, tuple(parameter.shape))

You will see names in the same shape as the tree:

blocks.0.up.weight
blocks.0.up.bias
blocks.0.down.weight
blocks.0.down.bias
blocks.1.up.weight
blocks.1.up.bias
blocks.1.down.weight
blocks.1.down.bias
blocks.2.up.weight
blocks.2.up.bias
blocks.2.down.weight
blocks.2.down.bias
head.weight
head.bias

We now have a deeper model.

But conceptually nothing new happened.

We kept composing modules.

That is why the idea scales.

4. nn.Module is a registration system

At this point it is worth becoming precise.

nn.Module does many things, but one of its most important jobs is maintaining registered model structure.

It knows about child modules.

It knows about parameters.

It knows about buffers.

That registered structure powers operations such as:

model.parameters()
model.named_parameters()
model.modules()
model.named_modules()
model.buffers()
model.state_dict()
model.to(device)
model.train()
model.eval()

These APIs become much less mysterious if you picture the module tree.

They are not independent tricks.

They operate on registered model state.

This is why registration bugs are so important.

A tensor can be perfectly valid mathematically and still be invisible to some of this machinery.

5. Autograd can see a tensor that nn.Module cannot see as a parameter

Consider this model:

class BrokenLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()

        self.weight = torch.randn(
            out_features,
            in_features,
            requires_grad=True,
        )

        self.bias = torch.zeros(
            out_features,
            requires_grad=True,
        )

    def forward(self, x):
        return x @ self.weight.T + self.bias

Run a backward pass:

model = BrokenLinear(4, 2)

x = torch.randn(8, 4)
y = torch.randint(0, 2, (8,))

logits = model(x)
loss = nn.functional.cross_entropy(logits, y)
loss.backward()

print(model.weight.grad is None)

You should get:

False

So autograd found the dependency between the loss and weight.

The tensor has a gradient.

Now ask the module for its parameters:

print(list(model.parameters()))

You get:

[]

This looks contradictory until you separate our two structures.

The computation graph knows:

loss depends on weight

But the module registration tree does not know:

weight is a trainable parameter belonging to this model

Those are different questions.

This distinction explains one of the most confusing bugs a new PyTorch programmer can meet:

A tensor can have a gradient and still be missing from model.parameters().

6. nn.Parameter says: this tensor belongs to the model as a parameter

Fix the class:

class ManualLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()

        self.weight = nn.Parameter(
            torch.randn(out_features, in_features) * 0.01
        )

        self.bias = nn.Parameter(
            torch.zeros(out_features)
        )

    def forward(self, x):
        return x @ self.weight.T + self.bias

Now:

model = ManualLinear(4, 2)

for name, parameter in model.named_parameters():
    print(name, parameter.shape)

Output:

weight torch.Size([2, 4])
bias torch.Size([2])

The tensors are now registered as parameters of the module.

That means an optimizer can discover them through:

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.1,
)

The optimizer does not recursively inspect arbitrary Python objects looking for tensors with gradients.

It receives whatever parameters you give it.

A common pattern is:

model.parameters()

So correct model registration becomes part of correct optimization.

A useful mental table

Think about three common kinds of tensor-like model state:

That table is not a collection of arbitrary rules.

It describes different relationships to the module tree.

7. The Python list bug proves why the tree matters

Now we can understand a famous PyTorch trap.

This looks like perfectly reasonable Python:

class BrokenMLP(nn.Module):
    def __init__(self, width=32, depth=4):
        super().__init__()

        self.layers = [
            nn.Linear(width, width)
            for _ in range(depth)
        ]

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

The forward pass works:

model = BrokenMLP()
x = torch.randn(8, 32)

print(model(x).shape)

But now:

print(sum(p.numel() for p in model.parameters()))

returns:

0

Why?

The list contains modules, but the plain Python list is not a registered PyTorch module container.

The layers participate in computation because your Python loop explicitly calls them.

But they are not children in the module tree.

Again our two structures disagree:

COMPUTATION

x
↓
layer 0
↓
layer 1
↓
layer 2

works

while:

MODULE TREE

BrokenMLP
└── layers: ordinary Python list

registered child modules: none

This is why the bug can be so deceptive.

The model runs.

The outputs have gradients.

But model.parameters() cannot see the layers.

8. ModuleList means: these modules are part of my registered structure

Fix it:

class MLP(nn.Module):
    def __init__(self, width=32, depth=4):
        super().__init__()

        self.layers = nn.ModuleList([
            nn.Linear(width, width)
            for _ in range(depth)
        ])

    def forward(self, x):
        for layer in self.layers:
            x = torch.relu(layer(x))
        return x

Now:

model = MLP()

print(sum(p.numel() for p in model.parameters()))

returns a non-zero parameter count.

And:

for name, _ in model.named_parameters():
    print(name)

produces names such as:

layers.0.weight
layers.0.bias
layers.1.weight
layers.1.bias
layers.2.weight
layers.2.bias
layers.3.weight
layers.3.bias

ModuleList is not merely a PyTorch-flavoured list.

It communicates structure to the framework:

The modules stored here belong to this module.

The same idea explains:

nn.ModuleDict
nn.ParameterList
nn.ParameterDict

Use ordinary Python containers when you merely want Python objects.

Use PyTorch registration-aware containers when the contents should belong to the model tree.

9. state_dict() is the module tree turned into names

Now inspect a nested model:

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()

        self.block = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 8),
        )


class Classifier(nn.Module):
    def __init__(self):
        super().__init__()

        self.encoder = Encoder()
        self.head = nn.Linear(8, 2)

    def forward(self, x):
        return self.head(self.encoder.block(x))

Print the state keys:

model = Classifier()

for key in model.state_dict():
    print(key)

Output:

encoder.block.0.weight
encoder.block.0.bias
encoder.block.2.weight
encoder.block.2.bias
head.weight
head.bias

Read those names as paths:

encoder
  ↓
block
  ↓
0
  ↓
weight

That is why checkpoint keys often tell you a great deal about the architecture that produced them.

It is also why changing module structure changes checkpoint compatibility.

Rename:

self.encoder

to:

self.backbone

and the state paths change.

Move a layer from one container to another and the state paths change.

Add or remove modules and the set of keys changes.

Then errors such as:

Missing key(s) in state_dict
Unexpected key(s) in state_dict

stop looking arbitrary.

They are disagreements between two registered structures.

10. Buffers are registered state that is not a trainable parameter

Not every tensor belonging to a model should be optimized.

Suppose a module keeps a running value:

class RunningMean(nn.Module):
    def __init__(self, features):
        super().__init__()

        self.register_buffer(
            "running_mean",
            torch.zeros(features),
        )

    def forward(self, x):
        return x - self.running_mean

Inspect it:

model = RunningMean(4)

print("parameters:", list(model.named_parameters()))
print("buffers:", list(model.named_buffers()))
print("state:", list(model.state_dict().keys()))

You should see:

parameters: []
buffers: [('running_mean', ...)]
state: ['running_mean']

The buffer belongs to the module’s registered state.

But it is not a trainable parameter.

This is useful for things such as:

running statistics
fixed masks
non-learned normalization state
position-related tensors
other state that should follow the module

The distinction becomes obvious when we ask the structural question:

Does this tensor belong to the model, and if so, what kind of state is it?

The answer can be:

trainable state       → Parameter
non-trainable state   → buffer
ordinary temporary data → ordinary tensor

Persistent and non-persistent buffers

Sometimes state should move with the module but does not need to be saved in checkpoints.

class PositionCache(nn.Module):
    def __init__(self, size):
        super().__init__()

        self.register_buffer(
            "cache",
            torch.arange(size),
            persistent=False,
        )

Now:

model = PositionCache(16)

print(list(model.named_buffers()))
print(list(model.state_dict().keys()))

The buffer is registered and visible through named_buffers(), but because it is non-persistent it is omitted from state_dict().

Again, PyTorch gives us a way to express the relationship between an object and the larger module.

11. Why model.to(device) can move a whole model

Suppose we have:

model = TinyDeepModel()

The model may contain many nested layers.

Yet we can write:

model = model.to("cuda")

when CUDA is available.

We do not manually move:

blocks.0.up.weight
blocks.0.up.bias
blocks.0.down.weight
...

PyTorch applies the device conversion through registered parameter and buffer state in the module hierarchy.

That is why registration bugs can become device bugs.

Consider:

class BadState(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = torch.ones(4)

scale is just an ordinary tensor attribute.

Compare it with:

class GoodState(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("scale", torch.ones(4))

In the second case, scale is registered model state and follows module device movement.

So if you ever see an error like:

Expected all tensors to be on the same device

one useful question is not merely:

Which tensor is on the wrong device?

It is:

Was this tensor actually registered as model state, or did I leave it outside the structure PyTorch knows how to move?

That question often gets you closer to the real bug.

12. train() and eval() also propagate through the module hierarchy

Some modules behave differently during training and inference.

Dropout is the classic example.

model = nn.Sequential(
    nn.Linear(8, 8),
    nn.ReLU(),
    nn.Dropout(0.5),
    nn.Linear(8, 2),
)

Call:

model.eval()

You do not need to find the dropout module manually and call eval() on it.

The training mode is propagated through the registered module hierarchy.

Inspect it:

for name, module in model.named_modules():
    print(name, module.training)

After:

model.train()

the modules report training mode.

After:

model.eval()

they report evaluation mode.

This is another reason the hierarchy matters.

A huge model can contain thousands of modules, yet one call at the root can change the mode of descendants.

That is only useful because the descendants are registered children.

13. Freezing a parameter does not remove it from the tree

Take a normal model:

model = TinyNet()

model.fc1.weight.requires_grad_(False)
model.fc1.bias.requires_grad_(False)

Now inspect:

for name, parameter in model.named_parameters():
    print(name, parameter.requires_grad)

The frozen parameters are still there.

They also remain in:

model.state_dict()

Why?

Because these are separate properties.

Registration answers:

Does this parameter belong to the model?

requires_grad answers:

Should autograd calculate gradients for it?

Freezing changes the second answer.

It does not change the first.

Once again, separating the model structure from the computation/gradient structure makes the behavior unsurprising.

14. The forward method is not the module tree

There is another subtle distinction worth making.

The module tree describes registered ownership.

forward() describes computation.

They often line up, but they do not have to be identical.

For example:

class SharedLayerModel(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.layer = nn.Linear(dim, dim)

    def forward(self, x):
        x = torch.relu(self.layer(x))
        x = torch.relu(self.layer(x))
        return x

The module tree contains one Linear module:

SharedLayerModel
└── layer: Linear

But the computation uses that layer twice:

x
↓
layer
↓
ReLU
↓
layer again
↓
ReLU

So do not confuse:

where modules are owned

with:

how data flows during a particular forward pass

This distinction becomes important in recurrent networks, weight sharing and many modern architectures.

The module tree can be static while the computation performed in forward() can be dynamic.

15. Now connect this back to autograd

In Step 02 we looked at autograd.

Suppose we run:

x = torch.randn(8, 16)
y = torch.randn(8, 1)

model = TinyDeepModel(dim=16)

prediction = model(x)
loss = ((prediction - y) ** 2).mean()

loss.backward()

There are now two structures worth keeping in your head.

The registered module hierarchy:

TinyDeepModel
├── blocks
│   ├── block 0
│   │   ├── up
│   │   └── down
│   ├── block 1
│   │   ├── up
│   │   └── down
│   └── block 2
│       ├── up
│       └── down
└── head

and the dynamic dependency graph produced by the actual tensor operations of the forward pass:

parameters
    ↓
linear operations
    ↓
ReLUs
    ↓
residual additions
    ↓
head
    ↓
prediction
    ↓
loss

Then:

model.parameters()

uses the registered model structure to expose trainable parameters.

And:

loss.backward()

uses autograd’s recorded dependency graph to compute gradients for tensors involved in the loss.

Then an optimizer bridges the two:

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3,
)

and later:

optimizer.step()

updates the parameters it was given using their gradients.

This is the whole training mechanism viewed structurally:

    graph LR
    MT[Module tree] -->|model.parameters| P[Registered parameters]
    P --> F[Forward computation]
    F --> L[Loss]
    L -->|backward| G[Gradients]
    G --> O[Optimizer step]
    O --> P
  

The model tree tells us what can be updated.

The autograd graph tells us how the loss depends on it.

The optimizer uses the resulting gradients to make the update.

16. A transformer is the same idea repeated farther

This is where the concept pays off.

When you later meet a transformer implementation, you might see something like:

class TransformerBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()

        self.attention = Attention(dim)
        self.mlp = MLP(dim)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

And then:

class Transformer(nn.Module):
    def __init__(self, dim, depth):
        super().__init__()

        self.blocks = nn.ModuleList([
            TransformerBlock(dim)
            for _ in range(depth)
        ])

Do not read that as:

advanced transformer magic

Read it as:

Transformer
    contains TransformerBlocks

TransformerBlock
    contains Attention, MLP and LayerNorm

Attention
    contains projections

projections
    are Linear modules

Linear modules
    contain Parameters

Parameters
    are Tensors

That is the same pattern we just built with TinyDeepModel.

There may be billions of parameters.

There may be dozens or hundreds of blocks.

There may be specialized attention implementations and complicated forward logic.

But the registered composition principle has not changed.

This is why understanding one small module tree gives you leverage over enormous models.

17. Why this is more useful than memorizing APIs

Suppose you forget exactly what ModuleList does.

If you memorized the API, you are stuck until you look it up.

If you understand the structure, you can reason toward the answer:

I have child modules in a collection.
PyTorch needs to know they belong to the parent.
A plain list is only a Python container.
Therefore I need a registration-aware module container.

Suppose a tensor does not move to CUDA with the rest of the model.

Reason structurally:

model.to(...) operates on registered model state
this tensor did not move
is it a Parameter or registered buffer?

Suppose an optimizer does not update something.

Reason structurally:

optimizer received model.parameters()
is my tensor actually a registered Parameter?

Suppose a checkpoint reports missing keys.

Reason structurally:

state_dict keys reflect registered state paths
did the module hierarchy change?

The mental model generates debugging questions for you.

That is much more valuable than remembering isolated rules.

18. A small debugging toolkit

When a model behaves strangely, inspect the structure directly.

What modules does PyTorch know about?

for name, module in model.named_modules():
    print(name, type(module).__name__)

What parameters does PyTorch know about?

for name, parameter in model.named_parameters():
    print(
        name,
        tuple(parameter.shape),
        "requires_grad=",
        parameter.requires_grad,
    )

Which parameters actually received gradients?

Call backward() first, then:

for name, parameter in model.named_parameters():
    print(
        name,
        "grad=",
        "None" if parameter.grad is None else "set",
    )

These answer different questions.

A parameter can be registered but receive no gradient because the forward computation never used it.

A tensor can receive a gradient but be unregistered because it was an ordinary tensor involved in the computation.

That distinction is worth repeating:

registered?   → module structure question
has gradient? → computation graph question

What persistent state will be saved?

for key, value in model.state_dict().items():
    print(key, tuple(value.shape))

What buffers exist?

for name, buffer in model.named_buffers():
    print(name, tuple(buffer.shape), buffer.device)

Once you inspect the structure rather than guessing, many mysterious PyTorch bugs become ordinary mismatches.

19. The practical registration checklist

If a parameter is not training, ask:

1. Is it an nn.Parameter?
2. Is it attached to a registered module?
3. If it is inside a collection, is that collection ModuleList/ModuleDict/etc.?
4. Does it appear in model.named_parameters()?
5. Was it included in the optimizer?
6. After backward(), does it have a gradient?

Notice the order.

Do not start by tuning the learning rate.

First prove the parameter exists in the structure you think you built.

If a tensor does not move devices, ask:

1. Is it a registered Parameter?
2. Is it a registered buffer?
3. Or is it just an ordinary tensor attribute?

If a checkpoint does not load, ask:

1. What keys are in the checkpoint?
2. What keys are in model.state_dict()?
3. Where do the module paths diverge?

These are all versions of the same debugging method:

Inspect the structure PyTorch sees.

20. A challenge: break the recursion on purpose

Here is a model that contains several bugs:

class BrokenModel(nn.Module):
    def __init__(self, dim=16):
        super().__init__()

        self.input = nn.Linear(dim, dim)

        self.blocks = [
            nn.Linear(dim, dim),
            nn.Linear(dim, dim),
        ]

        self.scale = torch.ones(dim)

        self.extra_weight = torch.randn(
            dim,
            dim,
            requires_grad=True,
        )

    def forward(self, x):
        x = self.input(x)

        for block in self.blocks:
            x = torch.relu(block(x))

        x = x * self.scale
        x = x @ self.extra_weight

        return x

Before changing the code, predict the answers to these questions:

Which tensors can participate in autograd?

Which weights appear in model.parameters()?

Which objects appear in state_dict()?

Which tensors move when model.to(device) is called?

Which tensors would an optimizer built from model.parameters() update?

Then inspect the model:

model = BrokenModel()

print("PARAMETERS")
for name, parameter in model.named_parameters():
    print(name, tuple(parameter.shape))

print("\nBUFFERS")
for name, buffer in model.named_buffers():
    print(name, tuple(buffer.shape))

print("\nSTATE")
for key in model.state_dict():
    print(key)

Now repair the model so that:

• the blocks are registered modules;
• scale is registered non-trainable state;
• extra_weight is a trainable parameter.

A possible repaired __init__ is:

class FixedModel(nn.Module):
    def __init__(self, dim=16):
        super().__init__()

        self.input = nn.Linear(dim, dim)

        self.blocks = nn.ModuleList([
            nn.Linear(dim, dim),
            nn.Linear(dim, dim),
        ])

        self.register_buffer(
            "scale",
            torch.ones(dim),
        )

        self.extra_weight = nn.Parameter(
            torch.randn(dim, dim)
        )

The forward method did not need to change.

The mathematics was already valid.

What changed was our description of ownership and state.

That is exactly the point of this chapter.

21. The one concept to keep

It is easy to come away from nn.Module remembering a list of APIs:

Parameter
ModuleList
register_buffer
state_dict
to
train
eval

Do not do that.

Remember the structure underneath them.

A PyTorch model is recursively composed from registered pieces.

A module can contain modules.

Those modules can contain more modules.

Eventually they contain parameters and buffers.

PyTorch can traverse that hierarchy to expose parameters, move registered state, switch module modes and serialize checkpoints.

Then, separately, the forward pass creates tensor dependencies that autograd can traverse backward from the loss.

Put those ideas together:

REGISTERED MODEL STRUCTURE
        ↓
what belongs to the model?
        ↓
parameters + buffers + submodules

            AND

DYNAMIC COMPUTATION GRAPH
        ↓
how did the loss depend on tensors?
        ↓
gradients

Then training becomes:

find the registered parameters
        ↓
run the computation
        ↓
produce the loss
        ↓
backpropagate through dependencies
        ↓
update the parameters

That is why this idea is so powerful.

Once you see PyTorch as composition plus traversal, many framework features stop looking like exceptions that have to be memorized.

They become consequences of the model you built.

And when something breaks, the first question becomes much better:

What structure does PyTorch actually see?

That question will follow us through the rest of the book.

In the next chapter we move from model structure to another piece of the training system that often looks simple until performance collapses:

getting data into the model.