engine.py

"""Training loop. Mirrors main.py's runTraining (same forward, loss, metric and RNG order, so a seeded run reproduces it exactly) and adds resume, CSV/W&B logging and fault tolerance."""
import csv
import logging
import math
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

import src.losses  # noqa: F401  registers losses
import src.models  # noqa: F401  registers models
import src.optim  # noqa: F401  registers optimizers / schedulers
from src.checkpoint import load_checkpoint, rng_state, save_checkpoint, set_rng_state
from src.data import build_loader, has_val
from src.metrics import epoch_metrics, slice_counts
from src.presence import presence_loss
from src.registry import build
from src.class_weights import compute_class_weights
from utils import tqdm_

LOG = logging.getLogger("ai4mi")


def build_model(cfg: dict, device: torch.device) -> nn.Module:
    return build("model", cfg["model"]["name"], in_channels=cfg["data"]["in_channels"],
                 num_classes=cfg["data"]["num_classes"], **cfg["model"]["kwargs"]).to(device)


def supervised_loss(out: torch.Tensor | list[torch.Tensor], gt: torch.Tensor, loss_fn, dist_maps=None) -> tuple[torch.Tensor, torch.Tensor]:
    """Loss and full-resolution probabilities of a model output.

    A tensor is supervised as is. A list (deep supervision, full resolution first) is supervised per
    scale against the nearest-neighbour downsampled one-hot GT, with weights 1, 1/2, 1/4, ... except the
    lowest-resolution scale, which gets 0, normalised to sum to 1.
    """
    def apply_loss(probs, target):
        if getattr(loss_fn, "requires_distance_maps", False):
            maps = dist_maps if dist_maps is not None and dist_maps.shape == target.shape else None
            return loss_fn(probs, target, maps)
        return loss_fn(probs, target)

    if not isinstance(out, list):
        probs = F.softmax(out, dim=1)
        return apply_loss(probs, gt), probs
    weights = [0.5**i for i in range(len(out) - 1)]
    total = sum(weights)
    loss = 0
    for logits, weight in zip(out, weights):
        target = F.interpolate(gt.float(), size=logits.shape[-2:], mode="nearest").to(gt.dtype)
        loss = loss + weight / total * apply_loss(F.softmax(logits, dim=1), target)
    return loss, F.softmax(out[0], dim=1)


def run_epoch(split: str, net: nn.Module, loader, loss_fn, optimizer, device, cfg: dict,
              desc: str) -> tuple[float, np.ndarray, list[str]]:
    training = optimizer is not None
    net.train(training)
    losses, counts, stems, failed = [], [], [], 0
    with torch.set_grad_enabled(training):
        for batch in tqdm_(loader, desc=desc, disable=not sys.stdout.isatty()):
            try:
                img, gt = batch["images"].to(device), batch["gts"].to(device)
                if training:
                    optimizer.zero_grad()
                dist_maps = batch.get("dist_maps")
                if dist_maps is not None:
                    dist_maps = dist_maps.to(device)
                loss, probs = supervised_loss(net(img), gt, loss_fn, dist_maps)
                if cfg["train"]["presence_weight"]:
                    loss = loss + cfg["train"]["presence_weight"] * presence_loss(net.presence_logits, gt)
                if hasattr(loss_fn, "combine") and "dist_maps" in batch:  # boundary_dice_ce: full-resolution output only, after the deep-supervision sum
                    loss = loss_fn.combine(loss, probs, batch["dist_maps"].to(device))
                if training:
                    loss.backward()
                    optimizer.step()
            except RuntimeError as exc:  # CUDA OOM, a malformed slice, ...: skip the batch, bounded
                failed += 1
                LOG.exception("%s batch failed (%d/%d allowed), stems %s", split, failed,
                              cfg["train"]["max_failed_batches"], batch["stems"][:4])
                if training:
                    optimizer.zero_grad(set_to_none=True)
                if failed > cfg["train"]["max_failed_batches"]:
                    raise RuntimeError(f"too many failed {split} batches") from exc
                continue
            losses.append(loss.item())
            counts.append(slice_counts(probs.argmax(dim=1), gt))
            stems.extend(batch["stems"])
    if not counts:
        raise RuntimeError(f"no {split} batches completed: {len(loader.dataset)} slices, {failed} failed")
    return float(np.mean(losses)), np.concatenate(counts), stems


def write_row(path: Path, row: dict) -> None:
    new = not path.exists()
    with path.open("a", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(row))
        if new:
            writer.writeheader()
        writer.writerow(row)


def truncate_rows(path: Path, before_epoch: int) -> None:
    """Drop CSV rows from epochs that will be re-run after a resume."""
    if not path.exists():
        return
    with path.open(newline="") as f:
        rows = [r for r in csv.DictReader(f) if int(r["epoch"]) < before_epoch]
    path.unlink()
    for r in rows:
        write_row(path, r)


def overlay_samples(net: nn.Module, dataset, device, n: int = 4):
    """Fixed, evenly spaced val slices so W&B shows the same examples every epoch."""
    idx = np.linspace(0, len(dataset) - 1, n).astype(int)
    items = [dataset[i] for i in idx]
    net.eval()
    with torch.no_grad():
        img = torch.stack([it["images"] for it in items]).to(device)
        pred = net(img).argmax(dim=1).cpu().numpy()
    gts = torch.stack([it["gts"] for it in items]).argmax(dim=1).numpy()
    return img[:, img.shape[1] // 2].cpu().numpy(), gts, pred, [it["stems"] for it in items]


def should_stop(epoch: int, best_epoch: int, patience: int) -> bool:
    """True once `patience` epochs have passed since the best one; patience 0 never stops."""
    return patience > 0 and epoch - best_epoch >= patience


def fit(cfg: dict, run_dir: Path, device: torch.device, resume: bool, wb) -> int:
    """Train to cfg.train.epochs, resuming from checkpoints/last.pt if asked. Returns the best epoch."""
    net = build_model(cfg, device)
    optimizer = build("optim", cfg["optim"]["name"], params=net.parameters(), **cfg["optim"]["kwargs"])
    scheduler = build("scheduler", cfg["scheduler"]["name"], optimizer=optimizer,
                      epochs=cfg["train"]["epochs"], **cfg["scheduler"]["kwargs"])
    final = not has_val(cfg)  # retains 0: all labelled patients train, no val split; the last epoch is the model
    splits = ("train",) if final else ("train", "val")
    loaders = {split: build_loader(cfg, split, device) for split in splits}
    loss_kwargs = dict(cfg["loss"]["kwargs"])
    if cfg["loss"]["name"] in ("weighted_ce", "weighted_ce_dice", "weighted_ce_dice_boundary"):
        loss_kwargs["weights"] = compute_class_weights(
            loaders["train"].dataset,
            num_classes=cfg["data"]["num_classes"],
        )
        LOG.info("Class weights: %s", loss_kwargs["weights"])
    if cfg["loss"]["name"] == "boundary_dice_ce":
        loss_kwargs["epochs"] = cfg["train"]["epochs"]  # the alpha schedule spans the whole run
    if cfg["loss"]["name"] == "boundary_dice_ce_scheduled":
        loss_kwargs["epochs"] = cfg["train"]["epochs"]  # the weight schedule runs over the configured epochs
    loss_fn = build("loss", cfg["loss"]["name"], num_classes=cfg["data"]["num_classes"],
                    **loss_kwargs)
    LOG.info("model %s: %.2fM parameters | train %d slices, val %s slices", cfg["model"]["name"],
             sum(p.numel() for p in net.parameters()) / 1e6,
             len(loaders["train"].dataset), "no" if final else len(loaders["val"].dataset))

    ckpt_dir, csv_path = run_dir / "checkpoints", run_dir / "epochs.csv"
    select = cfg["train"]["select_metric"]
    start, best, best_epoch = 0, -math.inf, -1
    if resume:
        state = load_checkpoint(ckpt_dir / "last.pt")
        net.load_state_dict(state["model"])
        optimizer.load_state_dict(state["optimizer"])
        if scheduler:
            scheduler.load_state_dict(state["scheduler"])
        set_rng_state(state["rng"])
        start, best, best_epoch = state["epoch"] + 1, state["best"], state["best_epoch"]
        truncate_rows(csv_path, start)
        LOG.info("resumed after epoch %d (best %s=%.4f at epoch %d)", state["epoch"], select, best, best_epoch)
    elif csv_path.exists():
        csv_path.unlink()  # leftovers from a run that died before its first checkpoint

    for epoch in range(start, cfg["train"]["epochs"]):
        t0 = time.time()
        if hasattr(loss_fn, "set_epoch"):
            loss_fn.set_epoch(epoch)
        lr = optimizer.param_groups[0]["lr"]
        row = {"epoch": epoch}
        for split, opt in ((s, optimizer if s == "train" else None) for s in splits):
            loss, counts, stems = run_epoch(split, net, loaders[split], loss_fn, opt, device, cfg,
                                            f">> {split:5s} ({epoch:4d})")
            row[f"{split}_loss"] = loss
            row |= epoch_metrics(split, counts, stems, cfg)
            if split == "val":
                val_counts, val_stems = counts, stems
        if scheduler:
            scheduler.step()

        improved = final or best_epoch < 0 or math.isnan(best) or row[select] > best  # final: every epoch overwrites best.pt
        if improved:
            best, best_epoch = row[select], epoch
            save_checkpoint(ckpt_dir / "best.pt", {"epoch": epoch, "model": net.state_dict(), "config": cfg})
            torch.save(net, ckpt_dir / "best_model.pkl")  # whole pickled net, as the submission asks
            if not final:
                np.savez_compressed(run_dir / "val_counts_best.npz", counts=val_counts, stems=np.array(val_stems))
        row |= {"lr": lr, "seconds": round(time.time() - t0, 1), "is_best": improved}
        write_row(csv_path, row)
        save_checkpoint(ckpt_dir / "last.pt", {
            "epoch": epoch, "model": net.state_dict(), "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict() if scheduler else None,
            "best": best, "best_epoch": best_epoch, "rng": rng_state()})

        wb.log(row, step=epoch)
        every = cfg["wandb"]["image_every"]
        if not final and wb.run and every and (epoch % every == 0 or epoch == cfg["train"]["epochs"] - 1):
            wb.log_overlays(*overlay_samples(net, loaders["val"].dataset, device),
                            cfg["data"]["class_names"], step=epoch)
        if final:
            LOG.info("epoch %3d | train loss %.4f | train dice fg %.4f | %.0fs",
                     epoch, row["train_loss"], row["train_dice_fg"], row["seconds"])
        else:
            LOG.info("epoch %3d | train loss %.4f | val loss %.4f | val dice fg %.4f (legacy %.4f)%s | %.0fs",
                     epoch, row["train_loss"], row["val_loss"], row["val_dice_fg"],
                     row["val_dice_legacy_fg"], "  *best*" if improved else "", row["seconds"])
        if should_stop(epoch, best_epoch, cfg["train"]["early_stop_patience"]):
            LOG.info("early stop: no %s gain for %d epochs (best epoch %d)", select,
                     cfg["train"]["early_stop_patience"], best_epoch)
            break

    return best_epoch