Skip to main content

Parameters, modules and an optimizer

A deliberately PyTorch-shaped library on top of the engine: Parameter, Module with a recursive parameters(), Linear, ReLU, Sequential and an SGD optimizer, with why gradients must be zeroed and why updates happen outside the graph.
DL01.6.3L2 Intermediate25 min
In review

The maths and code on this page are checked by automated tests on every change, but an independent reader has not signed it off yet. Spot a mistake or a confusing step? Tell us.

Before this lesson: DL01.6.2 Matrix products, reductions and broadcasting

What you will be able to do
  • Explain what a Parameter, a Module, parameters() and an optimizer are, and how they share the same tensor objects.
  • Implement Linear with its weight stored as (out, in), ReLU, Sequential and SGD with zero_grad() and step() in under a hundred lines.
  • Count the parameters of a multilayer perceptron from its layer sizes, and check the count of a 2-16-16-2 network.
  • Explain why gradients accumulate and must be zeroed between steps, and why the update must happen outside the graph.

The question

Every PyTorch training script has the same few lines: build model from layers, hand model.parameters() to an optimizer, then repeat optimizer.zero_grad(), loss.backward(), optimizer.step(). The engine of this module can already do backward. What are the objects around it, how does an optimizer find all the weights of a deep network, and what goes wrong if you leave out zero_grad?

From the last lesson

1. In the course convention \({Z = XW\T + \vect{1}\vect{b}\T}\), what are the gradients of \(W\) and \(\vect{b}\)?

Show answer

\({G\T X}\), shape (out, in), and the column sums of \(G\), shape (out,).

2. What is the backward rule of broadcasting?

Show answer

Sum the gradient over every axis that was stretched or added in front, so that it ends with the shape of the array.

Intuition

A network is a tree of parts. The root holds the layers, a layer holds its weight and bias, and the weight and bias are tensors the training is allowed to change. Walk the tree and collect those tensors, and you have a list the optimizer can update.

A tree diagram titled a 2-16-16-2 network. At the top a vermillion box Sequential, with arrows to five boxes below: Linear(2, 16), ReLU, Linear(16, 16), ReLU, Linear(16, 2). Each Linear box has two orange child boxes, weight and bias, with shapes (16, 2) and (16,), (16, 16) and (16,), (2, 16) and (2,); under each ReLU the words no parameters. Orange arrows from all six parameter boxes go down into a wide orange band reading model.parameters(): a list of 6 Parameters, 354 numbers. A green arrow from the band to a green box SGD(params, lr): step(), zero_grad() is labelled the optimizer holds the same objects.
DL01.6.3-F1

The library of this lesson, drawn as the tree it builds. parameters() walks the tree and returns the six tensors marked for learning. The optimizer keeps references to those same objects, so when it changes them, the network changes.

Nothing is copied. The optimizer’s list and the layers point at the same six objects, which is why step() can update the network without knowing anything about its layers.

Build it up

Step 1. A Parameter is a Tensor with a mark on it. It behaves exactly like a tensor; its type says “the training may change me”. A layer’s weight and bias are Parameters. Its inputs and its intermediate results are plain tensors.

Step 2. A Module holds Parameters and other Modules. Its parameters() method looks through its own attributes: a Parameter is collected, a Module is asked for its parameters in turn, and a list of Modules is searched item by item. That recursion is what lets the optimizer find the weights of a network of any depth.

Step 3. Linear, ReLU, Sequential. Linear(n_in, n_out) stores a weight of shape (n_out, n_in), the shape PyTorch’s nn.Linear uses, and a bias of shape (n_out,), and computes x @ W.T + b. ReLU has no parameters. Sequential calls its layers in order.

A horizontal bar chart with one orange bar per Parameter of the 2-16-16-2 network: layer 1 weight (16, 2), 32; layer 1 bias (16,), 16; layer 2 weight (16, 16), 256; layer 2 bias (16,), 16; layer 3 weight (2, 16), 32; layer 3 bias (2,), 2. The axis reads numbers in the Parameter, from 0 to 300, and a bold note says total: 354 equals 48 plus 272 plus 34.
DL01.6.3-F2

Where the 354 numbers of a 2-16-16-2 network live. The square hidden-to-hidden weight holds more than the rest put together.

Step 4. An optimizer holds the list and does two things. zero_grad() sets every parameter’s grad to zero. step() moves every parameter against its gradient, \({w \leftarrow w - \eta\,\nabla_w L}\), the gradient descent of ML01.5.

The formal statement

NoteThe parameter count and the update

A layer from \(n_{l-1}\) inputs to \(n_l\) outputs has a weight of shape \({(n_l, n_{l-1})}\) and a bias of shape \({(n_l,)}\), so \({n_l(n_{l-1} + 1)}\) parameters. A network with widths \({n_0, \dots, n_D}\) has

\[ \sum_{l=1}^{D} n_l\,(n_{l-1} + 1) \]

parameters. One step of SGD with learning rate \(\eta\) replaces every parameter \(\theta\) by \({\theta} - {\eta\,\cgrad{\nabla_\theta L}}\), using the gradients of the current mini-batch only.

For the 2-16-16-2 network the sum is \({16 \cdot 3} + {16 \cdot 17} + {2 \cdot 17} = {48} + {272} + {34} = {354}\).

Why the gradients accumulate. In DL01.6.1 the leaves of the graph keep what their grad holds, and every closure adds with +=. This is by design. A weight used twice in one forward pass, a shared weight, must collect both contributions (M05.8.3, DL01.5.3), and adding is the only way to do that when the contributions arrive one closure at a time. The same rule means that a second call of backward adds a second gradient on top of the first. PyTorch documents the same behaviour: backward accumulates in the leaves, and you may need to zero .grad first.

So the loop must clear the gradients before each backward pass. If it does not, step \(t\) uses the sum of all the gradients so far:

\[ \begin{aligned} \theta_{t+1} &= \theta_t - \eta \sum_{s \le t} \nabla L(\theta_s) \\ &\ne \theta_t - \eta\,\nabla L(\theta_t). \end{aligned} \]

Old gradients, computed at old parameters, keep pushing after they have stopped being true.

Why the update happens outside the graph. Written with tensor operations, w = w - lr * w.grad makes a new tensor whose parents are the old w and a constant. The name w now points at a result, not a Parameter, the optimizer still holds the old object, and every later forward pass would build its graph through a growing chain of past updates. The update is not part of the function being differentiated, so it is done on the arrays: w.data -= lr * w.grad changes the numbers inside the same object.

Two panels. Left, marked with a cross and titled w equals w minus lr times w.grad, a Tensor operation: an orange box old w with an arrow labelled minus to a vermillion box new w: has parents; a green box optimizer has a dashed arrow still pointing to old w; note: the optimizer still holds old w, and the graph keeps growing. Right, marked with a tick and titled w.data minus-equals lr times w.grad, a NumPy operation: one orange box w: same object, new numbers, no parents, with a solid green arrow from the optimizer box; note: the next forward pass starts a fresh graph from w.
DL01.6.3-F5

The update is not part of the model. Done with tensor operations it adds nodes to a graph and breaks the link between the optimizer and the network; done on .data it does neither.

The same idea gives a switch, no_grad(): inside it, operations compute their data but record no parents and no closure. An evaluation pass that will never be differentiated then builds no graph and keeps no arrays alive for a backward pass that never comes.

Worked example

One step of a tiny layer. Take Linear(2, 1) with weight \({(0.5, -1)}\) and bias 0.2, one example \({\vect{x} = (1, 2)}\) with target \({y = 1}\), the squared error and \({\eta = 0.1}\).

A table with columns stage, line and with the numbers. Forward, in vermillion: z equals w dot x plus b, 0.5 times 1 plus minus 1 times 2 plus 0.2 equals minus 1.3; L equals z minus y squared, minus 1.3 minus 1 squared equals 5.29, in purple. Backward, in green: delta z equals 2 times z minus y, 2 times minus 2.3 equals minus 4.6; w.grad plus-equals delta z times x, minus 4.6 times (1, 2) equals (minus 4.6, minus 9.2); b.grad plus-equals delta z, minus 4.6. Step, in green: w.data minus-equals lr times w.grad, (0.5, minus 1) plus (0.46, 0.92) equals (0.96, minus 0.08); b.data minus-equals lr times b.grad, 0.2 plus 0.46 equals 0.66. On a yellow row, again: z equals 1.46, L equals 0.2116.
DL01.6.3-F3

Forward, backward and one SGD step, with every number. The loss falls from 5.29 to 0.2116.

  • Forward: \({z} = {0.5 \cdot 1} + {(-1) \cdot 2} + {0.2} = {-1.3}\), and \({L = (-1.3 - 1)^2 = 5.29}\).
  • Backward: \({\delta_z = 2(z - y) = -4.6}\), so w.grad becomes \({-4.6\,(1, 2) = (-4.6, -9.2)}\) and b.grad becomes \({-4.6}\).
  • Step: \({w} \leftarrow {(0.5, -1)} - {0.1\,(-4.6, -9.2)} = {(0.96, -0.08)}\) and \({b \leftarrow 0.2 + 0.46 = 0.66}\).
  • Again: \({z = 0.96 - 0.16 + 0.66 = 1.46}\) and \({L = 0.46^2 = 0.2116}\).

Now the second step. Zeroed first, the new gradient is \({2 \times 0.46 \times (1, 2)} = {(0.92, 1.84)}\) for \(w\) and 0.92 for \(b\), and the step brings \(z\) to 0.908 and the loss to about 0.0085. Not zeroed, the grads hold \({(-4.6, -9.2)} + {(0.92, 1.84)} = {(-3.68, -7.36)}\) and \({-3.68}\): the old gradient still points up, so \(z\) shoots to 3.668 and the loss to about 7.12.

Code

The engine of DL01.6.1 and DL01.6.2, gathered into one class, is folded here.

The engine of DL01.6.1 and DL01.6.2 
import numpy as np

class Tensor:
    def __init__(self, data):
        self.data = np.asarray(data, dtype=float)
        self.grad = np.zeros_like(self.data)
        self.parents = ()
        self.backward_fn = None

    def record(self, parents, backward_fn):
        self.parents, self.backward_fn = parents, backward_fn
        return self

    @property
    def shape(self):
        return self.data.shape

    def __add__(self, other):
        other = lift(other)
        out = Tensor(self.data + other.data)
        def backward_fn():
            self.grad += unbroadcast(out.grad, self.shape)
            other.grad += unbroadcast(out.grad, other.shape)
        return out.record((self, other), backward_fn)

    def __sub__(self, other):
        other = lift(other)
        out = Tensor(self.data - other.data)
        def backward_fn():
            self.grad += unbroadcast(out.grad, self.shape)
            other.grad += unbroadcast(-out.grad, other.shape)
        return out.record((self, other), backward_fn)

    def __mul__(self, other):
        other = lift(other)
        out = Tensor(self.data * other.data)
        def backward_fn():
            self.grad += unbroadcast(other.data * out.grad, self.shape)
            other.grad += unbroadcast(self.data * out.grad, other.shape)
        return out.record((self, other), backward_fn)

    __radd__, __rmul__ = __add__, __mul__

    def __neg__(self):
        out = Tensor(-self.data)
        def backward_fn():
            self.grad += -out.grad
        return out.record((self,), backward_fn)

    def __pow__(self, k):
        out = Tensor(self.data ** k)
        def backward_fn():
            self.grad += k * self.data ** (k - 1) * out.grad
        return out.record((self,), backward_fn)

    def exp(self):
        out = Tensor(np.exp(self.data))
        def backward_fn():
            self.grad += out.data * out.grad
        return out.record((self,), backward_fn)

    def log(self):
        out = Tensor(np.log(self.data))
        def backward_fn():
            self.grad += out.grad / self.data
        return out.record((self,), backward_fn)

    def tanh(self):
        out = Tensor(np.tanh(self.data))
        def backward_fn():
            self.grad += (1 - out.data ** 2) * out.grad
        return out.record((self,), backward_fn)

    def relu(self):
        out = Tensor(np.maximum(self.data, 0))
        def backward_fn():
            self.grad += (self.data > 0) * out.grad
        return out.record((self,), backward_fn)

    def __matmul__(self, other):
        out = Tensor(self.data @ other.data)
        def backward_fn():
            self.grad += out.grad @ other.data.T
            other.grad += self.data.T @ out.grad
        return out.record((self, other), backward_fn)

    def sum(self, axis=None, keepdims=False):
        out = Tensor(self.data.sum(axis=axis, keepdims=keepdims))
        def backward_fn():
            g = out.grad if keepdims or axis is None else np.expand_dims(out.grad, axis)
            self.grad += np.broadcast_to(g, self.shape)
        return out.record((self,), backward_fn)

    def mean(self, axis=None, keepdims=False):
        total = self.sum(axis=axis, keepdims=keepdims)
        return total * (total.data.size / self.data.size)

    @property
    def T(self):
        out = Tensor(self.data.T)
        def backward_fn():
            self.grad += out.grad.T
        return out.record((self,), backward_fn)

    def topo_order(self):
        order, seen, stack = [], set(), [(self, False)]
        while stack:
            node, finished = stack.pop()
            if finished:
                order.append(node)
            elif id(node) not in seen:
                seen.add(id(node))
                stack.append((node, True))
                stack.extend((p, False) for p in node.parents if id(p) not in seen)
        return order

    def backward(self):
        order = self.topo_order()
        for node in order:
            if node.parents:
                node.grad = np.zeros_like(node.data)
        self.grad = np.ones_like(self.data)
        for node in reversed(order):
            if node.backward_fn:
                node.backward_fn()

def lift(x):
    return x if isinstance(x, Tensor) else Tensor(x)

def unbroadcast(grad, shape):
    while grad.ndim > len(shape):
        grad = grad.sum(axis=0)
    for axis, size in enumerate(shape):
        if size == 1 and grad.shape[axis] != 1:
            grad = grad.sum(axis=axis, keepdims=True)
    return grad

def cross_entropy(z, labels):
    n = z.shape[0]
    shifted = z.data - z.data.max(axis=1, keepdims=True)
    log_p = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
    out = Tensor(-log_p[np.arange(n), labels].mean())
    def backward_fn():
        g = np.exp(log_p)
        g[np.arange(n), labels] -= 1
        z.grad += g / n * out.grad
    return out.record((z,), backward_fn)

First the switch. record is the one place where an operation attaches its parents and closure, so the switch lives there.

import contextlib

Tensor.recording = True

def record(self, parents, backward_fn):
    if Tensor.recording:                         # under no_grad: data only, no graph
        self.parents, self.backward_fn = parents, backward_fn
    return self

Tensor.record = record

@contextlib.contextmanager
def no_grad():
    Tensor.recording = False
    try:
        yield
    finally:
        Tensor.recording = True

The library itself.

class Parameter(Tensor):
    """A Tensor that the optimizer is allowed to change."""

class Module:
    def parameters(self):
        """Every Parameter in this module and, recursively, in the modules it holds."""
        found = []
        for value in vars(self).values():
            for item in (value if isinstance(value, list) else [value]):
                if isinstance(item, Parameter):
                    found.append(item)
                elif isinstance(item, Module):
                    found.extend(item.parameters())
        return found

    def __call__(self, x):
        return self.forward(x)

class Linear(Module):
    def __init__(self, n_in, n_out, rng):
        bound = 1 / np.sqrt(n_in)                # PyTorch's range for nn.Linear; DL02.3 explains
        self.weight = Parameter(rng.uniform(-bound, bound, size=(n_out, n_in)))
        self.bias = Parameter(rng.uniform(-bound, bound, size=n_out))

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

class ReLU(Module):
    def forward(self, x):
        return x.relu()

class Sequential(Module):
    def __init__(self, *layers):
        self.layers = list(layers)

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

class SGD:
    def __init__(self, params, lr):
        self.params, self.lr = list(params), lr

    def zero_grad(self):
        for p in self.params:
            p.grad = np.zeros_like(p.data)

    def step(self):
        for p in self.params:
            p.data -= self.lr * p.grad           # on the array: outside the graph

A 2-16-16-2 network, its parameters, and a forward pass on four points:

rng = np.random.default_rng(0)
model = Sequential(Linear(2, 16, rng), ReLU(), Linear(16, 16, rng), ReLU(), Linear(16, 2, rng))
params = model.parameters()
print("shapes:", [p.shape for p in params])
print("parameters:", sum(p.data.size for p in params))
x = Tensor(rng.normal(size=(4, 2)))
print("logits shape:", model(x).shape)
shapes: [(16, 2), (16,), (16, 16), (16,), (2, 16), (2,)]
parameters: 354
logits shape: (4, 2)

The worked example, then its second step with and without zero_grad:

def tiny_run(zero_grad, steps):
    layer = Linear(2, 1, np.random.default_rng(0))
    layer.weight.data[:] = [[0.5, -1.0]]
    layer.bias.data[:] = [0.2]
    opt = SGD(layer.parameters(), lr=0.1)
    x, losses = Tensor([[1.0, 2.0]]), []
    for _ in range(steps):
        loss = ((layer(x) - 1.0) ** 2).mean()
        losses.append(float(loss.data))
        if zero_grad:
            opt.zero_grad()
        loss.backward()
        opt.step()
    return losses, layer

losses, layer = tiny_run(True, 1)
print("after one step: w =", np.round(layer.weight.data, 4).tolist(), " b =", np.round(layer.bias.data, 4).tolist())
print("with zero_grad:   ", np.round(tiny_run(True, 3)[0], 4).tolist())
print("without zero_grad:", np.round(tiny_run(False, 3)[0], 4).tolist())
after one step: w = [[0.96, -0.08]]  b = [0.66]
with zero_grad:    [5.29, 0.2116, 0.0085]
without zero_grad: [5.29, 0.2116, 7.1182]

Last, the switch at work: the same forward pass with and without a graph.

logits = model(x)
with no_grad():
    plain = model(x)
print("nodes recorded:", len(logits.topo_order()), " under no_grad:", len(plain.topo_order()))
print("same numbers:", np.array_equal(logits.data, plain.data))
nodes recorded: 18  under no_grad: 1
same numbers: True

The recorded pass has 18 nodes: the input, the six Parameters, and eleven results, three for each Linear (the transpose, the product and the addition of the bias) and one for each ReLU. Under no_grad the output is a lone tensor with no parents.

Explore

The module’s interactive, Broadcast and backward, is in DL01.6.2. Try the bias case there, shapes (2, 3) and (1, 3): it is the addition inside every Linear.

Pitfalls

Two loss curves against the step, from 0 to 11, on axes from 0 to 9. Left, marked with a cross and titled No zero_grad: the purple curve starts at 5.29, drops to 0.21, then jumps between about 0 and 7.5 for the rest of the run, with the note old gradients still push. Right, marked with a tick and titled zero_grad each step: the curve starts at 5.29 and falls to 0 within three steps, with the note the loss reaches 0.
DL01.6.3-F4

Wrong beside right, on the worked example. Without zero_grad, each step uses the sum of every gradient so far, and the loss jumps about instead of settling.

“backward replaces the old gradient.”

It adds to it. The engine, like PyTorch, accumulates in the leaves so that shared weights work, and the price is that the training loop must call zero_grad() before each backward pass. The bug is quiet: the loss still moves, just badly.

“w = w - lr * w.grad is the update.”

On tensors it makes a new node with parents. The optimizer keeps the old Parameter, so the model no longer changes, and the graph grows by a node per step. Update the array in place, w.data -= lr * w.grad, outside the graph.

“A layer’s weight has shape (in, out).”

In this library and in PyTorch it is (out, in), and the layer computes x @ W.T + b. Copying weights between a library that stores (in, out), such as scikit-learn’s coefs_, and one that stores (out, in) needs a transpose. With equal sizes in and out the shapes match either way, so the mistake can pass every shape check.

Check yourself

Recall

1. What does parameters() return, and how does it find the weights of a deep network?

Show answer

The list of every Parameter in the module and in the modules it holds; it looks through its attributes and asks each Module it finds for its own parameters, recursively.

2. How many parameters does a layer from \(n_{l-1}\) inputs to \(n_l\) outputs have?

Show answer

\({n_l(n_{l-1} + 1)}\): a weight of shape \({(n_l, n_{l-1})}\) and \(n_l\) biases.

3. Why do gradients accumulate across calls of backward?

Show answer

Closures add into grad and leaves keep their grad, which is what a weight used more than once needs; so a second backward pass adds to the first.

4. Why is the update done on .data instead of with tensor operations?

Show answer

The update is not part of the function being differentiated. With tensor operations it would create a new node with parents, leave the optimizer holding the old Parameter and grow the graph.

5. What does no_grad() change?

Show answer

Operations still compute their data but record no parents and no closure, so no graph is built.

Exercises

1. ★ Count the parameters of a 64-32-10 network, the course’s image classifier.

Show solution

\({32 \times (64 + 1)} + {10 \times (32 + 1)} = {2{,}080} + {330} = {2{,}410}\).

2. ★★ Continue the worked example for one more step without zero_grad and confirm the loss of about 7.12.

Show solution

After step 1, \({w = (0.96, -0.08)}\), \({b = 0.66}\) and \({z = 1.46}\). Step 2’s own gradient is \({\delta_z = 2(1.46 - 1) = 0.92}\), so it adds \({(0.92, 1.84)}\) and 0.92 to grads that still hold \({(-4.6, -9.2)}\) and \({-4.6}\): totals \({(-3.68, -7.36)}\) and \({-3.68}\). The step gives \({w} = {(0.96 + 0.368,\ -0.08 + 0.736)} = {(1.328, 0.656)}\) and \({b = 0.66 + 0.368 = 1.028}\), so \({z = 1.328 + 1.312 + 1.028 = 3.668}\) and \({L = 2.668^2 \approx 7.118}\).

3. ★★★ A weight matrix is shared: the same Parameter is used by two Linear layers in one forward pass. Explain why parameters() of a Sequential holding both layers lists it twice, what SGD.step() would then do, and how to fix it.

Show solution

parameters() collects each layer’s attributes, and both layers hold the same object, so it appears twice in the list. During backward its grad correctly collects both contributions, by accumulation. But step() walks the list and subtracts \({\eta\,\nabla}\) once per entry, so the shared weight moves twice as far as it should. The fix is to remove repeats by identity: keep a set of id(p) already seen and skip any parameter found again. PyTorch’s named_parameters does this by default (remove_duplicate=True).

Summary card

A library is a tree of Modules whose Parameters the optimizer holds.

The anchor figure again: Sequential over Linear, ReLU, Linear, ReLU, Linear, with each Linear's weight and bias collected by parameters() into a list of six Parameters, 354 numbers, held by an SGD optimizer.
DL01.6.3-F1

Walk the tree, collect, update.

  • Parameters per layer: \({n_l(n_{l-1} + 1)}\); 354 for 2-16-16-2.
  • Linear: weight (out, in), computes x @ W.T + b.
  • Each step: zero_grad(), backward(), step().
  • Update on .data; evaluate under no_grad().

Sources and next steps

Sources. Each was opened and checked against this lesson on 15 November 2026.

  • PyTorch 2.15, torch.nn.Linear: computes \({y = xA\T + b}\), with weight of shape (out_features, in_features) and bias of shape (out_features), both initialised from \({\mathcal{U}(-\sqrt{k}, \sqrt{k})}\) with \({k = 1/\text{in\_features}}\).
  • PyTorch 2.15, torch.Tensor.backward: the function accumulates gradients in the leaves, so .grad may need zeroing first; and Autograd mechanics: computations in no-grad mode behave as if no input requires a gradient and are not recorded; and torch.nn.Module: named_parameters removes duplicated parameters by default.
  • A. Karpathy, micrograd (MIT licence): a small neural network library with a PyTorch-like API on top of a scalar engine, trained with two hidden layers of 16 units on the moons; the design this lesson follows on arrays.
  • I. Goodfellow, Y. Bengio and A. Courville, Deep Learning, section 6.5.8: real implementations add each contribution into a single buffer instead of computing every term separately and summing at the end.

Next. DL01.6.4 Train with your engine and check it holds the engine to the hand-derived gradients of DL01.5, then trains with it.

This lesson unlocks

DL01.6.4 Train with your engine and check it

Your progress is saved in this browser only. No account, no tracking.
Back to top