datarekha

Autograd — why PyTorch exists

The computation graph, requires_grad, and the magic behind loss.backward() — demystified by building a tiny autograd from scratch.

7 min read Intermediate Deep Learning Lesson 2 of 28

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

  1. Start at the loss (a scalar) with gradient = 1.
  2. Walk to each operation in reverse topological order.
  3. At each op, multiply the upstream gradient by the local Jacobian, and pass it to the inputs.
  4. When you reach a leaf tensor (one with requires_grad=True and 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.

Forward builds the graph; backward applies the chain rulebackward: gradient flows right → leftx = 3x² = 92·x = 6sum = 15y = 16+ 1local ×6local ×2∂y/∂x = (1 × 6) + (1 × 2) = 2x + 2 = 8 at x = 3

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

0/2
Q1You wrap your inference code in `with torch.no_grad():` — what's the main benefit?
Q2You call `loss.backward()` twice in a row without `optimizer.zero_grad()`. What's in `param.grad` now?

Sign in to track your progress

Completed lessons, your XP, level, and streak save to your account — it's free and takes a few seconds.

Practice this in an interview

All questions

Related lessons

Explore further

Skip to content