PyTorch autograd and modules

Three rules to remember

Autograd has three habits that surprise people. This cell shows all three. (A with block runs its indented lines under a setting, then switches the setting back.)

import torch

w = torch.tensor([1.0, 2.0], requires_grad=True)

(w * 3).sum().backward()
(w * 3).sum().backward()
print("1. gradients accumulate:", w.grad)       # 3 + 3, not 3
w.grad = None                                   # what optimizer.zero_grad() does

with torch.no_grad():
    z = w * 3
print("2. no_grad records nothing:", z.requires_grad)

d = (w * 3).detach()
print("3. detach cuts a tensor out of the graph:", d.requires_grad)

1. Gradients add up. The first line prints tensor([6., 6.]), not 3. Each backward() adds its gradients to what's already in .grad. Think of a tip jar: each backward pass drops more in. This is on purpose. A weight used in two places should collect a gradient from each use, and adding is how it does that. So training loops clear gradients before every backward pass. Setting w.grad = None clears them. That's what optimizer.zero_grad() does, which you'll meet in the next lesson.
2. torch.no_grad() switches recording off. Nothing inside the block goes on the receipt, so z.requires_grad is False. Use it when you test a model, and when you update weights. It saves memory. It also keeps the weight update out of the record, since the update isn't part of the model's own sums.
3. .detach() gives the same numbers with no history. It cuts one tensor out of the record. Use it to log a loss for a chart, or to stop gradients flowing into one part of a model.

Models as nn.Module

A real model has thousands or billions of weights. You need a tidy way to keep them together. In PyTorch, a model is a class that inherits from nn.Module. Inherits means your class starts with everything nn.Module already has. You write that as class TinyClassifier(nn.Module):.

Think of a recipe card. The top lists the ingredients. The bottom says what to do with them. An nn.Module has the same two parts:

• __init__ creates the layers. These are the ingredients, and they hold the weights.
• forward says how data flows through the layers. That's the method.

Calling model(x) runs forward for you. And model.parameters() finds every weight in every layer, so you never list them by hand. (A parameter is any number the model learns: a weight or a bias.)

import torch
import torch.nn as nn

class TinyClassifier(nn.Module):
    def __init__(self, n_in, n_hidden, n_out):
        super().__init__()                      # required: sets up parameter tracking
        self.hidden = nn.Linear(n_in, n_hidden)
        self.out = nn.Linear(n_hidden, n_out)

    def forward(self, x):
        return self.out(torch.relu(self.hidden(x)))

torch.manual_seed(0)
model = TinyClassifier(4, 16, 3)
print(model)
for name, p in model.named_parameters():
    print(f"{name:<14} {tuple(p.shape)}")
print("total parameters:", sum(p.numel() for p in model.parameters()))

x = torch.randn(5, 4)                           # a batch of 5 examples
print("output shape:", model(x).shape)

What the output shows:

• print(model) lists the layers by name: hidden and out.
• named_parameters() gives each parameter with its name and shape. hidden.weight is (16, 4) and hidden.bias is (16,). out.weight is (3, 16) and out.bias is (3,).
• numel() counts the numbers in a tensor. The total is 16 × 4 + 16 + 3 × 16 + 3 = 64 + 16 + 48 + 3 = 131.
• A batch of 5 examples went in, with 4 numbers each. Out came 5 rows of 3 scores: shape [5, 3].

torch.relu is the ReLU activation. It keeps positive numbers and turns negative ones into 0.

super().__init__() must come first in __init__. It sets up the bookkeeping that tracks layers. Without it, PyTorch raises an error as soon as you add a layer.

What nn.Linear computes

nn.Linear(4, 16) is 16 neurons, and each one looks at the same 4 inputs. Each neuron has 4 weights and 1 bias. So the layer stores a weight table of shape (16, 4), one row per neuron. It also stores a bias of 16 numbers. Each output is one neuron's weighted sum plus its bias, like y in the first cell.

In symbols, the layer computes $x W^\top + \mathbf{b}$. Here $x$ is the batch of inputs, one row per example. $W$ is the weight table. The small $\top$ means "transposed": the table flipped so its rows become columns. Last, $\mathbf{b}$ is the bias.

In plain words: for each example, every neuron multiplies the inputs by its own weights, adds them up, and adds its bias.

The flip lines up the shapes. A (5, 4) batch times the flipped (4, 16) table gives (5, 16): one row per example, one column per neuron.

For simple stacks like this one there's a shortcut: nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3)). It runs the layers in order, so you don't write forward yourself.

import torch
import torch.nn as nn

layer = nn.Linear(2, 1)                  # 2 inputs, 1 output
with torch.no_grad():                    # set the numbers by hand
    layer.weight[0, 0] = 2.0
    layer.weight[0, 1] = -1.0
    layer.bias[0] = 0.5

x = torch.tensor([[3.0, 4.0]])           # a batch of 1 example, 2 inputs
print(layer(x).item())                   # 2 × 3 + (-1) × 4 + 0.5
print(sum(p.numel() for p in layer.parameters()))   # 2 weights + 1 bias

import torch.nn as nn

layer = nn.Linear(3, 2)
print(sum(p.numel() for p in layer.parameters()))

One training step, by hand

Now put the pieces together. One step of training has four parts:

1. Forward: run the model on the data and work out the loss. Autograd records as it goes.
2. Clear: wipe old gradients with model.zero_grad(), because gradients add up.
3. Backward: loss.backward() fills p.grad for every parameter p.
4. Update: move each parameter a small step against its gradient, inside no_grad().

The data here is made up. x holds 32 random examples with 4 numbers each. The label y counts how many of the first two numbers are positive: 0, 1 or 2. So there are 3 classes. The model's raw scores, one per class, are called logits. F.cross_entropy is the classification loss from the cross-entropy lesson.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
x = torch.randn(32, 4)
y = (x[:, 0] > 0).long() + (x[:, 1] > 0).long()      # a made-up 3-class target

for step in range(5):
    logits = model(x)
    loss = F.cross_entropy(logits, y)
    model.zero_grad()                                # clear old gradients
    loss.backward()                                  # fill p.grad for every parameter
    with torch.no_grad():
        for p in model.parameters():
            p -= 0.5 * p.grad                        # gradient descent
    print(f"step {step}  loss {loss.item():.4f}")

The line p -= 0.5 * p.grad is gradient descent. The 0.5 is the learning rate: how big a step to take. The minus sign moves each weight downhill, away from the direction that raises the loss.

In plain words: new weight = old weight − learning rate × gradient. A weight of 1.0 with a gradient of 0.2 becomes 1.0 − 0.5 × 0.2 = 0.9.

In this teaching version, the loss falls every step, from about 1.13 to about 0.82. Real PyTorch draws different random numbers for the weights and the data, so its losses will differ.

That loop is all of training. The next lesson hands the update to an optimizer (torch.optim) and adds the rest of the standard recipe.