Ch 2 — Training Deep Networks¶

optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

l1_loss = sum(p.abs().sum() for p in model.parameters())
total_loss = ce_loss + lambda_l1 * l1_loss

best_val_loss = float("inf")
patience_counter = 0
PATIENCE = 10

for epoch in range(max_epochs):
    train(...)
    val_loss = validate(...)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        patience_counter = 0
        torch.save(model.state_dict(), "best_model.pt")
    else:
        patience_counter += 1
        if patience_counter >= PATIENCE:
            print("Early stopping triggered.")
            break

"""
Production-quality PyTorch training loop.

Features:
- DataLoader with prefetching
- AdamW optimiser with weight decay
- Cosine annealing scheduler with linear warmup
- Gradient clipping
- Mixed precision training (torch.amp)
- Gradient accumulation
- Model checkpointing with early stopping
- Validation loop
"""

from __future__ import annotations

import math
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.cuda.amp import GradScaler, autocast
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader


# ---------------------------------------------------------------------------
# Scheduler: linear warmup + cosine decay
# ---------------------------------------------------------------------------

def get_cosine_schedule_with_warmup(
    optimiser: torch.optim.Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    eta_min_ratio: float = 0.0,
) -> LambdaLR:
    """Cosine decay schedule with linear warmup.

    Args:
        optimiser:           The optimiser to schedule.
        num_warmup_steps:    Number of linear warmup steps.
        num_training_steps:  Total number of training steps.
        eta_min_ratio:       Minimum LR as a fraction of peak LR.

    Returns:
        A LambdaLR scheduler.
    """

    def lr_lambda(current_step: int) -> float:
        # Warmup phase
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        # Cosine decay
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))
        return eta_min_ratio + (1.0 - eta_min_ratio) * cosine_decay

    return LambdaLR(optimiser, lr_lambda)


# ---------------------------------------------------------------------------
# Trainer
# ---------------------------------------------------------------------------

class Trainer:
    """Generic trainer for classification models.

    Args:
        model:              PyTorch model (nn.Module).
        train_loader:       DataLoader for training data.
        val_loader:         DataLoader for validation data.
        lr:                 Peak learning rate.
        weight_decay:       AdamW weight decay coefficient.
        max_epochs:         Maximum number of training epochs.
        grad_clip:          Max gradient norm for clipping (0 = disabled).
        accumulation_steps: Number of micro-batches to accumulate.
        patience:           Early stopping patience (epochs).
        checkpoint_path:    Path to save the best model checkpoint.
        use_amp:            Whether to use automatic mixed precision.
        device:             Device to train on.
    """

    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        lr: float = 3e-4,
        weight_decay: float = 1e-2,
        max_epochs: int = 100,
        grad_clip: float = 1.0,
        accumulation_steps: int = 1,
        patience: int = 10,
        checkpoint_path: str | Path = "best_model.pt",
        use_amp: bool = True,
        device: torch.device | None = None,
    ) -> None:
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.max_epochs = max_epochs
        self.grad_clip = grad_clip
        self.accumulation_steps = accumulation_steps
        self.patience = patience
        self.checkpoint_path = Path(checkpoint_path)
        self.device = device or torch.device(
            "cuda" if torch.cuda.is_available() else "cpu"
        )
        self.use_amp = use_amp and self.device.type == "cuda"

        self.model.to(self.device)

        # Optimiser: separate weight-decay and no-decay param groups
        decay_params = [
            p for n, p in model.named_parameters()
            if p.requires_grad and p.dim() >= 2
        ]
        no_decay_params = [
            p for n, p in model.named_parameters()
            if p.requires_grad and p.dim() < 2
        ]
        self.optimiser = torch.optim.AdamW(
            [
                {"params": decay_params, "weight_decay": weight_decay},
                {"params": no_decay_params, "weight_decay": 0.0},
            ],
            lr=lr,
        )

        # Scheduler: warmup for 5% of total steps
        total_steps = len(train_loader) * max_epochs // accumulation_steps
        warmup_steps = max(1, total_steps // 20)
        self.scheduler = get_cosine_schedule_with_warmup(
            self.optimiser, warmup_steps, total_steps
        )

        # Mixed precision
        self.scaler = GradScaler(enabled=self.use_amp)

        # State
        self.best_val_loss = float("inf")
        self._patience_counter = 0

    # ------------------------------------------------------------------
    # One training epoch
    # ------------------------------------------------------------------

    def _train_epoch(self) -> float:
        """Run one training epoch. Returns mean loss over all batches."""
        self.model.train()
        total_loss = 0.0
        self.optimiser.zero_grad()

        for step, (x_batch, y_batch) in enumerate(self.train_loader):
            x_batch = x_batch.to(self.device, non_blocking=True)
            y_batch = y_batch.to(self.device, non_blocking=True)

            with autocast(enabled=self.use_amp):
                logits = self.model(x_batch)
                loss = F.cross_entropy(logits, y_batch)
                # Scale by accumulation steps so effective LR doesn't change
                loss = loss / self.accumulation_steps

            self.scaler.scale(loss).backward()

            if (step + 1) % self.accumulation_steps == 0:
                # Unscale before gradient clipping
                self.scaler.unscale_(self.optimiser)

                if self.grad_clip > 0:
                    nn.utils.clip_grad_norm_(
                        self.model.parameters(), self.grad_clip
                    )

                self.scaler.step(self.optimiser)
                self.scaler.update()
                self.scheduler.step()
                self.optimiser.zero_grad()

            total_loss += loss.item() * self.accumulation_steps

        return total_loss / len(self.train_loader)

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    @torch.no_grad()
    def _validate(self) -> tuple[float, float]:
        """Evaluate on validation set. Returns (loss, accuracy)."""
        self.model.eval()
        total_loss, correct, total = 0.0, 0, 0

        for x_batch, y_batch in self.val_loader:
            x_batch = x_batch.to(self.device, non_blocking=True)
            y_batch = y_batch.to(self.device, non_blocking=True)

            with autocast(enabled=self.use_amp):
                logits = self.model(x_batch)
                loss = F.cross_entropy(logits, y_batch, reduction="sum")

            total_loss += loss.item()
            correct += (logits.argmax(dim=1) == y_batch).sum().item()
            total += y_batch.size(0)

        return total_loss / total, correct / total

    # ------------------------------------------------------------------
    # Checkpointing
    # ------------------------------------------------------------------

    def _save_checkpoint(self, epoch: int, val_loss: float) -> None:
        """Save model state dict and training metadata."""
        checkpoint: dict[str, Any] = {
            "epoch": epoch,
            "model_state": self.model.state_dict(),
            "optimiser_state": self.optimiser.state_dict(),
            "scheduler_state": self.scheduler.state_dict(),
            "val_loss": val_loss,
        }
        torch.save(checkpoint, self.checkpoint_path)

    def load_best(self) -> None:
        """Restore best checkpoint into the model."""
        ckpt = torch.load(self.checkpoint_path, map_location=self.device)
        self.model.load_state_dict(ckpt["model_state"])
        print(f"Loaded best model from epoch {ckpt['epoch']} "
              f"(val_loss={ckpt['val_loss']:.4f})")

    # ------------------------------------------------------------------
    # Main training loop
    # ------------------------------------------------------------------

    def fit(self) -> dict[str, list[float]]:
        """Train the model and return loss/accuracy history.

        Returns:
            Dictionary with keys 'train_loss', 'val_loss', 'val_acc'.
        """
        history: dict[str, list[float]] = {
            "train_loss": [], "val_loss": [], "val_acc": []
        }

        for epoch in range(1, self.max_epochs + 1):
            train_loss = self._train_epoch()
            val_loss, val_acc = self._validate()

            history["train_loss"].append(train_loss)
            history["val_loss"].append(val_loss)
            history["val_acc"].append(val_acc)

            lr_now = self.optimiser.param_groups[0]["lr"]
            print(
                f"Epoch {epoch:4d}/{self.max_epochs}  "
                f"train_loss={train_loss:.4f}  "
                f"val_loss={val_loss:.4f}  "
                f"val_acc={val_acc:.3f}  "
                f"lr={lr_now:.2e}"
            )

            # Checkpointing
            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                self._patience_counter = 0
                self._save_checkpoint(epoch, val_loss)
                print("  --> Saved new best checkpoint.")
            else:
                self._patience_counter += 1
                if self._patience_counter >= self.patience:
                    print(f"Early stopping at epoch {epoch}.")
                    break

        self.load_best()
        return history