Building models with nn.Module: Exercises — AI-ML deep dive

This is your self-testing arena for Building models with nn.Module. Every problem has a fully worked solution hidden inside a collapsible callout — try first, then reveal. Problems climb from recognising the API to synthesising whole architectures and reasoning about gradients.

Before you start, one tiny vocabulary check so no symbol sneaks in undefined:

Level 1 — Recognition

Exercise 1.1

Which of these is registered as a trainable parameter by PyTorch, and which is silently ignored?

class A(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.randn(5))     # (a)
        self.v = torch.randn(5)                    # (b)
        self.layers = [nn.Linear(3, 3)]            # (c)
        self.layers2 = nn.ModuleList([nn.Linear(3, 3)])  # (d)

Recall Solution

(a) registered — it is an nn.Parameter, so __setattr__ records it. ✅
(b) NOT registered — a plain torch.randn(5) tensor. It rides along on .to(device) only if you register it as a buffer; as written it is invisible to .parameters(). ❌
(c) NOT registered — a plain Python list. PyTorch's __setattr__ never looks inside ordinary lists, so the nn.Linear inside is a ghost: it exists, but .parameters() never finds it and the optimizer never trains it. ❌
(d) registered — nn.ModuleList is the special container that does register everything inside. ✅

Exercise 1.2

What single method call gives you every learnable tensor of a model, and what does it return?

Recall Solution

model.parameters() returns an iterator over all nn.Parameter tensors, collected by recursively walking the submodule tree. To count numbers: sum(p.numel() for p in model.parameters()). To also see names: model.named_parameters().

Level 2 — Application

Exercise 2.1

Count the parameters of nn.Linear(100, 50).

Recall Solution

Weights: $50 \times 100 = 5000$ . Bias: $50$ . Total $= 5050$ . ✅

Exercise 2.2

Build a classifier nn.Linear(784, 256) → ReLU → nn.Linear(256, 10). How many parameters total?

Recall Solution

Layer 1: $256 \times 784 + 256 = 200704 + 256 = 200960$ .
ReLU is stateless — zero parameters.
Layer 2: $10 \times 256 + 10 = 2560 + 10 = 2570$ .
Total $= 200960 + 2570 = 203530$ . ✅

Exercise 2.3

A DeepNetwork(input_dim=64, hidden_dim=32, num_layers=3, output_dim=8) follows the parent's pattern: input_layer = Linear(64,32), then a ModuleList of num_layers copies of Linear(32,32), then output_layer = Linear(32,8). Total parameters?

Recall Solution

input_layer: $32 \times 64 + 32 = 2048 + 32 = 2080$ .
Each hidden Linear(32,32): $32 \times 32 + 32 = 1024 + 32 = 1056$ . Three of them: $3 \times 1056 = 3168$ .
output_layer: $8 \times 32 + 8 = 256 + 8 = 264$ .
Total $= 2080 + 3168 + 264 = 5512$ . ✅

Level 3 — Analysis

Exercise 3.1

Two students build the "same" 2-hidden-layer net. Student A uses a Python list; Student B uses nn.ModuleList:

# A:
self.layers = [nn.Linear(10,10), nn.Linear(10,10)]
# B:
self.layers = nn.ModuleList([nn.Linear(10,10), nn.Linear(10,10)])

Both call the layers correctly in forward. After training, whose hidden layers actually improved, and what does sum(p.numel() for p in model.parameters()) print for each?

Recall Solution

Each Linear(10,10) has $10 \times 10 + 10 = 110$ params.

Student B registers both → .parameters() sees $2 \times 110 = 220$ . The optimizer updates them; they train. ✅
Student A's list is invisible → .parameters() returns 0 (no other params exist). The optimizer receives an empty list; those layers keep their random init forever — the forward pass still runs, loss may even wobble due to randomness, but nothing learns. ❌
Prints: B → 220, A → 0.

Exercise 3.2

In the parent's ResidualBlock, the forward is out = F(x); out = out + identity; out = relu(out). Suppose a beginner "simplifies" by removing the skip: out = relu(F(x)). Using the gradient rule from the parent, explain what quantity is lost.

Recall Solution

With the skip, the gradient reaching the input is $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial out} (\frac{\partial F ( x )}{\partial x} + 1) .$ Remove the skip and you lose the $+ 1$ : $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial out} \cdot \frac{\partial F ( x )}{\partial x} .$ What's lost: the guaranteed " $+ 1$ highway." If $\frac{\partial F ( x )}{\partial x} \approx 0$ (dead layer), the skip version still passes gradient $\approx \frac{\partial L}{\partial out} \cdot 1$ , but the stripped version passes $\approx 0$ — the classic vanishing gradient. The picture:

Level 4 — Synthesis

Exercise 4.1

Design a module MLPBlock(dim) whose forward is $out = x + W_{2} σ (W_{1} x)$ where $σ$ is ReLU, $W_{1}$ maps $d im \to 4 d im$ (with bias) and $W_{2}$ maps $4 d im \to d im$ (with bias) — a transformer-style feed-forward block with a residual. Write it, then compute the parameter count for $d im = 8$ .

Recall Solution

class MLPBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, 4 * dim)   # W1
        self.fc2 = nn.Linear(4 * dim, dim)   # W2
        self.act = nn.ReLU()
    def forward(self, x):
        return x + self.fc2(self.act(self.fc1(x)))

The x + residual adds no parameters (it's just addition). Count for $d im = 8$ :

fc1: $32 \times 8 + 32 = 256 + 32 = 288$ .
fc2: $8 \times 32 + 8 = 256 + 8 = 264$ .
Total $= 288 + 264 = 552$ . ✅

Exercise 4.2

Build TowerNet(dims) that takes a list dims = [d0, d1, ..., dn] and chains Linear(d0,d1), ReLU, Linear(d1,d2), ReLU, ..., Linear(d(n-1), dn) — ReLU between every consecutive pair but not after the last layer. Use the correct container. Then count params for dims = [16, 8, 4, 2].

Recall Solution

class TowerNet(nn.Module):
    def __init__(self, dims):
        super().__init__()
        self.linears = nn.ModuleList(
            [nn.Linear(dims[i], dims[i+1]) for i in range(len(dims) - 1)]
        )
        self.act = nn.ReLU()
    def forward(self, x):
        for i, lin in enumerate(self.linears):
            x = lin(x)
            if i < len(self.linears) - 1:   # no ReLU after the last
                x = self.act(x)
        return x

nn.ModuleList (not a plain list) so all layers register. Count for [16,8,4,2]:

Linear(16,8): $8 \times 16 + 8 = 128 + 8 = 136$ .
Linear(8,4): $4 \times 8 + 4 = 32 + 4 = 36$ .
Linear(4,2): $2 \times 4 + 2 = 8 + 2 = 10$ .
Total $= 136 + 36 + 10 = 182$ . ✅

Level 5 — Mastery

Exercise 5.1

You have SimpleClassifier(input_dim, hidden_dim, output_dim) from the parent (three Linear layers fc1: input→hidden, fc2: hidden→hidden, fc3: hidden→output). Derive a closed-form parameter count $P (i, h, o)$ as a formula in $i, h, o$ , then verify it reproduces the parent's $117121$ for $i = 784, h = 128, o = 1$ .

Recall Solution

Sum the three layers using $o u t \times in + o u t$ each: $P (i, h, o) = f c 1 (hi + h) + f c 2 (h^{2} + h) + f c 3 (o h + o) .$ Group it: $P (i, h, o) = h (i + 1) + h (h + 1) + o (h + 1) .$ Plug in $i = 784, h = 128, o = 1$ :

$f c 1 = 128 \cdot 784 + 128 = 100352 + 128 = 100480$ .
$f c 2 = 12 8^{2} + 128 = 16384 + 128 = 16512$ .
$f c 3 = 1 \cdot 128 + 1 = 129$ .
Total $= 100480 + 16512 + 129 = 117121$ . ✅ Matches the parent.

Exercise 5.2

Capacity vs. cost trade-off. For SimpleClassifier, if $i$ and $o$ are fixed, the count grows fastest in which variable, and why? Using $P (i, h, o) = h (i + 1) + h (h + 1) + o (h + 1)$ , find the dominant term for large $h$ and confirm the growth rate.

Recall Solution

Expand: $P = hi + h + h^{2} + h + o h + o = h^{2} + h (i + o + 2) + o$ .

The $h^{2}$ term (from fc2, the hidden-to-hidden layer) dominates for large $h$ — parameter count grows quadratically in the hidden width.
Why: fc2 connects $h$ inputs to $h$ outputs, an $h \times h$ grid, so doubling $h$ roughly quadruples its size, while fc1 and fc3 only grow linearly in $h$ .
Sanity check the rate: at $h = 128$ , $h^{2} = 16384$ ; at $h = 256$ , $h^{2} = 65536 = 4 \times 16384$ . Doubling $h$ multiplied the dominant term by $4$ . ✅

Recall Quick self-check reveals

Params of nn.Linear(100,50) ::: $5050$ Params of Linear(784,256)→ReLU→Linear(256,10) ::: $203530$ DeepNetwork(64,32,3,8) total ::: $5512$ MLPBlock(8) total ::: $552$ TowerNet([16,8,4,2]) total ::: $182$ Dominant growth of SimpleClassifier in hidden width $h$ ::: quadratic, $O (h^{2})$ , from fc2