Gradually increase LR from near-zero for the first N steps. Stabilizes training for transformers and large batch sizes. Implemented in OneCycleLR or custom schedulers.

Yes for small models and debugging. GPU is essential for reasonable training time on images or large data.

Training Loops

Q: What loss function for classification?

CrossEntropyLoss for multi-class (includes softmax). BCEWithLogitsLoss for multi-label binary. Pass raw logits, not softmax outputs.

Q: What is weight decay?

L2 regularization applied to weights in AdamW. weight_decay=0.01 is a common starting point. Reduces overfitting on small datasets.

Q: What learning rate scheduler should I use?

CosineAnnealingLR - smooth decay, popular default. OneCycleLR - fast convergence with warmup. ReduceLROnPlateau - reduce when val metric stalls.

Q: Should I shuffle every epoch?

Yes for training (shuffle=True in DataLoader). Shuffling prevents the model from learning batch order artifacts.

Q: How do I log training metrics?

Print per epoch for simple scripts. Use TensorBoard (torch.utils.tensorboard), W&B, or MLflow for production.

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.
AdamW optimizer 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_mode without 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_size and 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?

CrossEntropyLoss for multi-class (includes softmax).
BCEWithLogitsLoss for 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.01 is 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"] + 1

What 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=True in 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 OneCycleLR or 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.

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+.

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.
AdamW optimizer 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_mode without 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_size and 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?

CrossEntropyLoss for multi-class (includes softmax).
BCEWithLogitsLoss for 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.01 is 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"] + 1

What 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=True in 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 OneCycleLR or 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.

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+.

Training Loops

Search Documentation

Training Loops