Training loops, loss functions & optimizers

The training loop is where learning actually happens. A well-written loop is reusable across projects, logs the right numbers and makes experiments easy to compare. This lesson covers the three decisions that most influence training: the loss function, the optimizer and the learning-rate schedule.

Choosing a loss function

Task	Output layer	Loss	Label format
Regression (e.g. price in INR)	1 linear unit	nn.MSELoss or nn.L1Loss	float tensor
Robust regression with outliers	1 linear unit	nn.HuberLoss	float tensor
Binary (fraud / not fraud)	1 logit	nn.BCEWithLogitsLoss	float 0.0 / 1.0
Multi-class (10 product types)	K logits	nn.CrossEntropyLoss	integer class index
Multi-label (tags on a saree photo)	K logits	nn.BCEWithLogitsLoss	K floats of 0/1

Two practical tips:

• For imbalanced data (1 fraud in 500 transactions), pass pos_weight to BCEWithLogitsLoss or weight to CrossEntropyLoss so rare classes count more.
• MSE punishes large errors heavily; L1 (mean absolute error) is easier to explain to business users ("we are off by Rs 230 on average").

Optimizers

Optimizer	Idea	When to use
SGD	Step against the gradient	Simple baselines
SGD + momentum	Keeps a running velocity, smoothing noisy steps	Strong for CNNs with a good schedule
Adam	Per-parameter adaptive step sizes using running mean and variance of gradients	Fast, forgiving default
AdamW	Adam with decoupled weight decay	Default for transformers and most modern work

Start with AdamW at lr=1e-3 (or 3e-4 for bigger models). If loss explodes or becomes NaN, lower the learning rate by 10x.

Learning-rate schedules

A fixed learning rate is rarely optimal. Common schedules:

• StepLR – multiply the rate by 0.1 every N epochs.
• CosineAnnealingLR – smoothly decays to near zero.
• OneCycleLR – warms up then anneals; excellent for short training runs.
• ReduceLROnPlateau – lowers the rate when validation loss stops improving.

A reusable training loop

import torch
import torch.nn as nn

def fit(model, train_dl, val_dl, epochs=10, lr=1e-3, device="cpu"):
    model.to(device)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-2)
    sched = torch.optim.lr_scheduler.OneCycleLR(
        opt, max_lr=lr, epochs=epochs, steps_per_epoch=len(train_dl)
    )
    history = []
    for epoch in range(epochs):
        model.train()
        train_loss = 0.0
        for xb, yb in train_dl:
            xb, yb = xb.to(device), yb.to(device)
            opt.zero_grad(set_to_none=True)
            loss = loss_fn(model(xb), yb)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            opt.step()
            sched.step()                      # OneCycle steps every batch
            train_loss += loss.item() * xb.size(0)

        model.eval()
        val_loss, correct = 0.0, 0
        with torch.no_grad():
            for xb, yb in val_dl:
                xb, yb = xb.to(device), yb.to(device)
                logits = model(xb)
                val_loss += loss_fn(logits, yb).item() * xb.size(0)
                correct += (logits.argmax(1) == yb).sum().item()

        n_train, n_val = len(train_dl.dataset), len(val_dl.dataset)
        row = {"epoch": epoch + 1,
               "train_loss": train_loss / n_train,
               "val_loss": val_loss / n_val,
               "val_acc": correct / n_val}
        history.append(row)
        print(row)
    return history

Details worth noticing:

• zero_grad(set_to_none=True) is slightly faster than filling with zeros.
• Gradient clipping caps the size of updates and prevents occasional exploding gradients, especially in RNNs.
• Schedulers such as OneCycleLR step per batch; StepLR and CosineAnnealingLR are usually stepped once per epoch.

Mixed precision for speed

On modern NVIDIA GPUs, automatic mixed precision runs many operations in 16-bit floats, often giving 1.5–3x speed-ups:

scaler = torch.amp.GradScaler("cuda")
for xb, yb in train_dl:
    xb, yb = xb.to("cuda"), yb.to("cuda")
    opt.zero_grad()
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(xb), yb)
    scaler.scale(loss).backward()
    scaler.step(opt)
    scaler.update()

Reading the loss curves

Pattern	Likely cause	Fix
Loss flat from the start	LR too low, bug in labels	Raise LR, check a batch by eye
Loss becomes NaN	LR too high, bad inputs	Lower LR, clip gradients, check for inf
Train falls, val rises	Overfitting	Regularise, augment, stop early
Both high and flat	Underfitting	Bigger model, train longer

A great sanity check: train on a single batch for 100 steps. A correct pipeline should drive its loss close to zero. If it cannot, there is a bug.

Key takeaways

• Match the loss to the task and label format; use logits with the "WithLogits" and cross-entropy losses.
• AdamW plus a schedule is a strong default.
• Clip gradients, log train and validation metrics every epoch, and overfit one batch first.