Putting It Together — A Production-Grade Loop

Everything in one place — the loop you'll copy and adapt

This lesson assembles every pattern from the track into one runnable training loop. Read it once carefully; you'll be cribbing from variants of this for the rest of your career.

What's in the recipe:

Device selection (CUDA → MPS → CPU fallback)
AdamW with weight decay split (no decay on bias / norm)
Cosine LR with warmup
Mixed-precision (bfloat16 on Ampere+ / Apple Silicon)
Gradient clipping
tqdm progress bar with running loss
Per-epoch validation
Best-checkpoint saving
NaN protection

What's intentionally NOT in here (each gets its own dedicated track or lesson):

Distributed training (DDP / FSDP) — Track 7
torch.compile() — Track 7
Experiment tracking (W&B / MLflow) — Track 8

Code

The whole loop, end to end·python

import math
import torch
import torch.nn as nn
from torch.amp import autocast
from tqdm import tqdm

# 1. Device
device = (
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)

# 2. Model
model = MyModel().to(device)

# 3. Loss
criterion = nn.CrossEntropyLoss()

# 4. Optimizer with weight-decay split
decay, no_decay = [], []
for n, p in model.named_parameters():
    if not p.requires_grad: continue
    if p.dim() < 2 or any(k in n for k in ('bias', 'norm')):
        no_decay.append(p)
    else:
        decay.append(p)
optimizer = torch.optim.AdamW(
    [{'params': decay, 'weight_decay': 0.01},
     {'params': no_decay, 'weight_decay': 0.0}],
    lr=1e-4,
)

# 5. Scheduler — warmup + cosine
total_steps = len(train_loader) * num_epochs
warmup_steps = total_steps // 20    # 5% warmup
def lr_lambda(step):
    if step < warmup_steps:
        return step / warmup_steps
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return 0.5 * (1.0 + math.cos(math.pi * progress))
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

...continued: train + validate + checkpoint·python

best_val = float('inf')
global_step = 0

for epoch in range(num_epochs):
    # ---- TRAIN ----
    model.train()
    pbar = tqdm(train_loader, desc=f"epoch {epoch:02d}")
    running = 0.0

    for batch_x, batch_y in pbar:
        batch_x = batch_x.to(device, non_blocking=True)
        batch_y = batch_y.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)

        with autocast(device_type=device, dtype=torch.bfloat16) if device != 'cpu' else nullcontext():
            output = model(batch_x)
            loss = criterion(output, batch_y)

        if not torch.isfinite(loss):
            print(f"step {global_step}: non-finite loss; skipping")
            continue

        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()
        global_step += 1

        running = 0.95 * running + 0.05 * loss.item() if global_step > 1 else loss.item()
        pbar.set_postfix(loss=f"{running:.4f}", lr=f"{scheduler.get_last_lr()[0]:.2e}")

    # ---- VALIDATE ----
    model.eval()
    val_loss = 0.0
    val_n = 0
    with torch.inference_mode():
        for vx, vy in val_loader:
            vx, vy = vx.to(device), vy.to(device)
            val_loss += criterion(model(vx), vy).item() * vx.size(0)
            val_n += vx.size(0)
    val_loss /= val_n
    print(f"epoch {epoch:02d}  val_loss={val_loss:.4f}")

    # ---- CHECKPOINT ----
    if val_loss < best_val:
        best_val = val_loss
        torch.save({
            'epoch': epoch,
            'val_loss': val_loss,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
        }, 'best.pt')

nullcontext shim — autocast is no-op on CPU·python

from contextlib import nullcontext

# autocast doesn't apply on CPU. Use nullcontext as a no-op stand-in
# so the same training code runs cleanly on CPU for debugging.

device = "cpu"  # for example
amp_ctx = autocast(device_type=device, dtype=torch.bfloat16) if device != 'cpu' else nullcontext()

with amp_ctx:
    out = model(x)
    loss = criterion(out, y)

Exercise

Adapt this loop to your favorite small dataset (MNIST / CIFAR-10 / FashionMNIST). Run for 5 epochs. Confirm: (1) val loss decreases over epochs, (2) the best checkpoint loads cleanly into a fresh model, (3) bfloat16 autocast actually engages on your hardware (test by removing it and timing — should be slower).

Putting It Together — A Production-Grade Loop

Everything in one place — the loop you'll copy and adapt

Code

External links

Exercise

Progress

Comments 0