🎓 A1Learner
🔍

PyTorch Checking Gradient Flow

Gradient flow is the movement of gradients from the loss back through each layer. A gradient norm is the length of one gradient tensor, showing whether the loss can still steer that layer.

One layer, one norm

One healthy layer gets a clear non-zero norm after backward. The next snippet shows the opposite case.

torch.manual_seed(0)
layer = nn.Linear(3, 2)
x = torch.randn(4, 3)
loss = layer(x).sum()
loss.backward()
print("Weight grad norm:", round(layer.weight.grad.norm().item(), 4))

A dead ReLU path

Zero weights and a bias of -1 make every pre-activation negative. ReLU turns them all to zero, and the weight gradient dies with them.

layer = nn.Linear(2, 2)
with torch.no_grad():
    layer.weight.zero_()
    layer.bias.fill_(-1.0)
x = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
out = torch.relu(layer(x))
print("Outputs all zero:", bool((out == 0).all()))
out.sum().backward()
print("Weight grad norm:", layer.weight.grad.norm().item())

Read the whole table

named_parameters names every parameter in order. One loop prints the full table for a small model.

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(3, 4), nn.ReLU(), nn.Linear(4, 2))
loss = model(torch.randn(5, 3)).sum()
loss.backward()
for name, param in model.named_parameters():
    print(name, round(param.grad.norm().item(), 4))

Example

import torch
from torch import nn


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.main1 = nn.Linear(4, 8)
        self.gate = nn.Linear(4, 4)
        self.dead = nn.Linear(4, 3)
        self.main2 = nn.Linear(8, 3)
        with torch.no_grad():
            self.gate.weight.zero_()
            self.gate.bias.fill_(-1.0)

    def forward(self, x):
        main = torch.relu(self.main1(x))
        dead = self.dead(torch.relu(self.gate(x)))
        return self.main2(main) + dead


torch.manual_seed(0)
net = Net()
x = torch.randn(6, 4)
y = torch.tensor([0, 1, 2, 1, 0, 2])

with torch.no_grad():
    gate_out = torch.relu(net.gate(x))
print("Gate output all zero:", bool((gate_out == 0).all()))

loss = nn.CrossEntropyLoss()(net(x), y)
loss.backward()
print("Loss:", round(loss.item(), 4))
print("Parameter | grad norm")
for name, param in net.named_parameters():
    print(name, round(param.grad.norm().item(), 4))

Key points

  • Gradient norm: the length of a gradient tensor, read with param.grad.norm()
  • Use one loss.backward() call, then read each parameter's .grad
  • A gradient norm of 0.0 means the loss cannot move that parameter
  • A ReLU whose inputs are all negative outputs only zeros and passes no gradient back
  • A bias can keep a gradient with all-zero inputs, because the bias is added to every output
  • The norm table gives one row per parameter, so dead layers stand out as zero rows

In this example, a small network carries one deliberately dead branch. The gate layer has zero weights and a bias of -1, so its outputs are all negative. ReLU turns those outputs to zero, and the first print confirms the gate output is all zero. The dead layer receives only zeros, so its weight gradient is exactly 0.0. The gate rows are also 0.0, because ReLU passes no gradient back through the zeros. The dead bias row reads 0.1547, because the bias is added after the zero inputs. The main path rows stay non-zero, so the table separates the live path from the dead one.

Note

A parameter that backward never reached has grad None, which is a different finding from a 0.0 norm

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

One row of a gradient norm table reads 0.0. What does it mean?