Training Loops
A training loop repeats forward pass, loss computation, backward pass, and optimizer step across batches and epochs. Schedulers adjust the learning rate; checkpoints preserve progress for resume and deployment.
Recipe
Quick-reference recipe card - copy-paste ready.
for epoch in range(num_epochs):
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
optimizer.zero_grad(set_to_none=True)
loss = loss_fn(model(x), y)
loss.backward()
optimizer.step()
scheduler.step()When to reach for this:
- Every PyTorch model needs a training loop (unless using Lightning).
- Implementing validation passes, early stopping, and checkpointing.
- Tuning learning rate schedules and gradient clipping.
- Logging metrics per epoch or per step.
Working Example
"""training_loops.py - complete train/val loop with checkpointing."""
from __future__ import annotations
import torch
import torch.nn as nn
from torch.optim.lr_scheduler import CosineAnnealingLR
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
dataset = datasets.FashionMNIST(root="./data", train=True, download=True, transform=transform)
train_ds, val_ds = random_split(dataset, [50_000, 10_000])
train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=2, pin_memory=True)
val_loader = DataLoader(val_ds, batch_size=256, shuffle=False, num_workers=2)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=10)
loss_fn = nn.CrossEntropyLoss()
best_val_loss = float("inf")
for epoch in range(10):
model.train()
train_loss = 0.0
for x, y in train_loader:
x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
optimizer.zero_grad(set_to_none=True)
loss = loss_fn(model(x), y)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
train_loss += loss.item() * x.size(0)
scheduler.step()
model.eval()
val_loss = 0.0
correct = 0
with torch.inference_mode():
for x, y in val_loader:
x, y = x.to(device), y.to(device)
logits = model(x)
val_loss += loss_fn(logits, y).item() * x.size(0)
correct += (logits.argmax(1) == y).sum().item()
val_loss /= len(val_ds)
acc = correct / len(val_ds)
print(f"epoch {epoch+1}: val_loss={val_loss:.4f} acc={acc:.3f} lr={scheduler.get_last_lr()[0]:.6f}")
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save({"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict()}, "best.pt")What this demonstrates:
- Separate train and validation phases with correct mode switching.
AdamWoptimizer with weight decay and cosine LR schedule.- Gradient clipping for training stability.
- Checkpoint saving when validation loss improves.
Deep Dive
How It Works
- Forward pass computes predictions and loss.
- Backward pass computes gradients via autograd.
- Optimizer step updates parameters using gradients and learning rate.
- Scheduler modifies LR based on epoch or step count.
- Validation runs in
inference_modewithout gradient tracking.
Optimizer Selection
| Optimizer | Use When | Notes |
|---|---|---|
| AdamW | Default for most models | Decoupled weight decay |
| SGD + momentum | CNNs with careful tuning | Often best with LR schedule |
| Adam | Quick prototyping | Weight decay interacts with adaptive LR |
| Lion | Research/experimentation | Memory-efficient alternative |
Python Notes
# Gradient accumulation for large effective batch size
accum_steps = 4
for i, (x, y) in enumerate(train_loader):
loss = loss_fn(model(x), y) / accum_steps
loss.backward()
if (i + 1) % accum_steps == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)Gotchas
- Not calling model.eval() for validation - dropout randomizes and batch norm uses batch stats. Fix:
model.eval()before validation;model.train()before training. - Logging loss without averaging - raw batch losses are noisy. Fix: accumulate
loss.item() * batch_sizeand divide by dataset size. - Stepping scheduler every batch vs epoch - wrong schedule shape. Fix: match scheduler type (
step()per epoch for CosineAnnealingLR). - No checkpoint on best validation - lose the best model if training continues and overfits. Fix: save on validation improvement.
- Learning rate too high - loss NaNs or diverges. Fix: start with 1e-3 for Adam, 0.1 for SGD; use warmup.
- Forgetting to move data to device - CPU/GPU mismatch errors. Fix:
x.to(device, non_blocking=True)in the loop.
Alternatives
| Alternative | Use When | Don't Use When |
|---|---|---|
| Manual loop | Learning, full control | Repetitive boilerplate across projects |
| PyTorch Lightning | Production training structure | Debugging autograd fundamentals |
| HuggingFace Trainer | Transformer fine-tuning | Custom CNN architectures |
torch.compile | Speed up the forward/backward | Debugging training instability |
FAQs
What loss function for classification?
CrossEntropyLossfor multi-class (includes softmax).BCEWithLogitsLossfor multi-label binary.- Pass raw logits, not softmax outputs.
How many epochs should I train?
- Until validation metric plateaus (early stopping).
- 10-100 for small datasets; monitor val loss every epoch.
What is weight decay?
- L2 regularization applied to weights in AdamW.
weight_decay=0.01is a common starting point.- Reduces overfitting on small datasets.
When should I use gradient accumulation?
- When GPU memory limits batch size but you need a larger effective batch.
- Divide loss by accumulation steps before backward.
- Step optimizer every N batches.
How do I resume from a checkpoint?
ckpt = torch.load("best.pt", weights_only=True)
model.load_state_dict(ckpt["model"])
optimizer.load_state_dict(ckpt["optimizer"])
start_epoch = ckpt["epoch"] + 1What learning rate scheduler should I use?
CosineAnnealingLR- smooth decay, popular default.OneCycleLR- fast convergence with warmup.ReduceLROnPlateau- reduce when val metric stalls.
How do I detect overfitting?
- Training loss decreases while validation loss increases.
- Add dropout, weight decay, data augmentation, or early stopping.
Should I shuffle every epoch?
- Yes for training (
shuffle=Truein DataLoader). - Shuffling prevents the model from learning batch order artifacts.
What is early stopping?
- Stop training when validation metric does not improve for N epochs.
- Restore weights from the best checkpoint.
How do I log training metrics?
- Print per epoch for simple scripts.
- Use TensorBoard (
torch.utils.tensorboard), W&B, or MLflow for production.
What is warmup?
- Gradually increase LR from near-zero for the first N steps.
- Stabilizes training for transformers and large batch sizes.
- Implemented in
OneCycleLRor custom schedulers.
Can I train on CPU?
- Yes for small models and debugging.
- GPU is essential for reasonable training time on images or large data.
Related
- Autograd - backward pass mechanics
- Datasets & DataLoaders - batch input
- GPUs & Mixed Precision - faster training
- PyTorch Lightning - structured loops
- Saving, Loading & Export - checkpoint format
Stack versions: This page was written for Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, and uv 0.6+.