PyTorch Custom Training Loops

This closing note of the PyTorch category covers when and why to write a fully custom training loop — rather than relying on a high-level trainer abstraction — and the practical patterns (gradient accumulation, custom logging, mixed precision) that come up when doing so.

When a Custom Loop Is the Right Choice

High-level training abstractions (like Hugging Face's Trainer, or PyTorch Lightning) handle the standard training loop pattern for you, but a fully custom loop — writing every step manually, exactly as covered throughout this category — gives complete control, which matters for: non-standard training procedures (like GAN's alternating generator/discriminator updates from GAN, or RLHF's multi-stage process), highly specific logging/debugging needs, or simply understanding exactly what's happening at every step, which is often valuable for learning and debugging even when a higher-level tool would otherwise suffice.

Gradient Accumulation — Simulating a Larger Batch Size

accumulation_steps = 4   # effectively simulates a 4x larger batch size
optimizer.zero_grad()

for i, (x_batch, y_batch) in enumerate(train_loader):
    predictions = model(x_batch)
    loss = loss_fn(predictions, y_batch) / accumulation_steps   # scale down, since gradients will SUM
    loss.backward()   # gradients ACCUMULATE across these steps (no zero_grad() in between)

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()      # apply the accumulated gradient
        optimizer.zero_grad()   # reset for the next accumulation cycle

This directly exploits autograd's default gradient-accumulation behavior (from Autograd) — instead of treating it as something to avoid via zero_grad() every step, gradient accumulation deliberately uses it to simulate training with a larger effective batch size than would otherwise fit in memory.

Custom Logging and Metric Tracking

metrics_history = {"train_loss": [], "val_loss": [], "val_accuracy": []}

for epoch in range(num_epochs):
    train_loss = run_training_epoch(model, train_loader, optimizer, loss_fn)
    val_loss, val_acc = validate(model, val_loader, loss_fn, device)

    metrics_history["train_loss"].append(train_loss)
    metrics_history["val_loss"].append(val_loss)
    metrics_history["val_accuracy"].append(val_acc)

    # full custom control: log to a file, a dashboard, or a tool like MLflow (covered in Production DL & MLOps)

Common Mistakes

• Forgetting to scale the loss by accumulation_steps in gradient accumulation — without this, the accumulated gradient's effective magnitude is accumulation_steps times larger than intended, distorting the effective learning rate.
• Calling optimizer.zero_grad() on every single batch during gradient accumulation, rather than only after the accumulation cycle completes — this defeats the entire purpose, since gradients need to accumulate across those batches.

Interview Relevance

Q: "How does gradient accumulation let you train with an effectively larger batch size than fits in GPU memory?" Instead of calling optimizer.zero_grad() and optimizer.step() after every single batch, gradients are allowed to accumulate (summing, via autograd's default behavior) across several consecutive smaller batches, with the loss scaled down by the number of accumulation steps to keep the effective gradient magnitude correct. Only after several batches' worth of gradients have accumulated is optimizer.step() finally called — simulating the effect of one large batch using several small ones sequentially.

Key Takeaways — PyTorch

• Tensors, autograd, and computational graphs are PyTorch's foundational mechanics — dynamic graph construction is what makes variable control flow (RNNs, conditionals) work naturally.
• nn.Module, Dataset, and DataLoader together form the standard structure of every real PyTorch project — a model, a data source, and a batching/loading pipeline.
• The standard training and validation loops assemble every concept from this hub into runnable code, with device management and gradient bookkeeping as the main practical gotchas.
• Custom training loops give full control when needed — gradient accumulation, GANs' alternating updates, and specialized logging all require dropping down to this level.

Next: TensorFlow & Keras covers the same practical territory — tensors, layers, training loops — in TensorFlow's Sequential and Functional APIs, with a direct comparison to everything just covered in PyTorch.

Practice Question

Why might you choose to write a fully custom training loop for implementing GAN training, rather than using a high-level, standard-pattern trainer abstraction?

Related DL Notes

• PyTorch Training Loop
• Autograd
• GAN