Leaf vs Intermediate Tensors and .retain_grad()

Only leaves get .grad — by default

Autograd makes a strong distinction between two kinds of tensors:

Leaf tensors — created by you (or by nn.Parameter). They are inputs to the graph, not outputs of any op. Their .grad is the gradient of the loss with respect to them.
Intermediate tensors — outputs of operations. PyTorch's default behavior is to NOT store gradients on them, because storing intermediates everywhere would balloon memory.

If you're debugging and want to see the gradient on an intermediate (the output of layer 3, say), you have two options:

x.retain_grad() — call this on the intermediate before backward, and PyTorch will store its gradient too.
x.register_hook(callback) — register a function that gets called with the gradient as it flows through. Doesn't store anything; lets you log or modify.

In a normal training loop you don't touch either. They're both for the moment when you go "huh, why is this layer not learning" and you reach for diagnostics.

Code

Why .grad is None on intermediates·python

import torch

a = torch.tensor(2.0, requires_grad=True)   # LEAF
b = a * 3                                    # INTERMEDIATE
c = b + 1                                    # INTERMEDIATE
c.backward()

print(a.is_leaf, a.grad)   # True tensor(3.)
print(b.is_leaf, b.grad)   # False None  (intermediate, .grad not stored)
print(c.is_leaf, c.grad)   # False None

retain_grad — diagnostic-only·python

import torch

a = torch.tensor(2.0, requires_grad=True)
b = a * 3
b.retain_grad()                       # tell autograd to keep b's grad
c = b + 1
c.backward()

print(b.grad)                         # tensor(1.) — now available

register_hook — modify or log gradients in-flight·python

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 2),
)

stats = {}
def make_hook(name):
    def _hook(grad):
        stats[name] = (grad.abs().mean().item(), grad.std().item())
    return _hook

for name, p in model.named_parameters():
    p.register_hook(make_hook(name))

x = torch.randn(16, 4)
y = torch.randn(16, 2)
loss = nn.functional.mse_loss(model(x), y)
loss.backward()

for name, (mean, std) in stats.items():
    print(f"{name}: |grad|.mean={mean:.4f}  std={std:.4f}")

Exercise

Build a 3-layer MLP. Attach a register_hook to each parameter that records the gradient mean and std. Train for one batch on random data. Print the stats. Then change the activation from ReLU to Sigmoid and rerun — observe the gradient magnitudes shrink toward zero in earlier layers (vanishing gradients).

Leaf vs Intermediate Tensors and .retain_grad()

Only leaves get .grad — by default

Code

External links

Exercise

Progress

Comments 0