Training Loop

The training loop is where every concept from the previous nine categories comes together into one runnable procedure — this note assembles the complete, standard structure, annotated piece by piece.

The Complete Structure

import torch
import torch.nn as nn
import torch.optim as optim

model = MyModel()
optimizer = optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)

for epoch in range(num_epochs):                  # outer loop: one pass per epoch
    model.train()                                  # sets the model to TRAINING mode (matters for dropout, batch norm)
    running_loss = 0.0

    for X_batch, y_batch in train_loader:           # inner loop: one iteration per batch
        optimizer.zero_grad()                        # step 1: clear old gradients
        predictions = model(X_batch)                  # step 2: forward pass
        loss = loss_fn(predictions, y_batch)
        loss.backward()                                # step 3: backward pass -- computes gradients
        optimizer.step()                                # step 4: update weights

        running_loss += loss.item()

    scheduler.step()                                  # update the learning rate schedule, once per epoch

    avg_train_loss = running_loss / len(train_loader)
    print(f"Epoch {epoch+1}: train_loss={avg_train_loss:.4f}")

Why model.train() Matters

Some layers — dropout (randomly zeroing activations) and batch normalization (using batch statistics vs. running averages) — behave differently during training versus evaluation. Calling model.train() tells PyTorch to use the training-time behavior for these layers; forgetting it (or forgetting the corresponding model.eval() before validation/inference) is a genuinely common and confusing bug, covered in full in the Regularization and Normalization categories.

Tracking Loss — Per-Batch vs Per-Epoch

Where This Loop Connects to Earlier Categories

• Steps 2–4 are exactly the four-step cycle from Backpropagation Weight Updates.
• The optimizer and scheduler are exactly the objects covered throughout the Optimization & LR Scheduling category.
• The outer epoch loop and inner batch loop map directly onto the epoch/batch/iteration definitions from Epoch, Batch, Iteration.

Common Mistakes

• Placing scheduler.step() inside the inner batch loop for a scheduler designed to be stepped per-epoch (like StepLR) — this applies the decay far more often than intended.
• Forgetting model.train() at the start of each epoch, especially after having called model.eval() for a validation pass the epoch before — the model can silently stay in evaluation mode for training, disabling dropout and using stale batch norm statistics.

Interview Relevance

Q: "Walk through the four inner-loop steps of a standard PyTorch training iteration, in order, and explain why the order matters." Zero gradients (clear any leftover accumulation from the previous batch), forward pass (compute predictions and loss), backward pass (compute gradients for every parameter via backpropagation), optimizer step (apply the update using those gradients). Reversing steps 3 and 4, or skipping step 1, produces incorrect or unstable training, as detailed in Backpropagation Weight Updates.

Practice Question

Why is per-batch loss noisier than per-epoch average loss, in terms of what each one is actually measuring?

Related DL Notes

• Validation Loop
• Backpropagation Weight Updates
• PyTorch Training Loop