Deep Learning

deep-learning — train a neural net in PyTorch without the silent bugs

You own training and understanding neural networks in PyTorch: the loop, autograd, mixed precision, optimizers/schedulers, multi-GPU, and the reproducibility/checkpoint hygiene that separates a real result from a lucky one. Nets from scratch, vision, custom architectures — all here. This is PyTorch-first by design: JAX and TensorFlow are real and fine, but the patterns, APIs, and gotchas below are Torch's.

Version reality (verify at author time). Current stable is PyTorch 2.x — ~2.13 as of mid-2026 (pytorch.org/get-started, releases move fast; don't hard-pin a minor). Everything below is stable 2.x API. The one namespace shift to know: AMP now lives under torch.amp (torch.amp.GradScaler("cuda")), not the old torch.cuda.amp.*.

Am I in the right skill?

1. PyTorch essentials

Three objects carry everything.

• Tensor — an n-d array on a device (cpu/cuda/mps) with a dtype. requires_grad=True makes autograd track ops on it. .to(device) / .detach() / .item() are the moves you use constantly; .item() pulls a Python scalar and drops the graph (see the loop bugs below).
• autograd — a tape. Every op on a requires_grad tensor records a node; loss.backward() walks it and accumulates into each leaf's .grad. "Accumulates" is the word that bites people (§2). Wrap read-only regions in torch.no_grad() to skip taping.
• nn.Module — the model container. __init__ registers submodules/params; forward defines compute. model.parameters() feeds the optimizer; model.train() / model.eval() flip train-vs-eval behavior for Dropout and BatchNorm.

import torch, torch.nn as nn
device = "cuda" if torch.cuda.is_available() else "cpu"

class MLP(nn.Module):
    def __init__(self, d_in, d_h, d_out, p=0.1):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d_in, d_h), nn.ReLU(),
                                 nn.Dropout(p), nn.Linear(d_h, d_out))
    def forward(self, x):
        return self.net(x)

model = MLP(784, 256, 10).to(device)

torch.compile — wrap the model once and get a JIT-optimized (Dynamo→Inductor) graph, often a real speedup with no code change: model = torch.compile(model). Modes: "default" (balanced), "reduce-overhead" (CUDA graphs, best for small batches), "max-autotune" (Triton autotuning, slowest to warm up). First step(s) are slow — that's compilation, not your model. Keep it after .to(device) and before the loop. (pytorch.org/docs/stable/generated/torch.compile.html)

2. The training loop — and its classic bugs

The whole game is five ordered steps. Get the order or the resets wrong and it fails silently — no exception, just a model that won't learn.

model.train()                                    # Dropout/BN in TRAIN mode
for epoch in range(epochs):
    for x, y in train_loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad(set_to_none=True)    # clear LAST step's grads
        out  = model(x)                          # forward
        loss = loss_fn(out, y)                    # loss  (raw logits in, e.g. CrossEntropyLoss)
        loss.backward()                          # autograd fills .grad (ACCUMULATES)
        optimizer.step()                         # update weights from .grad
        scheduler.step()                         # if a per-step scheduler

The load-bearing traps, stated next to the lines they govern:

• Missing optimizer.zero_grad() → gradients silently add up across batches. PyTorch accumulates .grad by design (so you can split a batch); if you never zero it, every step updates on the sum of all prior batches' grads and training destabilizes. Zero every step. (pytorch.org/tutorials optimization tutorial; torch.optim docs.) set_to_none=True is the default in modern Torch and is slightly faster.
• CrossEntropyLoss eats raw logits. It applies log_softmax internally — do not put a Softmax/Sigmoid before it or you double-activate and learning stalls. Same for BCEWithLogitsLoss.
• Accumulating loss without .item() leaks memory. total += loss keeps the whole autograd graph alive across the epoch. Use total += loss.item() (or .detach()).
• backward() twice on one graph raises unless you meant it (retain_graph=True) — usually a sign you forgot to zero or re-ran forward.

3. Mixed precision (AMP)

Run forward/loss in low precision for speed and memory, keep a master copy for the update. Two pieces: torch.autocast (picks per-op precision) + torch.amp.GradScaler (rescues fp16 gradients from underflow). Verified against pytorch.org/docs/stable/notes/amp_examples.html.

scaler = torch.amp.GradScaler("cuda")            # current namespace (not torch.cuda.amp)
for x, y in train_loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        out  = model(x)                          # forward + loss ONLY under autocast
        loss = loss_fn(out, y)
    scaler.scale(loss).backward()                # backward OUTSIDE autocast, on scaled loss
    scaler.step(optimizer)                        # unscales, skips step if inf/NaN grads
    scaler.update()                               # adjusts the scale for next step

bf16 vs fp16 — the choice that matters:

• With bf16 you can skip GradScaler entirely (GradScaler(enabled=False) or just don't use it) — its whole job is fp16 underflow. With fp16 the scaler is not optional.
• Gradient clipping under AMP: scaler.unscale_(optimizer) first, then torch.nn.utils.clip_grad_norm_(...), then scaler.step. Clipping scaled grads clips the wrong magnitude.
• Autocast wraps forward + loss only. Never wrap backward().

4. Optimizers and LR schedules

• torch.optim.AdamW is the default for most modern nets (transformers, and increasingly vision): Adam with decoupled weight decay (correct L2, unlike plain Adam(weight_decay=)). SGD + momentum + Nesterov still wins on many CNNs and generalizes well — worth trying.
• Don't weight-decay everything. Biases and Norm (LayerNorm/BatchNorm) weights should get weight_decay=0; use two param groups. Decaying them hurts with no upside.
• The LR is the hyperparameter. Too high → loss NaNs or oscillates; too low → crawls. A warmup then cosine decay is the reliable modern recipe.

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01, betas=(0.9, 0.999))

# Warmup (linear ramp) THEN cosine decay, composed with SequentialLR — verified 2.x API.
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
warmup = LinearLR(opt, start_factor=0.1, total_iters=500)          # 500 steps ramp-in
cosine = CosineAnnealingLR(opt, T_max=total_steps - 500, eta_min=3e-5)
scheduler = SequentialLR(opt, [warmup, cosine], milestones=[500])  # switch at step 500
# call scheduler.step() every STEP here (T_max/total_iters are in steps).

Alternatively OneCycleLR(opt, max_lr=..., total_steps=...) bakes warmup+anneal into one scheduler. Match step() cadence to the scheduler's units — per-step schedulers stepped once per epoch (or vice-versa) silently give the wrong curve. (pytorch.org/docs/stable/optim.html)