Cheat sheets

Debugging Neural Network Training

Find silent training bugs from the symptom. Checks, thresholds, init rules, clipping and AMP order, profilers.

  • Concepts and terms
  • Tasks and commands
  • Annotated code
  • Diagrams
  • Defaults and numbers
  • Troubleshooting
  • Questions and tips

Source

Symptom, cause, first fix#

SymptomCauseFix
Loss stuck at -log(1/num_classes)Model predicts a uniform distributionCheck the data pipeline and the labels
Loss NaN after a few stepsLearning rate too highReduce LR by 10x
Loss NaN immediatelylog(0) or division by zeroAdd epsilon to log and division
Loss oscillates wildlyLR too high or batch too smallReduce LR, increase batch size
Loss decreases, then plateausLR too high for the fine-tuning phaseAdd a cosine or step decay schedule
Train acc high, test acc lowOverfittingAdd dropout, weight decay, more data
Train acc = test acc = chanceModel learns nothingRun the overfit-one-batch test
Gradients all zeroDead ReLUs or a detached graphSwitch to LeakyReLU, check .requires_grad

From Lesson 03.13

Common PyTorch bugs#

SymptomCauseFix
Loss oscillates, gradients accumulateNo optimizer.zero_grad()Call it before loss.backward()
Test accuracy varies between runsNo model.eval() at test timeAdd model.eval() and torch.no_grad()
Wrong results, no errorSilent broadcasting of shapesPrint shapes after every operation
RuntimeError: expected CUDA tensorCPU and GPU mismatchUse .to(device) on model AND data
Graph grows forever, OOMTensors not detachedUse .detach() or with torch.no_grad()
RuntimeError: modified by in-place operationIn-place op breaks autogradReplace x += 1 with x = x + 1
Loss stuck at random-chance levelInputs not normalizedNormalize inputs to mean=0, std=1
Cross-entropy expects Long, got FloatLabels have the wrong dtypeCast labels: labels.long()

From Lesson 03.13, Lesson 00.12

Thresholds the checks use#

SettingDefaultNote
fraction_zero> 0.5Dead neurons. Use LeakyReLU or a lower LR
Activation std< 1e-6Collapsed activations
Gradient abs_mean< 1e-7Vanishing gradient in that layer
Gradient abs_mean> 100Exploding gradient in that layer
First / last grad ratio> 100Vanishing gradients through depth
rel_diff< 1e-5Gradient check passes. > 1e-3 is a bug
Overfit final_loss> 0.1FAIL: model or loop is broken
start_lr1e-7LR finder sweeps up to end_lr=10

From 03.13 code/debug_neural_nets.py, Lesson 03.13, 00.12 code/debug_tools.py

One AMP training step#

self.optimizer.zero_grad(set_to_none=True)
with torch.amp.autocast(device_type=self.device_type, dtype=self.amp_dtype):
loss = self._loss_fn(self.model(inputs), targets)
if not torch.isfinite(loss).all().item():skip, no scaler.update()
...
self.scaler.scale(loss).backward()
self.scaler.unscale_(self.optimizer)before the clip
if has_non_finite_grad(self.model.parameters()):
self.scaler.update()halves the scale
pre_clip, post_clip = clip_global_l2_norm(list(self.model.parameters()), self.max_norm)
self.scaler.step(self.optimizer)
self.scaler.update()every step
scale = self.scaler.get_scale()log next to lr and norm

From 19.45 code/main.py, Lesson 19.45

Loss not decreasing: check in this order#

1Learning ratetoo high, loss oscillates or explodes. Too low,loss barely moves.2Learning ratetry 3 rates 10x apart, for example 1e-2, 1e-3,1e-4.3Gradientsall zeros means dead ReLUs or vanishinggradients.4GradientsNaN or Inf means exploding gradients.5Data pipelineshuffled labels give random-chance accuracy.6Data pipelinea preprocessing bug makes the model learn noise.7Architecturetoo small underfits. Too deep is hard tooptimize.

From Lesson 03.13

Diagnostic functions#

TaskCode
Shape, dtype, device, NaN of a tensordebug_print(name, tensor)
Map every layer shapecheck_shapes(model, sample_input)
Find NaN or Inf gradientsdetect_nan(model, loss, step)
Find tensors on another devicecheck_devices(model, *tensors)
Total norm, large or zero gradscheck_gradient_health(model)
Prove the loop can learnoverfit_one_batch(model, x_batch, y_batch, criterion)
Sweep the learning ratefind_learning_rate(model, x_data, y_data, criterion)
Compare backprop to finite differencesgradient_check(model, x, y, criterion, eps=1e-4)

From 00.12 code/debug_tools.py, 03.13 code/debug_neural_nets.py

Hook a debugger into a model#

debugger1 = NetworkDebugger(model1)hooks Linear, Conv2d, ReLU, LeakyReLU
debugger1.record_loss(loss.item())call once per step
debugger1.print_report()loss, activation, gradient checks
debugger1.remove_hooks()
with torch.autograd.detect_anomaly():built-in NaN and Inf detection
loss.backward()

From 03.13 code/debug_neural_nets.py, Lesson 03.13

Initialization by activation#

Zero init
all neurons stay identical. One effective neuron for any width.
Variance rule
\mathrm{Var}(z) = \text{fan}_{\text{in}} \cdot \mathrm{Var}(w) \cdot \mathrm{Var}(x)
Xavier
\mathrm{Var}(w) = 2/(\text{fan}_{\text{in}} + \text{fan}_{\text{out}}) for sigmoid and tanh.
Kaiming
\mathrm{Var}(w) = 2/\text{fan}_{\text{in}} for ReLU, Leaky ReLU, GELU and Swish.
Residual scaling
scale residual layers by 1/\sqrt{2N} for N layers.
Random N(0,1)
explodes by layer 5. N(0,0.01): vanishes by layer 10.
Check
activation magnitudes stay between 0.5 and 2.0 through all layers.
PyTorch
nn.Linear(512, 256) defaults to Kaiming uniform.

From Lesson 03.08, 03.08 code/main.py

Clip and scale step order#

finitecleanprocessForward inautocastloss in FP32processisfinite loss?no: skip stepgpuscale(loss).backwardFP16 gradsprocessunscale_optimizertrue gradientsprocessNaN or Inf grad?yes: skip, updatehalvesprocessclip_grad_norm_max_norm 1.0processscaler.stepoptimizer stepstorescaler.updatedoubles or halvesfinitecleanprocessForward in autocastloss in FP32processisfinite loss?no: skip stepgpuscale(loss).backwardFP16 gradsprocessunscale_ optimizertrue gradientsprocessNaN or Inf grad?yes: skip, updatehalvesprocessclip_grad_norm_max_norm 1.0processscaler.stepoptimizer stepstorescaler.updatedoubles or halves

From Lesson 19.45

Profile and inspect#

CommandWhat it does
python -m cProfile -s cumtime train.pyEvery function call, sorted by cumulative time
kernprof -l -v train.pyLine by line time for @profile functions
python -m memory_profiler your_script.pyLine by line memory use
tracemalloc.take_snapshot()CPU memory by line after tracemalloc.start()
torch.cuda.memory_summary()GPU memory allocated and cached
tensorboard --logdir=runsLoss curves, weight and gradient histograms
breakpoint()Stop in pdb. p outputs.shape, c continue, q quit
python3 code/main.pyClip and AMP demo. At least one row is a SKIP

From Lesson 00.12, Lesson 19.45

Numbers to watch#

What does a healthy GradScaler scale do?
It climbs in powers of two to near 2^17 or 2^18.
When is the skip rate an alert?
Above 5 percent over a 1,000-step rolling window.
The scale never changes?
The skip path forgets scaler.update().
Data loading takes 60% of the time?
Set num_workers > 0 in the DataLoader.
Out of memory?
Reduce the batch size first. Then use torch.cuda.empty_cache().
Test accuracy is 99%?
Check for data leakage. Split first, preprocess second.
Train loss goes down and val loss goes up?
The model overfits.

From Lesson 19.45, Lesson 00.12, Lesson 03.13

Checklist before full training#

  • Run the overfit-one-batch test. If it fails, stop.
  • Print the model summary. Verify the parameter count.
  • Run one forward pass with random data. Check the output shape.
  • Train for 5 epochs. Verify that the loss decreases.
  • Check activation statistics and gradient flow.
  • Print 5 random samples with labels.
  • Log loss, learning rate and gradient norms during training.

From Lesson 03.13, Lesson 00.12