Debugging Neural Networks

Normal code crashes when something's wrong. Neural networks don't. They just quietly train badly, and they give you no error message. That makes them some of the hardest things to debug. Experienced engineers don't guess. They follow a method, and you're about to learn it.

• Neural networks fail silently
• No error message
• Follow a method

The debugging ladder

Climb this ladder, one rung at a time. First, check your data, by looking at real batches, with their labels. Second, check the first loss, which should match what random guessing predicts. Third, overfit one small batch. Fourth, watch your gradients. Only then should you try bigger models, or tune settings. Most bugs fall at the first three rungs.

1. Look at datareal batches, real labels
2. Check first losslog of class count
3. Overfit one batchthe great test
4. Watch gradientsare they flowing?

The overfit test

Here's the single most useful test in deep learning. Take one tiny batch, and try to memorise it. A healthy model and training loop should drive the loss close to zero. These labels are pure random, and the model can still memorise them. Run it. The loss falls to almost nothing. If this fails, there's a bug in your code, in your loss, or in your labels, and more data won't fix it. It takes one minute, and it can save you days.

import torch
from torch import nn

torch.manual_seed(0)
x = torch.rand(16, 20)
y = torch.randint(0, 3, (16,))

model = nn.Sequential(nn.Linear(20, 128), nn.ReLU(), nn.Linear(128, 3))
opt = torch.optim.Adam(model.parameters(), lr=0.01)

for step in range(201):
    loss = nn.functional.cross_entropy(model(x), y)
    opt.zero_grad(); loss.backward(); opt.step()
    if step in (0, 50, 200):
        print(f"step {step:>3}  loss {loss.item():.4f}")

step   0  loss 1.1026
step  50  loss 0.0049
step 200  loss 0.0005

Looking at gradients

When training is stuck, look inside. After one backward pass, we print the size of the weights and the gradients in each layer. Run it. Healthy gradients are small, but not zero. If some layer's gradient is exactly zero, then it's disconnected, or dead. If they're huge, you'll get an explosion, and a loss of not a number. The last line checks for those bad values. Gradient clipping is the usual fix for explosions.

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
x, y = torch.rand(8, 4), torch.randint(0, 2, (8,))

loss = nn.functional.cross_entropy(model(x), y)
loss.backward()

for name, p in model.named_parameters():
    size = p.grad.abs().mean().item()
    print(f"{name:<9} weight {p.abs().mean().item():.3f}  gradient {size:.5f}")

print("Any bad numbers?", any(torch.isnan(p).any().item() for p in model.parameters()))

0.weight  weight 0.265  gradient 0.03175
0.bias    weight 0.228  gradient 0.07218
2.weight  weight 0.232  gradient 0.06466
2.bias    weight 0.283  gradient 0.28781
Any bad numbers? False

Common silent bugs

Here are the classic bugs. Forgetting to call zero grad. Forgetting model dot eval when testing. Applying softmax before cross entropy, which already does it. Forgetting to shuffle training data. Labels that don't line up with images. And the wrong normalisation. Keep this list handy, because every engineer makes these mistakes, even after years of work.

Summary