Autograd — why PyTorch exists
The computation graph, requires_grad, and the magic behind loss.backward() — demystified by building a tiny autograd from scratch.
What you'll learn
- What `requires_grad=True` actually tracks
- How `loss.backward()` walks the computation graph
- When to use `.detach()` and `torch.no_grad()`
Before you start
Before autograd, training a neural network meant writing the backward pass by hand — page after page of chain-rule calculus, debugged with finite differences. Autograd is the feature that turned deep learning from a research project into an engineering discipline. You write the forward pass; PyTorch gives you backward for free.
The mental model
Autograd (automatic differentiation) means PyTorch records every math operation you do, then automatically computes derivatives by replaying that record in reverse.
Every tensor with requires_grad=True is a node in a graph — it tells
PyTorch “I might need the derivative with respect to this value later.”
Every operation on it adds an edge. When you call .backward() on a scalar,
PyTorch walks the graph backward — applying the chain rule at each node
— and accumulates the gradient into each input’s .grad attribute.
import torch
x = torch.tensor(3.0, requires_grad=True)
y = x ** 2 + 2 * x + 1
y.backward() # dy/dx at x=3 is 2*3 + 2 = 8
print(x.grad) # tensor(8.)
You wrote y = x**2 + 2*x + 1. PyTorch silently built a graph of pow,
mul, add operations and knows the derivative of each. Backward
multiplies them together in the right order.
Build your own (NumPy)
Let’s build a tiny autograd to see what’s happening. Each Node holds a
value, knows how it was made, and can propagate gradients backward.
import numpy as np
class Node:
def __init__(self, value, parents=(), op=""):
self.value = float(value)
self.parents = parents # list of (parent_node, local_grad)
self.op = op
self.grad = 0.0
def __add__(self, other):
other = other if isinstance(other, Node) else Node(other)
# d(a+b)/da = 1, d(a+b)/db = 1
return Node(self.value + other.value,
parents=[(self, 1.0), (other, 1.0)], op="+")
def __mul__(self, other):
other = other if isinstance(other, Node) else Node(other)
# d(a*b)/da = b, d(a*b)/db = a
return Node(self.value * other.value,
parents=[(self, other.value), (other, self.value)], op="*")
def __pow__(self, n):
# d(a**n)/da = n * a**(n-1)
return Node(self.value ** n,
parents=[(self, n * self.value ** (n - 1))], op=f"**{n}")
def backward(self, upstream=1.0):
self.grad += upstream
for parent, local in self.parents:
parent.backward(upstream * local)
# Now: same computation as before, but with our handmade autograd.
x = Node(3.0)
y = x ** 2 + x * 2 + Node(1.0)
y.backward()
print(f"y = {y.value}") # 16
print(f"dy/dx = {x.grad}") # 8.0
y = 16.0
dy/dx = 8.0
That’s the whole concept. PyTorch’s autograd is the same idea — every op records (parent, local_gradient) — but written in C++, vectorized over tensors, with hundreds of operators defined, and a smart graph traversal that handles diamond shapes (a node used twice).
When you DON’T want gradients
Two situations come up constantly:
Evaluation / inference. You’re not training — you just want predictions. Tracking the graph wastes memory and time.
model.eval()
with torch.no_grad():
pred = model(x)
torch.no_grad() disables graph construction inside its block. Memory
drops, speed goes up. Forget this and your inference server will OOM
under load.
Detaching a branch. Sometimes you want a value to be treated as a constant — the gradient should NOT flow through it. The classic case is target networks in reinforcement learning, or fixing a teacher model’s output during distillation.
teacher_logits = teacher(x).detach() # no gradient flows back into teacher
student_logits = student(x)
loss = kl_div(student_logits, teacher_logits)
loss.backward() # only updates student
detach() returns a new tensor sharing the same data but disconnected
from the graph. Use it any time you want “the value but not the
history.”
What backward() actually does, step by step
- Start at the loss (a scalar) with gradient = 1.
- Walk to each operation in reverse topological order.
- At each op, multiply the upstream gradient by the local Jacobian, and pass it to the inputs.
- When you reach a leaf tensor (one with
requires_grad=Trueand no parents — a model parameter), add the accumulated gradient to.grad.
After backward, every model parameter has a .grad of the same shape as
itself. The optimizer reads those and updates the weights.
In one breath
- Autograd records every operation on a tensor with
requires_grad, building a computation graph as the forward pass runs. loss.backward()walks that graph in reverse, multiplying local gradients (the chain rule) and accumulating each leaf’s gradient into.grad.- You write only the forward pass; the backward is free — the feature that turned deep learning into engineering.
- Wrap inference in
torch.no_grad()to skip graph-building (saves memory and time); use.detach()to treat a value as a constant the gradient can’t flow through. backward()only works on a scalar — reduce a vector loss with.mean()or.sum()first.
Quick check
Quick check
Practice this in an interview
All questionsThe forward pass feeds an input through every layer in sequence: each layer computes a linear transform followed by an activation, caching the intermediate values needed later for backpropagation. The final layer produces a prediction, which is compared to the label via a loss function.
PyTorch accumulates gradients by default, adding new gradients to whatever is already stored in each parameter's .grad. If you do not zero them out each iteration, gradients from previous batches mix with the current batch and corrupt the weight updates. zero_grad() resets gradients to zero so each step uses only the current batch's signal.
Gradient accumulation sums gradients over multiple small forward-backward passes before calling the optimizer, simulating a larger effective batch size without requiring the memory to hold it all at once. It is the standard workaround when the desired batch size does not fit in GPU memory.
Backpropagation is the algorithm that computes the gradient of the loss with respect to every parameter by applying the chain rule layer by layer in reverse. It turns a single backward pass through the computation graph into exact gradients for all weights simultaneously.