How autograd works

PyTorch · Expert · 16 min

Understanding the engine explains in-place errors, memory use in backward, and how to compute Hessians or meta-gradients.

During the forward pass, every differentiable op creates a node (grad_fn) holding references to the tensors it needs for its derivative (the saved tensors) and edges to the nodes of its inputs (grad_fn.next_functions). User-created tensors with requires_grad=True are leaves; they get an AccumulateGrad node that writes into .grad. Non-leaf tensors do not keep .grad unless you call retain_grad().

backward() runs a topological traversal from the output, computing vector-Jacobian products node by node, so the cost of a full gradient is a small constant times the forward cost. Saved tensors are freed as soon as their node has run, which is why the default graph is single-use.

In-place safety comes from version counters. Every tensor has a _version that in-place ops increment. When a node is executed, it checks that each saved tensor still has the version it saw; if not, you get one of the variables needed for gradient computation has been modified by an inplace operation.

torch.autograd.grad(outputs, inputs, create_graph=True) returns gradients that are themselves part of a graph, so you can differentiate again: Hessian-vector products, gradient penalties and meta-learning. For a functional, composable API, torch.func provides grad, vmap, jacrev and hessian.

Saved tensors dominate training memory; torch.autograd.graph.saved_tensors_hooks lets you compress or offload them.

Example

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x * x
z = y.sin()
print(z.grad_fn)                     # <SinBackward0 object at ...>
print(z.grad_fn.next_functions)      # ((<MulBackward0 object at ...>, 0),)
print(x.is_leaf, y.is_leaf)          # True False

(g,) = torch.autograd.grad(z, x, create_graph=True)   # dz/dx = 2x cos(x^2)
(h,) = torch.autograd.grad(g, x)                       # second derivative
print(round(g.item(), 4), round(h.item(), 4))         # -2.6146 10.8016

a = torch.ones(3, requires_grad=True)
b = a.exp()                 # exp saves its output for backward
b.add_(1)                   # in-place bump of b's version counter
try:
    b.sum().backward()
except RuntimeError as e:
    print(str(e).split(":")[0])
# one of the variables needed for gradient computation has been modified by an inplace operation

Remember

  • The graph is grad_fn nodes with saved tensors and edges to inputs
  • Version counters detect in-place edits of saved tensors
  • create_graph=True enables higher-order derivatives

Watch out for

  • Using retain_graph=True to silence errors, which hides a design problem and leaks memory.

Check yourself

  1. Why does modifying a tensor in place sometimes break backward()?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: pytorch.org/docs/stable/notes/autograd.html

Builds on

Unlocks

Nothing builds on this yet.