D4.3 · LESSON 3 OF 6

Modules and parameters

Translate an inspected NumPy model into tensors, modules, automatic differentiation and a clear training loop.

THE IDEA

Modules own their parameters

A nn.Module is a container for computation that registers what it owns: its nn.Parameter tensors and its child modules. That registration is what lets model.parameters() hand every trainable tensor to the optimizer, state_dict() save them under stable names, and model.eval() reach every layer.

real PyTorch (run it locally)

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()                 # must come first
        self.fc1 = nn.Linear(10, 10)       # assigned as an attribute: registered
        self.fc2 = nn.Linear(10, 3)
    def forward(self, x):
        return self.fc2(self.fc1(x).relu())

A plain tensor you attach to a module is not a parameter, even if it has requires_grad=True. Only nn.Parameter values and child modules are registered, so the optimizer would never see the plain tensor.

Predict first. You store an ordinary tensor on a module as self.scale = torch.ones(3). Will model.parameters() include it?

Counting and naming parameters

Linear(in, out) holds a weight of shape (out, in) and a bias of shape (out,). An activation such as ReLU has no parameters at all. In a Sequential the layers are named by position, so a Linear, ReLU, Linear stack has the keys 0.weight, 0.bias, 2.weight, 2.bias.

hidden width h10

freeze the first layer (requires_grad_(False))

name	layer	weight	bias	parameters	trainable?
0	Linear(10, 10)	100	10	110	yes
1	ReLU()	-	-	0	n/a
2	Linear(10, 3)	30	3	33	yes
total / trainable				143	143

model = nn.Sequential(nn.Linear(10, 10), nn.ReLU(), nn.Linear(10, 3)) Linear(in, out) stores weight of shape (out, in) and bias of shape (out,) sum(p.numel() for p in model.parameters()) → 143 state_dict() keys: 0.weight, 0.bias, 2.weight, 2.bias

A 10 → h → 3 network. Move h, freeze the first layer, and watch the counts and the names.

Chapter context

PyTorch provides tensor operations and differentiation infrastructure. The framework removes bookkeeping, but it does not remove the need to understand dimensions, parameter ownership, device placement and evaluation behavior.

TRY IT YOURSELF

Part 1. Do it hands-on

PyTorch itself cannot run in a browser here, so the Python parts use a small PyTorch-style library that is already loaded. It has the same names and call shapes (torch.tensor, nn.Linear, loss.backward(), optimizer.step(), DataLoader) and the same behaviour for everything these lessons teach: shapes, accumulating .grad, graphs freed after backward(), parameter registration. It is not PyTorch: initial values and random numbers differ, and "cuda" is only a label. Every code block marked real PyTorch is for you to run on your own machine.

Do the parameter arithmetic in the lab, then build the same kind of network in code.

hidden width h4

freeze the first layer (requires_grad_(False))

name	layer	weight	bias	parameters	trainable?
0	Linear(10, 4)	40	4	44	yes
1	ReLU()	-	-	0	n/a
2	Linear(4, 3)	12	3	15	yes
total / trainable				59	59

model = nn.Sequential(nn.Linear(10, 4), nn.ReLU(), nn.Linear(4, 3)) Linear(in, out) stores weight of shape (out, in) and bias of shape (out,) sum(p.numel() for p in model.parameters()) → 59 state_dict() keys: 0.weight, 0.bias, 2.weight, 2.bias

keys of state_dict() parameters in the ReLU layer

To do: Linear(10, h) → ReLU → Linear(h, 3): choose h so the network has exactly 143 parameters
To do: Freeze the first Linear layer (requires_grad_(False)) so that exactly 33 parameters remain trainable
To do: Type the keys that state_dict() returns for nn.Sequential(Linear, ReLU, Linear), separated by commas
To do: How many parameters does the ReLU layer have?

Part 2. Do it in code

MLP(sizes): sizes is a list like [4, 8, 3] (input, hidden..., output). Build nn.Linear layers between consecutive sizes with nn.ReLU() between them (no ReLU after the last layer) and store them in one nn.Sequential called self.net. forward(x) returns self.net(x).
count_parameters(model, trainable_only=False): the number of scalar parameters. With trainable_only=True count only parameters whose requires_grad is True.
freeze(module): set requires_grad = False on every parameter of the module.

import torch
from torch import nn

class MLP(nn.Module):
    def __init__(self, sizes):
        super().__init__()
        # build self.net
        pass

    def forward(self, x):
        pass

def count_parameters(model, trainable_only=False):
    pass

def freeze(module):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Is an arbitrary unregistered tensor automatically a trainable module parameter?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

For a batch X of shape (B,d), a dense output with h features has shape (B,h). torch.nn.Linear stores weight with shape (h,d) and computes X times weight-transpose plus bias. Parameter count is d*h+h. PyTorch execution belongs in a linked notebook or supported runtime; the chapter does not assume browser PyTorch.

Worked example

A batch has shape (2,4). A Linear(4,3) layer stores a (3,4) weight and a length-three bias. Its output is (2,3). The parameter count is 4×3+3=15. This differs from the batch's element count: model parameters are shared across both examples.

Reference code

The following example uses the Python standard library and does not call a remote service.

batch, inputs, outputs = 2,4,3
weight_shape = (outputs,inputs)
output_shape = (batch,outputs)
assert weight_shape == (3,4)
assert output_shape == (2,3)
assert inputs*outputs+outputs == 15

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

eval mode and no_grad are independent. Model parameters do not multiply with batch size. A tensor detached for logging must not replace the tensor used to compute the training loss.