Autograd and your first PyTorch training loop

The big idea

In the last lesson you computed nine gradients by hand. A real network has millions of parameters and dozens of layer types. Nobody writes those backward passes by hand. Deep learning frameworks such as PyTorch do it for you with autogradAutomatic differentiation: software that records the operations performed on tensors and applies the chain rule in reverse to compute exact gradients..

The idea is a flight recorder. As your ordinary-looking Python code runs the forward pass, every operation on a tracked tensor is written down: what it was, which tensors went in, and how to compute its local gradient. When you call loss.backward(), the framework plays the recording in reverse, applying the chain rule at each step, and leaves each parameter's gradient in its .grad field.

With that in place, training any network comes down to the same short loop: get a batch, run the model, compute the loss, clear old gradients, backpropagate, update. This lesson walks through that loop line by line, then shows a tiny working autograd engine in plain numpy that implements the same interface, so nothing is magic.

How it works

Step 1: Tensors

A tensorA multi-dimensional array of numbers (a scalar, vector, matrix or higher) with a shape, a data type and a device. The basic data object in deep learning frameworks. is an n-dimensional array, like a numpy array, with three extra properties that matter here: a data type (usually 32-bit floats for training), a device (CPU or GPU), and whether it requires gradients. Parameters are created with requires_grad=True; input data usually is not.

Step 2: The recorded graph

Every operation involving a tensor that requires gradients produces a new tensor that remembers its parents and its backward rule. In PyTorch you can see this as loss.grad_fn. The graph is built fresh on every forward pass, which is why you can use normal Python if statements and loops inside a model: the recording simply follows whatever path the code took.

Step 3: backward, zero_grad and step

• loss.backward() walks the graph from the loss back to every leaf tensor with requires_grad=True, and adds dL/dparam into param.grad.
• Gradients are added, not overwritten. That is useful for accumulating over several small batches, but it means you must clear them every step with optimizer.zero_grad(). Forgetting this is the classic first bug.
• optimizer.step() updates every parameter it was given using its .grad. For plain SGD that is p -= lr * p.grad; Adam and friends ( Optimizers: SGD, momentum, RMSProp, Adam, AdamW ) use the same call.

Step 4: Mini-batches and the DataLoader

Computing the loss over the whole dataset for every step is slow and unnecessary. A mini-batchA small random subset of the training data (often 32 to 1024 examples) used to compute one gradient step. of, say, 32 examples gives a noisy but useful estimate of the gradient at a fraction of the cost. A DataLoader wraps a dataset, shuffles it each epochOne full pass over the training data., and yields batches. With 256 examples and batch size 32, one epoch is 8 update steps.

Step 5: The canonical loop

Here it is in PyTorch, for a small regression network. The code below is the standard pattern; the runnable version further down implements the same interface in numpy so you can see what happens inside each call.

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

device = "cuda" if torch.cuda.is_available() else "cpu"
X = torch.rand(256, 1) * 4 - 2
Y = torch.sin(2 * X) + 0.1 * torch.randn_like(X)
loader = DataLoader(TensorDataset(X, Y), batch_size=32, shuffle=True)

model = nn.Sequential(nn.Linear(1, 32), nn.ReLU(), nn.Linear(32, 1)).to(device)
opt = torch.optim.SGD(model.parameters(), lr=0.02)
loss_fn = nn.MSELoss()

for epoch in range(300):
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)   # data must live where the model lives
        loss = loss_fn(model(xb), yb)           # forward pass and loss
        opt.zero_grad()                         # clear old gradients
        loss.backward()                         # autograd: fill p.grad for every parameter
        opt.step()                              # update the parameters

Step 6: Devices, train and eval

Tensors live on a deviceWhere a tensor's memory and computation live, such as the CPU or a particular GPU. Operations require all their tensors on the same device.. Move the model once with model.to(device) and each batch with .to(device); mixing devices raises an error. Two more lines you will see everywhere: model.train() and model.eval() switch layers such as dropout and batch norm between training and inference behaviour, and wrapping evaluation in with torch.no_grad(): stops the recorder, saving memory and time when you do not need gradients.

Going deeper

Code: a tiny autograd engine with the same loop

import numpy as np

class Tensor:
    """A tiny tensor that records operations so .backward() can apply the chain rule."""
    def __init__(self, data, requires_grad=False, parents=(), backward=None):
        self.data, self.grad = np.asarray(data, float), None
        self.requires_grad, self._parents, self._backward = requires_grad, parents, backward
    def __matmul__(self, o): return op(self.data @ o.data, (self, o), lambda g: (g @ o.data.T, self.data.T @ g))
    def __add__(self, o):    return op(self.data + o.data, (self, o), lambda g: (g, g.sum(0)))  # o: bias row
    def __sub__(self, o):    return op(self.data - o.data, (self, o), lambda g: (g, -g))
    def relu(self):          return op(np.maximum(0, self.data), (self,), lambda g: (g * (self.data > 0),))
    def square_mean(self):   return op((self.data ** 2).mean(), (self,), lambda g: (g * 2 * self.data / self.data.size,))
    def backward(self):
        order, seen = [], set()
        def visit(t):  # topological order: every node after the nodes it depends on
            if id(t) not in seen:
                seen.add(id(t)); [visit(p) for p in t._parents]; order.append(t)
        visit(self)
        self.grad = np.ones_like(self.data)
        for t in reversed(order):  # walk the graph backwards, passing gradients to parents
            if t._backward:
                for p, g in zip(t._parents, t._backward(t.grad)):
                    if p.requires_grad or p._parents:
                        p.grad = g if p.grad is None else p.grad + g  # sum over all uses

def op(data, parents, backward):
    return Tensor(data, parents=parents, backward=backward)

class SGD:
    def __init__(self, params, lr): self.params, self.lr = params, lr
    def zero_grad(self): [setattr(p, "grad", None) for p in self.params]
    def step(self):
        for p in self.params: p.data -= self.lr * p.grad

# Data: y = sin(2x) + noise, served in shuffled mini-batches like a DataLoader
rng = np.random.default_rng(0)
X = rng.uniform(-2, 2, (256, 1)); Y = np.sin(2 * X) + rng.normal(0, 0.1, X.shape)
def loader(bs=32):
    idx = rng.permutation(len(X))
    for i in range(0, len(X), bs):
        yield Tensor(X[idx[i:i + bs]]), Tensor(Y[idx[i:i + bs]])

# Model: 1 -> 32 -> 1 with ReLU; parameters are leaf tensors with requires_grad=True
W1, b1 = Tensor(rng.normal(0, 1, (1, 32)), True), Tensor(np.zeros(32), True)
W2, b2 = Tensor(rng.normal(0, 0.2, (32, 1)), True), Tensor(np.zeros(1), True)
model = lambda x: ((x @ W1 + b1).relu() @ W2) + b2
opt = SGD([W1, b1, W2, b2], lr=0.02)

for epoch in range(1, 301):
    for xb, yb in loader():
        loss = (model(xb) - yb).square_mean()  # forward pass and loss
        opt.zero_grad()                         # clear old gradients
        loss.backward()                         # backpropagation fills .grad
        opt.step()                              # gradient descent update
    if epoch in (1, 10, 50, 100, 300):
        full = (model(Tensor(X)) - Tensor(Y)).square_mean().data
        print(f"epoch {epoch:>3}: training MSE {full:.4f}")
print("noise floor (variance of the added noise): 0.0100")

epoch   1: training MSE 0.3872
epoch  10: training MSE 0.3726
epoch  50: training MSE 0.1327
epoch 100: training MSE 0.0332
epoch 300: training MSE 0.0127
noise floor (variance of the added noise): 0.0100

The Tensor class records, for each operation, its parents and a function that turns the upstream gradient into gradients for each parent: for a matrix multiply x @ W, that is g @ W.T for x and x.T @ g for W, the matrix rules from Backpropagation . backward() sorts the graph so every node comes after its inputs, then walks it in reverse, adding gradients into .grad. The training loop at the bottom has exactly the shape of the PyTorch one: batches from a loader, forward, loss, zero_grad, backward, step.

The training MSE falls from 0.387 after the first epoch to 0.0127 after 300, close to the noise floor of 0.0100: the variance of the noise added to the targets, which no model can predict.

What the framework adds on top

Real frameworks add hundreds of operations with hand-optimized forward and backward kernels, GPU execution, mixed precision, memory savings such as freeing intermediate values as soon as the backward pass has used them, and the ability to compile the graph for speed. But the contract you program against is the one in this file: build the forward computation from tracked operations, call backward on a scalar loss, read .grad.

Common shapes of the loop

• Validation: after each epoch, run model.eval() and a no-grad pass over held-out data, then model.train() again.
• Gradient accumulation: call backward() on several small batches before one step() and one zero_grad(), to simulate a larger batch than fits in memory.
• Gradient clipping: rescale gradients between backward() and step() when they explode ( Vanishing and exploding gradients, residual connections ).
• Checkpointing: save model.state_dict() and the optimizer's state regularly so you can resume.

Trade-offs

Common mistakes

• Forgetting optimizer.zero_grad(), so gradients from every previous batch add up.
• Calling backward() on a tensor that is not a scalar (reduce the loss to one number first).
• Leaving the model on the GPU and the batch on the CPU, or the other way round.
• Evaluating without model.eval() and torch.no_grad(), getting dropout noise and wasting memory.
• Using loss (a tensor attached to the graph) in a running total, which keeps every step's graph alive; use loss.item().

In the wild

This loop, with the same five lines at its core, trains everything from a small classifier on a laptop to large language models on thousands of GPUs. At scale, the loop is wrapped in data-parallel and model-parallel machinery, mixed precision and frequent checkpointing, but forward, loss, zero_grad, backward and step are still recognizable in the code. Reading any training script starts by finding them.

Check yourself

Recap