PyTorch training loop

Milestone 2: from pieces to a training repo

This lesson closes the foundations module by assembling everything you have learned — splits, loss, curves, embeddings intuition, autograd — into a clean, reproducible PyTorch training project. By the end, your portfolio should have a folder that a reviewer can clone, run, and understand in five minutes.

A production-quality training loop is not just a for loop over epochs. It includes data loading, device placement, logging, checkpointing, evaluation mode, and a README that documents how to reproduce results. You are building the template you will reuse when you train language models in the next module.

Callout — the loop is the product: Teams iterate on datasets, architectures, and hyperparameters — but the training harness stays stable. Invest in a readable loop once; reuse it everywhere.

Anatomy of a PyTorch training loop

Here is the canonical structure. Read it top to bottom; each block maps to a concept from earlier lessons.

import torch
from torch.utils.data import DataLoader

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

best_val_loss = float("inf")

for epoch in range(num_epochs):
    # --- 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()
        logits = model(batch_x)
        loss = criterion(logits, batch_y)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    # --- VALIDATE ---
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            logits = model(batch_x)
            val_loss += criterion(logits, batch_y).item()

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

    # --- CHECKPOINT ---
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save({
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "val_loss": val_loss,
        }, "checkpoints/best.pt")

Key details easy to get wrong:

Building the MLP

For Milestone 2, a multilayer perceptron on a tabular or MNIST-style task is sufficient. Structure your model as an nn.Module:

class MLP(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, num_classes):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(input_dim, hidden_dim),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.1),
            torch.nn.Linear(hidden_dim, num_classes),
        )

    def forward(self, x):
        return self.net(x)

Principles:

• Keep forward simple — just the computation graph.
• Register layers in __init__ so model.parameters() finds them.
• Match output dimension to num_classes for cross-entropy; use 1 output + BCEWithLogitsLoss for binary.

Dataset and DataLoader

from torch.utils.data import Dataset, DataLoader, random_split

train_ds, val_ds = random_split(full_dataset, [0.8, 0.2])
train_loader = DataLoader(train_ds, batch_size=64, shuffle=True)
val_loader = DataLoader(val_ds, batch_size=64, shuffle=False)

Use shuffle=True only for training. Set num_workers > 0 on Linux/macOS for faster loading once your pipeline is stable (start with 0 while debugging).

For tabular data, consider loading from CSV with pandas and wrapping in a custom Dataset. For images, torchvision.datasets provides MNIST, CIFAR, etc. out of the box.

Checkpointing: save more than weights

A checkpoint should let you resume training and run inference without retraining.

Minimum viable checkpoint:

torch.save(model.state_dict(), "model.pt")

Better checkpoint (recommended):

torch.save({
    "model_state_dict": model.state_dict(),
    "optimizer_state_dict": optimizer.state_dict(),
    "epoch": epoch,
    "val_loss": val_loss,
    "config": {"lr": 1e-3, "hidden_dim": 128},
}, "checkpoints/best.pt")

Loading:

checkpoint = torch.load("checkpoints/best.pt", map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

For inference-only deployment, you may ship just state_dict plus a config JSON. For research and course projects, save the full checkpoint so you can resume after Colab disconnects.

Callout — map_location matters: Loading a GPU checkpoint on CPU requires map_location=torch.device("cpu") or you get device errors.

Logging and reproducibility

At minimum, log per-epoch train and val loss to a JSON file (you did this in lesson 2.2). Better:

• Fix random seeds (torch.manual_seed, numpy, Python random) for reproducibility.
• Save hyperparameters in a config.yaml or argparse flags.
• Print parameter count: sum(p.numel() for p in model.parameters()).
• Record wall-clock time per epoch.

Your README should include:

pip install -r requirements.txt
python train.py --epochs 20 --lr 1e-3
python eval.py --checkpoint checkpoints/best.pt

A reviewer should reproduce your val loss within reasonable tolerance given the same seed and hardware.

Common failure modes

Add a --smoke flag that runs 2 batches on CPU for CI or quick sanity checks.

Engineering problem (staff framing)

The training loop is the ML unit of work: batch→forward→loss→backward→step→eval→ckpt.

Diagram — Training loop

flowchart TD
  DL[DataLoader] --> Dev[To device] --> Fwd --> Loss --> Bwd --> Opt --> Log --> DL
  Log --> Ckpt[Best val ckpt]

Precise definitions & mental model

Epoch/step, eval mode, checkpoint contents, seeding limits.

Tradeoffs — when to use what

Hand-rolled (debuggable) vs Trainer (fast, opaque).

Failure modes (interview + on-call)

Dropout left on; labels on CPU; save weights without tokenizer.

Production & OSS practices

CPU smoke test, resume, config YAML — hand-roll once.

Deep dive (FAANG / OSS bar)

Device and dtype checklist

• Move model and batch to the same device.
• Prefer torch.autocast + GradScaler on CUDA for speed/memory.
• Call optimizer.zero_grad(set_to_none=True) for perf.
• torch.compile later — correctness first.

Eval loop hygiene

Never leave BatchNorm/Dropout in train mode during eval. For generative eval, separate "teacher-forced CE" from "sample-and-grade" metrics.

Micro-project: End-to-end MLP

In m2/mlp/ (or your course-portfolio monorepo), ship a complete mini-project:

Required files:

• train.py — full training loop with train/val split, logging, checkpointing
• model.py — MLP definition
• eval.py — load checkpoint, report val accuracy/loss
• requirements.txt — pinned torch version
• README.md — setup, run commands, expected metrics, hardware used

Required behavior:

1. Train an MLP on a real dataset (MNIST, Fashion-MNIST, or a tabular CSV from lesson 2.1).
2. Save best checkpoint by val loss to checkpoints/best.pt.
3. eval.py loads the checkpoint and prints metrics without retraining.
4. Log history.json with per-epoch losses; optional curves.png.

Milestone 2 acceptance: A hiring manager (or future you) can open the README, run two commands, and see a trained model with documented results. This folder becomes the template for the mini-LLM training script in the next module.

Optional stretch goals:

• Early stopping when val loss plateaus.
• Weights & Biases or TensorBoard logging.
• Export to ONNX for inference demo.

Checklist

• train.py runs end-to-end on CPU or GPU
• Best checkpoint saved and reloadable via eval.py
• README with reproduce instructions and final val metric
• Module README updated with what you built and what broke
• Milestone 2 training-loop section is portfolio-ready