Intro to PyTorch

1. Intuition

In backpropagation you computed gradients by hand: run the forward pass, then walk backwards multiplying local derivatives. That is mechanical, so a computer can do it if it remembers what happened.

PyTorch remembers. A tensor is an n-dimensional array, like a NumPy array, that can live on a CPU or GPU. If a tensor has requires_grad=True (all model parameters do), every operation on it creates a new tensor that records which operation made it and from which inputs. Chained together, these records form a graph from the parameters to the loss.

loss.backward() walks that graph from the loss back to the parameters. At each node it applies that operation's local derivative, multiplies by the gradient flowing in (the chain rule), and adds the result into the inputs' .grad. Note the word adds: gradients accumulate. That's handy (you can sum gradients over several small batches) and dangerous (if you forget to reset them, each step uses the sum of all previous gradients).

On top of autograd:

• nn.Module is a container that registers its parameters, so model.parameters() finds them all.
• An optimizer (SGD, AdamW) holds references to those parameters and, on step(), updates each one using its .grad.
• A DataLoader cuts the dataset into shuffled mini-batches.

The training loop is the same five steps every time: forward, loss, clear old gradients, backward, step.

2. Visualization

Synthetic two-moons data, 480 training and 120 validation points. Left: one missing line (opt.zero_grad()) and every step uses the running sum of all earlier gradients. Right: what the correct loop learned, a curved boundary no linear model could draw.

Try it (predict, then run the script):

1. Without zero_grad, the effective gradient is a sum that keeps growing. AdamW divides by a running estimate of gradient size, which partly hides this. Predict what happens if you switch the optimizer to plain torch.optim.SGD(..., lr=0.1).
2. Remove model.eval() from accuracy. Predict whether validation accuracy changes from run to run, and why.
3. Change hidden=32 to hidden=4. Predict the new parameter count before running.
4. Print model(xtr).shape and ytr.shape. Why does F.cross_entropy accept a [480, 2] float tensor and a [480] integer tensor?

3. The math

This lesson is about a mechanism, so this section walks one tiny computation through autograd with every number.

Symbols

Forward pass: build the graph

With $x = 2$, $w = 3$, $b = 1$, $y = 10$, autograd records four nodes:

Backward pass: walk it in reverse

Start with $\bar{L} = 1$. Each step multiplies the incoming gradient by the local derivative:

So w.grad = -12 and b.grad = -6, which both the scratch engine and PyTorch print.

Accumulation and the update

Call backward again on a freshly computed loss without clearing, and each .grad gets the new gradient added: $\bar{w} = -12 + (-12) = -24$. That's what zero_grad() prevents.

One SGD step with $\eta = 0.01$ and the correct gradients:

New prediction $3.12 \times 2 + 1.06 = 7.30$, new loss $(7.30 - 10)^2 = 7.29$, down from 9.

Shapes in the real model

The script's MLP maps $[B, 2] \to [B, 32] \to [B, 32] \to [B, 2]$. nn.Linear(in, out) stores a weight of shape $[\text{out}, \text{in}]$ and computes $x W^\top + b$. Parameters:

F.cross_entropy(logits, y) takes logits $[B, K]$ (float) and class indices $[B]$ (int64), and returns the mean loss, a scalar, which is what .backward() needs. At step 0 it should be near $\ln 2 = 0.693$ for 2 classes; the script prints 0.682.

4. Implementation

Build it from scratch

A Value holds a number, its gradient, the values it came from, and a function that pushes its gradient to those parents. Addition and multiplication each record their local derivatives.

class Value:
    """A number that remembers how it was computed, so it can send gradients back."""

    def __init__(self, data, parents=(), backward=lambda: None):
        self.data, self.grad = data, 0.0
        self._parents, self._backward = parents, backward

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other))
        def backward():                       # d(a+b)/da = 1, d(a+b)/db = 1
            self.grad += out.grad
            other.grad += out.grad
        out._backward = backward
        return out

Predict: why += rather than = in backward? (Hint: what if the same value is used twice, as in a + a?)

Multiplication swaps in the other factor; power uses the power rule. Subtraction is built from the other two.

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other))
        def backward():                       # d(ab)/da = b, d(ab)/db = a
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = backward
        return out

    def __sub__(self, other):
        return self + (other * -1 if isinstance(other, Value) else -other)

    def __pow__(self, k):
        out = Value(self.data**k, (self,))
        def backward():                       # d(a^k)/da = k a^(k-1)
            self.grad += k * self.data ** (k - 1) * out.grad
        out._backward = backward
        return out

Predict: for c = a * b with a = 3, b = 2, and c.grad = 1, what are a.grad and b.grad?

backward() sorts the graph so each node runs after every node that uses it, seeds the loss with gradient 1, and walks in reverse.

    def backward(self):
        order, seen = [], set()
        def visit(v):                         # topological order: parents before children
            if v not in seen:
                seen.add(v)
                for p in v._parents:
                    visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0                       # dL/dL = 1
        for v in reversed(order):             # walk from the loss back to the leaves
            v._backward()

Predict: what goes wrong if you call L.backward() twice on the same L? (Every intermediate .grad still holds its first-pass value, so the second pass adds on top at every node: w.grad ends at −60, not −24. PyTorch frees the graph after backward(), so a second call raises an error unless you pass retain_graph=True.)

That is all of autograd. PyTorch does the same thing on tensors, with each operation's backward written in C++ or CUDA.

Use the library

Same graph, same numbers:

wt = torch.tensor(3.0, requires_grad=True)
bt = torch.tensor(1.0, requires_grad=True)
Lt = (wt * 2.0 + bt - 10.0) ** 2
Lt.backward()                                # wt.grad = -12, bt.grad = -6

A model is an nn.Module. Layers assigned as attributes are registered, so model.parameters() finds their weights. The last layer outputs raw logits, because F.cross_entropy applies log-softmax itself (see loss functions).

class MLP(nn.Module):
    def __init__(self, d_in=2, hidden=32, d_out=2):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden, d_out),           # raw logits: no softmax here
        )

    def forward(self, x):                       # x: [batch, d_in] -> logits: [batch, d_out]
        return self.net(x)

The training loop, five numbered steps. model.train() turns dropout (and batch-norm statistics updates) on.

def train(model, loader, epochs, lr=1e-2, zero_grad=True):
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    losses = []
    for _ in range(epochs):
        model.train()                           # dropout on
        total = 0.0
        for xb, yb in loader:
            logits = model(xb)                  # 1. forward
            loss = F.cross_entropy(logits, yb)  # 2. loss
            if zero_grad:
                opt.zero_grad()                 # 3. clear old gradients
            loss.backward()                     # 4. backward: fill .grad
            opt.step()                          # 5. update parameters
            total += loss.item() * len(xb)
        losses.append(total / len(loader.dataset))
    return losses

Evaluation turns off both dropout (model.eval()) and graph recording (torch.no_grad(), here as a decorator), which saves memory and time.

@torch.no_grad()
def accuracy(model, x, y):
    model.eval()                                # dropout off
    return (model(x).argmax(1) == y).float().mean().item()

.item() pulls a Python number out of a one-element tensor; use it for logging so you don't keep the graph alive by storing the loss tensor. Save weights with torch.save(model.state_dict(), path) and restore with model.load_state_dict(torch.load(path, weights_only=True)).

Beyond CPU. Move the model and every batch to the same device: device = "cuda" if torch.cuda.is_available() else "cpu", then model.to(device) and xb.to(device) (Apple silicon uses "mps"). torch.compile(model) (PyTorch 2.x) traces and fuses the model into faster kernels without changing your code; for real training add mixed precision with torch.autocast (see numerical stability).

Run it

uv run code/10-neural-networks/intro_to_pytorch.py

PyTorch is declared in the script's header (PEP 723 inline metadata), so uv installs it into a cached environment on first run; that download takes a minute or two. Key output:

L = 9.0, dL/dw = -12.0, dL/db = -6.0
L = 9.0, dL/dw = -12.0, dL/db = -6.0
after a second backward without zeroing: dL/dw = -24.0 (gradients accumulate)
X + bias =
tensor([[10., 21., 32.],
        [13., 24., 35.]])
weight (3, 2), bias (3,)  (out_features, in_features)
parameters: 1218  (2*32+32 + 32*32+32 + 32*2+2)
step-0 loss 0.6815 (ln 2 = 0.6931)
correct loop:      final train loss 0.0406, val accuracy 1.000
no zero_grad():    final train loss 0.4439, val accuracy 0.775
train mode, same input twice -> identical? False;  eval mode -> identical? True

Exact losses can differ slightly across PyTorch versions and CPUs; the pattern won't.

5. Engineering

Where the time goes. PyTorch is fast because each tensor operation runs as one optimized kernel over the whole batch. Python loops over individual numbers (like the scratch Value) are millions of times slower. Keep work in batched tensor ops, and keep the GPU busy: DataLoader(num_workers=..., pin_memory=True) overlaps data loading with compute.

Memory. Autograd keeps every intermediate activation until backward, so activation memory grows with batch size and depth. Use torch.no_grad() (or torch.inference_mode()) for evaluation, and gradient checkpointing for very deep models. Accumulating total += loss (a tensor) instead of loss.item() keeps every step's graph alive and leaks memory.

Reproducibility. Seed torch.manual_seed, the DataLoader's generator, and NumPy/Python if used. On GPU, some kernels are non-deterministic; torch.use_deterministic_algorithms(True) trades speed for repeatability.

Gradient accumulation on purpose. To simulate a batch of 256 on a GPU that fits 64, call backward() on loss / 4 for four micro-batches, then step() and zero_grad() once. This is the useful side of accumulation.

Checkpoints. Save model.state_dict(), optimizer.state_dict(), the scheduler, the step, and the RNG states, so a resumed run continues exactly. Load with weights_only=True for files you didn't create; arbitrary pickles can run code.

Watch out

Failure modes. Forgetting opt.zero_grad() (gradients accumulate across steps); forgetting model.eval() (dropout stays on and batch norm uses batch statistics at inference); softmax before F.cross_entropy; model on GPU and data on CPU (a device mismatch error); calling .numpy() on a tensor that requires grad (use .detach() first); in-place operations on tensors autograd still needs.

Common mistakes

• Putting zero_grad() in the wrong place (it must run before each backward(), not between backward() and step()).
• Thinking torch.no_grad() and model.eval() do the same thing. One stops graph recording; the other changes layer behaviour. Inference needs both.
• Using integer tensors as model inputs (torch.arange(6) is int64; nn.Linear needs float).
• Leaving shuffle=True on the validation loader and then comparing per-batch metrics across runs.

6. Knowledge check

Practice exercise. With $x = 1$, $w = 2$, $b = -1$, $y = 4$, compute $L = (wx + b - y)^2$ and both gradients by hand. Then compute $w$ and $b$ after one SGD step with $\eta = 0.1$.

Implementation challenge. Add a relu() method to Value, build a 2-input neuron $\text{relu}(w_1 x_1 + w_2 x_2 + b)$, and check all three gradients against PyTorch for one input where the ReLU is active and one where it is not.

def relu(self):
    out = Value(max(0.0, self.data), (self,))
    def backward():
        self.grad += (self.data > 0) * out.grad      # slope 1 if active, else 0
    out._backward = backward
    return out
Value.relu = relu

Summary

• A tensor with requires_grad=True records the operations applied to it; backward() runs the chain rule over that record and adds into .grad.
• Gradients accumulate, so call zero_grad() once per step before backward().
• nn.Module registers parameters; nn.Linear(in, out) has weight shape [out, in].
• The loop is forward, loss, zero_grad, backward, step. Evaluation uses model.eval() and torch.no_grad().
• Output raw logits and let F.cross_entropy apply the softmax.

Next: Gradient descent

Related: Backpropagation · Loss functions · Debugging neural networks · Numerical stability

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Introduction to PyTorch.

Interview angle