03 · The Training Loop 🔁

5 min read

This is the payoff. Every concept in the course so far snaps together here into the loop that makes a model learn.

1. The 'Why'

You now have all the pieces: a model (Module 03), batched data (Module 04), a loss (01), and an optimizer (02). The training loop is the small piece of code that orchestrates them into learning. Its job is repetitive by design: for every batch of data, predict, measure the error, compute gradients, and take an optimization step — then do it again, batch after batch, epoch after epoch, until the model is good. This loop is the single most-written pattern in all of deep learning, and once it's in your fingers you'll write it almost without thinking.

What makes the loop worth studying carefully is that its correctness depends on getting a specific sequence right, and on managing the model's mode. The canonical order — zero gradients, forward, loss, backward, step — is not arbitrary: shuffle it and you'll silently accumulate stale gradients or update with the wrong ones (the Module 02 gotcha made real). And because layers like Dropout behave differently while learning versus while being evaluated, you must switch the model into train() mode for the training pass and eval() mode for validation. This sub-module assembles the complete, correct loop — including a validation pass — that you'll reuse for every model in Phase 3.

2. Core Concepts

Epoch vs. batch. One batch is a single group of samples; one epoch is one full pass over the whole training set. The loop is two nested loops: for epoch: for batch:.

The canonical five steps (per batch). optimizer.zero_grad() → model(x) (forward) → criterion(pred, y) (loss) → loss.backward() → optimizer.step(). Memorize this order.

model.train() vs. model.eval(). Call model.train() before the training pass so Dropout/BatchNorm behave for learning; call model.eval() before validation so they behave for inference. These set a mode flag — they don't move data or touch gradients.

torch.no_grad() for validation. During the validation pass you're not learning, so wrap it in torch.no_grad() to skip graph-building — faster and lighter. Real evaluation uses both eval() and no_grad() (as previewed in Modules 02–03).

Move data to the device. The DataLoader yields CPU tensors; send each batch (and the model) to the same device inside the loop.

3. Code in Action

The minimal training loop

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

model = nn.Linear(4, 3)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

for epoch in range(10):                       # outer loop: epochs
    model.train()                             # training mode ON
    for batch_X, batch_y in train_loader:     # inner loop: batches
        batch_X = batch_X.to(device)          # data to the model's device
        batch_y = batch_y.to(device)

        optimizer.zero_grad()                 # 1. clear old gradients
        logits = model(batch_X)               # 2. forward pass
        loss = criterion(logits, batch_y)     # 3. compute loss
        loss.backward()                       # 4. backward pass
        optimizer.step()                      # 5. update weights

Adding a validation pass

import torch

for epoch in range(10):
    # ---- Train ----
    model.train()
    train_loss = 0.0
    for batch_X, batch_y in train_loader:
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(batch_X), batch_y)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()             # .item() -> Python float, detached

    # ---- Validate ----
    model.eval()                              # inference mode (Dropout off, etc.)
    val_loss = 0.0
    with torch.no_grad():                     # no graph, no gradients
        for batch_X, batch_y in val_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            val_loss += criterion(model(batch_X), batch_y).item()

    print(f"epoch {epoch}: "
          f"train {train_loss/len(train_loader):.4f} | "
          f"val {val_loss/len(val_loader):.4f}")

A reusable structure

def train_one_epoch(model, loader, criterion, optimizer, device):
    model.train()
    total = 0.0
    for X, y in loader:
        X, y = X.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(X), y)
        loss.backward()
        optimizer.step()
        total += loss.item()
    return total / len(loader)                # average loss for the epoch

4. Common Pitfalls

Pitfall 1 — Wrong step order. Calling optimizer.step() before loss.backward(), or forgetting zero_grad(), breaks learning silently. Burn the order in: zero → forward → loss → backward → step.

Pitfall 2 — Forgetting model.train() / model.eval(). If you validate without eval(), Dropout stays active and your validation numbers are noisy and pessimistic; if you train while stuck in eval(), regularization is off. Set the mode explicitly at the top of each phase.

Pitfall 3 — Accumulating the loss tensor instead of its value. Writing train_loss += loss (not loss.item()) keeps the whole computational graph alive across the epoch, leaking memory until you crash. Always accumulate loss.item() (or loss.detach()).

5. Further Reading & Watch List

⬅️ Prev: 02 · Optimizers · ➡️ Next: 04 · Evaluation & Metrics