Training Loop — Module 6.2

How This Lesson Fits the Module

Module 6.1 taught neurons, loss, optimizers, and batches. The training loop is where those pieces run in sequence—every epoch, every batch—until weights converge. Every other lesson in 6.2 (validation, checkpoints, schedulers) wraps around this core loop.

PyTorch does not hide training behind a single .fit() call. You own the loop, which means you also own reproducibility, logging, and failure recovery.

Module 6.1 Bridge You already know forward propagation, backpropagation, loss functions, and Adam. This lesson wires them into a repeatable epoch/batch cycle.

Learning Objectives

By the end of this lesson, students should be able to:

Describe the five-step batch cycle: forward, loss, backward, step, zero_grad.
Structure an outer epoch loop over a DataLoader.
Set model.train() and move tensors to the correct device.
Log per-batch and per-epoch training loss for debugging.
Explain why optimizer.zero_grad() must precede each backward pass.
Identify common bugs: forgotten .to(device), gradients not cleared, loss not scalar.

The Canonical Training Loop

Training repeats two nested loops: an epoch sweeps the full dataset once; inside each epoch, batches feed the GPU memory budget. Each batch runs the same five operations.

Step	PyTorch call	Purpose
1. Forward	`outputs = model(inputs)`	Compute predictions
2. Loss	`loss = criterion(outputs, targets)`	Scalar objective to minimize
3. Backward	`loss.backward()`	Accumulate gradients in `.grad`
4. Update	`optimizer.step()`	Apply optimizer rule to weights
5. Clear	`optimizer.zero_grad()`	Reset gradients before next batch

Minimal PyTorch Training Loop

Below is the pattern every production trainer extends. Keep it readable—complexity belongs in helper functions, not nested logic.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyClassifier().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    for batch_idx, (inputs, targets) in enumerate(loader):
        inputs, targets = inputs.to(device), targets.to(device)

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    return running_loss / len(loader)

for epoch in range(1, num_epochs + 1):
    train_loss = train_one_epoch(model, train_loader, criterion, optimizer, device)
    print(f"Epoch {epoch:03d} | train_loss={train_loss:.4f}")

Critical Mistake — Forgetting zero_grad()

PyTorch accumulates gradients by default. Skipping optimizer.zero_grad() adds new gradients onto old ones, producing nonsense updates. Call it once per batch (or once per accumulation group if using gradient accumulation).

DataLoader and Device Placement

DataLoader handles shuffling, batching, and optional multi-worker loading. Tensors returned by the loader live on CPU until you explicitly .to(device). Pin memory (pin_memory=True) speeds host-to-GPU copies when training on CUDA.

train_loader = DataLoader(
    train_dataset,
    batch_size=128,
    shuffle=True,
    num_workers=4,
    pin_memory=torch.cuda.is_available(),
)

# Always move BOTH inputs and targets
inputs = inputs.to(device, non_blocking=True)
targets = targets.to(device, non_blocking=True)

Logging and Reproducibility Hooks

Engineering-grade loops log epoch number, learning rate, and loss. Set seeds for torch, random, and numpy at the start of the script. Record the git commit hash and hyperparameters in the same log line as metrics.

Engineering Habit — Extract train_one_epoch

Keep the outer script thin: one function for training, one for validation (next lesson), one for saving checkpoints. Testable functions beat 200-line scripts.

Knowledge Check

Short Answer: What are the five steps in a training batch? Answer: Forward, loss, backward, optimizer step, zero_grad.
True/False: model.train() is optional during training. Answer: False—it enables dropout and batch-norm training behavior.
Multiple Choice: Gradients accumulate unless you: (a) call zero_grad, (b) call model.eval(), (c) use no_grad. Answer: (a).
Short Answer: Why shuffle the training loader? Answer: Reduces correlation between consecutive batches, improving SGD convergence.
Short Answer: What does loss.item() return? Answer: A Python float detached from the computation graph.
True/False: DataLoader automatically puts batches on GPU. Answer: False—you must call .to(device).
Multiple Choice: An epoch is: (a) one batch, (b) one full pass over training data, (c) validation. Answer: (b).
Short Answer: When should optimizer.step() run relative to loss.backward()? Answer: After backward, once gradients exist.
Short Answer: What does pin_memory=True help with? Answer: Faster CPU→GPU tensor transfers during training.
Multiple Choice: Training loss should be logged: (a) never, (b) per epoch at minimum, (c) only on test set. Answer: (b).

Key Takeaways

The training loop is forward → loss → backward → step → zero_grad, nested in epoch and batch loops.
Always call model.train(), move data to device, and clear gradients each batch.
Extract epoch logic into functions; the outer script orchestrates epochs and logging.
Next: Validation Loop—evaluate without updating weights.

Trainer’s Guide

Hands-on idea: Deliberately omit zero_grad() once and plot loss—students see divergence immediately. Fix it and compare curves.

Discussion prompt: Where would you insert gradient clipping or mixed precision in this loop? (Covered later in the module.)

What’s Next Open Validation Loop to add an evaluation pass each epoch without weight updates.