Autograd
PyTorch autograd automatically computes gradients of scalar outputs with respect to tensor inputs. It powers all neural network training through reverse-mode automatic differentiation.
Recipe
Quick-reference recipe card - copy-paste ready.
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x
y.backward()
print(x.grad) # dy/dx at x=2: 2*2+3 = 7
# Inference: no gradient tracking
with torch.no_grad():
pred = model(x)When to reach for this:
- Understanding why
loss.backward()andoptimizer.step()work. - Debugging gradient flow in custom layers or loss functions.
- Implementing custom autograd functions for non-standard operations.
- Disabling gradients during inference or evaluation.
Working Example
"""autograd.py - gradient computation and gradient checking."""
import torch
import torch.nn as nn
# Simple computation graph
w = torch.tensor([1.0, 2.0], requires_grad=True)
b = torch.tensor(0.5, requires_grad=True)
x = torch.tensor([3.0, 4.0])
# Forward: y = w . x + b, loss = y^2
y = (w * x).sum() + b
loss = y ** 2
loss.backward()
print("w.grad:", w.grad) # d(loss)/dw
print("b.grad:", b.grad) # d(loss)/db
# Model with autograd
model = nn.Linear(3, 1)
x_batch = torch.randn(16, 3)
target = torch.randn(16, 1)
model.zero_grad(set_to_none=True)
pred = model(x_batch)
loss = nn.functional.mse_loss(pred, target)
loss.backward()
for name, param in model.named_parameters():
print(f"{name}: grad norm = {param.grad.norm().item():.4f}")What this demonstrates:
- Leaf tensors (
w,b) accumulate gradients in.grad. loss.backward()traverses the graph in reverse.- Model parameters receive gradients automatically.
zero_grad(set_to_none=True)clears gradients efficiently.
Deep Dive
How It Works
- Each tensor operation records inputs and the backward function on a computation graph.
backward()applies the chain rule from the output back to leaves.- Only tensors with
requires_grad=Trueaccumulate.grad. - Non-leaf tensor gradients are freed after backward unless
retain_graph=True. torch.inference_mode()is stricter and faster thanno_grad()for inference.
Gradient Controls
| Context | Effect | Use |
|---|---|---|
requires_grad=True | Track operations | Model parameters, inputs needing grad |
torch.no_grad() | Disable tracking | Inference, metric computation |
torch.inference_mode() | Stricter no-grad | Production inference |
.detach() | Break from graph | Stop gradients through a tensor |
retain_graph=True | Keep graph after backward | Multiple backward calls |
Python Notes
# Custom autograd function
class MyReLU(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_inputGotchas
- Forgetting
zero_grad()- gradients accumulate across steps. Fix: calloptimizer.zero_grad(set_to_none=True)before each backward. - Calling backward twice without retain_graph - graph is freed after first backward. Fix:
loss.backward(retain_graph=True)or separate forward passes. - In-place operations on tensors needing grad - breaks autograd tracking. Fix: avoid
x.add_(1)on tracked tensors; usex = x + 1. - GPU memory from saved graphs - keeping references to loss tensors retains the graph. Fix: use
.item()for logging;del lossafter backward. - Gradients on integer tensors - only floating/complex tensors support autograd. Fix: cast to float before operations needing gradients.
- Mixing numpy and torch - in-place numpy changes on shared memory corrupt gradients. Fix:
.detach().cpu().numpy()for export.
Alternatives
| Alternative | Use When | Don't Use When |
|---|---|---|
| PyTorch autograd | Default for neural nets | You need symbolic math (use JAX) |
torch.func (functorch) | Per-sample gradients, vmap | Simple model training |
| Manual gradients | Teaching, verification | Production model code |
| JAX autograd | Functional transforms, TPU | Existing PyTorch ecosystem |
FAQs
What is a computation graph?
- A directed acyclic graph of tensor operations.
- Built dynamically on each forward pass (define-by-run).
- Freed after backward unless retained.
Why does loss have to be a scalar for backward?
backward()computes gradients of a scalar output.- For vector outputs, pass a gradient argument:
y.backward(gradient=torch.ones_like(y)). - Loss functions reduce to a scalar automatically.
What is the difference between no_grad and inference_mode?
inference_modedisables more autograd infrastructure - faster.- Use
inference_modefor production inference. no_gradis fine for validation loops during training.
How do I check if gradients are flowing?
for name, p in model.named_parameters():
if p.grad is None:
print(f"no grad: {name}")- Missing gradients indicate disconnected graph or frozen layers.
What does requires_grad=False on parameters do?
- Freezes the parameter - no gradient computed or updated.
- Use for transfer learning on early layers.
- Set
param.requires_grad = Falseormodel.freeze()patterns.
What is gradient clipping?
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)- Caps gradient magnitude before
optimizer.step(). - Prevents exploding gradients in RNNs and deep networks.
Can I compute second derivatives?
loss.backward(create_graph=True)
grad_grad = torch.autograd.grad(loss, model.parameters(), create_graph=True)create_graph=Truekeeps the graph for higher-order gradients.- Needed for some meta-learning algorithms.
What is set_to_none in zero_grad?
optimizer.zero_grad(set_to_none=True)sets.gradto None instead of zeroing.- Slightly faster and reduces memory on large models.
- Default recommended pattern in PyTorch 2.x.
How does autograd interact with torch.compile?
- Compiled models still use autograd for training.
- Inference-only compilation can skip autograd entirely.
- Test compiled training loops for numerical equivalence.
What are leaf tensors?
- Tensors created directly by the user (not from operations).
- Only leaf tensors have
.gradpopulated by default. - Use
retain_grad()on non-leaf tensors if needed.
Why do my gradients explode?
- Learning rate too high, deep unnormalized networks, or recurrent architectures.
- Fix with gradient clipping, lower LR, layer normalization.
How do I disable grad for part of the forward pass?
with torch.no_grad():
features = frozen_encoder(x)
output = trainable_head(features)- Common pattern in transfer learning.
Related
- PyTorch Basics - tensor fundamentals
- Training Loops - backward in practice
- Building Models with nn.Module - nn.Parameter
- Transfer Learning & Fine-Tuning - freezing layers
- GPUs & Mixed Precision - autocast and grad scaling
Stack versions: This page was written for Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, and uv 0.6+.