Learning-rate schedules and regularisation

StepLR, cosine annealing, warmup, ReduceLROnPlateau, early stopping, dropout, label smoothing, weight decay and gradient clipping.

Schedules and warmup

import math
import torch
from torch.optim.lr_scheduler import (
    StepLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau, LambdaLR
)

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)

# warmup then cosine decay: the standard transformer schedule
def warmup_cosine(step, warmup=500, total=20_000, base_lr=1e-3, min_lr=1e-5):
    if step < warmup:
        return step / max(1, warmup)
    progress = (step - warmup) / max(1, total - warmup)
    factor = 0.5 * (1 + math.cos(math.pi * progress))
    return min_lr / base_lr + (1 - min_lr / base_lr) * factor

sched = LambdaLR(opt, lr_lambda=warmup_cosine)

# or step decay, still common for vision
step_sched = StepLR(opt, step_size=30, gamma=0.1)

# or one cycle, which often trains faster than a fixed rate
one_cycle = OneCycleLR(opt, max_lr=3e-3, total_steps=10_000, pct_start=0.3)

# the plateau scheduler must be stepped with the metric, not every iteration
plateau = ReduceLROnPlateau(opt, mode="min", factor=0.5, patience=3, min_lr=1e-6)

for epoch in range(10):
    train_one_epoch(model, train_loader, opt, sched)     # sched.step() per iteration
    val_loss = evaluate(model, val_loader)
    plateau.step(val_loss)                                # the metric version
    print(epoch, val_loss, opt.param_groups[0]["lr"])

Scheduler	Steps on	Use when
`StepLR`	Iterations or epochs	Vision, easy to reason about
`CosineAnnealingLR`	Iterations	Long runs; smooth and predictable
`LambdaLR`	Iterations	Custom warmup plus decay
`OneCycleLR`	Iterations	Fixed budget, fast convergence
`ReduceLROnPlateau`	Validation metric	You do not know the right schedule

⚠️

A scheduler's step count depends on how often you call it. StepLR parameters mean epochs if you call it per epoch and batches if you call it per batch. Mixing the two is the most common cause of a learning rate that reaches min_lr in the first minute.

Regularisation techniques

import torch
import torch.nn as nn

criterion = nn.CrossEntropyLoss(label_smoothing=0.1)     # softens targets

def train_one_epoch(model, loader, opt, sched, clip=1.0):
    model.train()
    total, seen = 0.0, 0
    for x, y in loader:
        x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        opt.zero_grad(set_to_none=True)
        loss = criterion(model(x), y)
        loss.backward()

        # clip the global norm: prevents a single bad batch from wrecking the run
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip)

        opt.step()
        sched.step()
        total += loss.item() * y.size(0)
        seen += y.size(0)
    return total / seen

# early stopping, tracked by hand so you keep the state of the best epoch
class EarlyStopping:
    def __init__(self, patience=5, min_delta=1e-4):
        self.patience = patience
        self.min_delta = min_delta
        self.best = float("inf")
        self.bad = 0
        self.stop = False

    def step(self, value):
        if value < self.best - self.min_delta:
            self.best = value
            self.bad = 0
        else:
            self.bad += 1
            if self.bad >= self.patience:
                self.stop = True
        return self.stop

Dropout is active only in model.train(). Forgetting to switch to model.eval() during validation changes the metric and can make a model look worse than it is.
Label smoothing prevents a model from driving logits to infinity on a noisy label and consistently helps text classification; it costs a little top-1 confidence calibration.
Gradient clipping is not a regulariser in the shrinkage sense; it prevents catastrophic steps. The global-norm version (clip_grad_norm_) is standard.
Mixup and cutmix are augmentation-level regularisers that usually beat explicit weight decay on image tasks, at the cost of longer training to converge.

Putting the pieces in the right order

# the order of operations in a step matters
def step(model, batch, opt, sched, scaler=None):
    x, y = batch
    opt.zero_grad(set_to_none=True)          # 1. clear, before backward
    with torch.autocast("cuda", enabled=scaler is not None):
        out = model(x)                       # 2. forward under autocast
        loss = criterion(out, y)

    if scaler is not None:
        scaler.scale(loss).backward()        # 3. scaled backward
        scaler.unscale_(opt)                 # 4. unscale before clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(opt)                     # 5. skip the step if inf/nan
        scaler.update()
    else:
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()

    sched.step()                             # 6. schedule after the optimiser
    return float(loss.detach())

A short checklist that catches most issues: gradients cleared before backward, backward before clipping, clipping before the optimiser step, and the schedule stepped after the step — except a plateau scheduler, which is stepped once per epoch with the validation metric.

FAQ

Do I need a scheduler if I use Adam?

It helps on any run long enough to plateau. Warmup matters especially for transformer-like models, where early large updates on poorly-initialised attention weights destabilise training.

Should I use weight decay and dropout together?

Yes, but tune them independently. Increase decay first when the training loss falls far below the validation loss; reduce dropout if the training loss stops falling before the model has fit the data.

The training loop Loss functions and optimisers in practice

Last refreshed 2026-09-18.

Learning-rate schedules and regularisation

Schedules and warmup

Regularisation techniques

Putting the pieces in the right order

FAQ

Related