The Training Loop, Optimizers & Schedulers

1. The Canonical 5-Step Training Loop

Every forward-backward pass in deep learning revolves around five fundamental operations executed in sequence:

2. Selecting Loss Functions: Avoid Numerical Instability

PyTorch incorporates numerical safeguards into its loss implementations. The most frequent error beginners make is manually applying torch.softmax before passing predictions into nn.CrossEntropyLoss.

3. Optimizers: Why AdamW is the Gold Standard

Standard Adam implements $L_2$ regularization by adding weight decay directly into the gradient calculation. As proven by Loshchilov & Hutter (2019), this causes weights with large gradients to decay less than weights with small gradients.AdamW decouples weight decay from the gradient update step, significantly improving generalization.

import torch.optim as optim

# Modern AdamW Optimizer setup
optimizer = optim.AdamW(
    model.parameters(),
    lr=3e-4,              # Andrej Karpathy's golden starting learning rate
    betas=(0.9, 0.999),   # Moving average decay rates for gradients & variance
    eps=1e-8,
    weight_decay=0.01     # Decoupled L2 regularization
)

# Traditional SGD with Momentum and Nesterov acceleration
sgd_optimizer = optim.SGD(
    model.parameters(),
    lr=0.01,
    momentum=0.9,
    nesterov=True,
    weight_decay=1e-4
)

4. Dynamic Schedulers & Gradient Clipping

A static learning rate frequently overshoots local minima or gets stuck in saddle points. Use CosineAnnealingLR for smooth decay and clip_grad_norm_to prevent catastrophic gradient explosion:

from torch.optim.lr_scheduler import CosineAnnealingLR
import torch.nn.utils as utils

# Decay learning rate according to a cosine curve down to eta_min
scheduler = CosineAnnealingLR(
    optimizer,
    T_max=100,           # Maximum number of training epochs
    eta_min=1e-6         # Minimum learning rate target
)

5. Complete Production Training & Validation Script

Here is a complete, reusable training engine incorporating validation loss calculation, accuracy tracking, gradient clipping, and scheduler stepping:

import torch
import torch.nn as nn

def train_one_epoch(model, dataloader, criterion, optimizer, device, max_grad_norm=1.0):
    model.train()  # Activate dropout and batchnorm updates
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)

        # 1. Zero out previous gradients
        optimizer.zero_grad(set_to_none=True)  # Faster than zeroing tensors

        # 2. Forward pass
        outputs = model(inputs)

        # 3. Compute loss
        loss = criterion(outputs, targets)

        # 4. Backward pass (Autograd)
        loss.backward()

        # Gradient Clipping (prevents explosion in deep architectures)
        nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)

        # 5. Step optimizer
        optimizer.step()

        # Track metrics
        running_loss += loss.item() * inputs.size(0)
        preds = torch.argmax(outputs, dim=-1)
        correct += (preds == targets).sum().item()
        total += targets.size(0)

    epoch_loss = running_loss / total
    epoch_acc = correct / total
    return epoch_loss, epoch_acc

@torch.inference_mode()
def validate(model, dataloader, criterion, device):
    model.eval()  # Freeze dropout and batchnorm statistics
    running_loss = 0.0
    correct = 0
    total = 0

    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device), targets.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, targets)

        running_loss += loss.item() * inputs.size(0)
        preds = torch.argmax(outputs, dim=-1)
        correct += (preds == targets).sum().item()
        total += targets.size(0)

    val_loss = running_loss / total
    val_acc = correct / total
    return val_loss, val_acc

# High-Performance Training Loop
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
criterion = nn.CrossEntropyLoss()

for epoch in range(1, 11):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, device)
    val_loss, val_acc = validate(model, val_loader, criterion, device)
    
    # Step scheduler after epoch
    scheduler.step()
    current_lr = optimizer.param_groups[0]['lr']

    print(f"Epoch [{epoch:02d}/10] | Train Loss: {train_loss:.4f} Acc: {train_acc*100:.1f}% | "
          f"Val Loss: {val_loss:.4f} Acc: {val_acc*100:.1f}% | LR: {current_lr:.6f}")