noaa-goes-visualization/train.py

#!/usr/bin/env python
"""Train the learned filler.  Runs on the training host, inside the ROCm container.

    train.py --shards /data/shards --manifest /data/shards/manifest.json --out /data/runs/v4

The loss lives in **asinh space**, where the shards already are: it is bounded,
smooth, and supervises every pixel of the bands the sample lost -- the ones the
renderer would fill; the target's other bands are inputs, not answers.  The previous trainer optimised the display
mapping instead, which clamps at a per-band `vmin` sitting *above* the median pixel
-- most of the frame had exactly zero gradient, and the annulus straddling the
threshold carried a 25x spike with a hard discontinuity (measured: |dL/dpred| of 0 /
49 / 2-7 below, at, and above the threshold).  Three runs diverged on that loss.
A display-space term is kept, softened so its worst slope is ~10, because the bench
scores display PSNR and the loss should feel the same per-band weighting; reported
PSNR uses the *exact* hard bench mapping so the numbers are comparable to the
baseline table (optical_flow 44.79 dB at gap 1, 34.47 at gap 300).

Validation runs on a fixed, evenly spaced subset of held-out days **every 250
steps**, not per epoch: the last collapse happened entirely inside one 41-minute
epoch and was invisible until it ended.  Training aborts automatically if validation
PSNR falls well below its running best twice in a row.
"""

import argparse
import json
import os
import random
import sys
import time

import numpy as np
import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from suvi import (align, availability as availability_module, dataset, metrics,
                  model, pool as pool_module, samples, segments as segments_module)

#: Targets a pool worker draws from one slice before moving to another: enough to
#: amortise opening it, few enough that a batch mixes several slices.
POOL_VISIT = 16
#: Seconds a pool worker waits for a first slice before giving up.
POOL_WAIT = 1800.0

#: Weight of the gradient-matching term (asinh space): penalises blur, not just error.
GRADIENT_WEIGHT = 0.5
#: Weight of the softened display-space term.
DISPLAY_WEIGHT = 0.25
#: Charbonnier knee; small enough not to blunt real errors.
CHARBONNIER_EPS = 1e-3
#: Softening of the display gamma for the *loss*.  At 5e-3 the worst-case slope of
#: the mapping is ~10 against the 37,500 of the 1e-8 the old trainer used; the
#: mapping still sends 0 to 0 exactly.  The reported PSNR never uses this -- metrics
#: use the hard bench mapping.
DISPLAY_EPS = 5e-3
#: Steps between validation passes.
EVAL_EVERY = 250
#: Abort when validation PSNR sits this far (dB) below its best, twice in a row.
ABORT_DROP_DB = 1.5
#: Stepped plateau decay (ReduceLROnPlateau).  A pre-committed anneal (OneCycle)
#: starved the data that arrived late in the pool rotation, and a one-shot
#: cool-down freezes the model on the val set's say-so alone.  Instead: hold LR
#: while validation improves, and each time it makes no new best for
#: PLATEAU_EVALS evals, halve it.  Training thus runs at the highest LR still
#: producing progress; reaching the frozen threshold takes six consecutive
#: full plateau windows, so the corpus is seen several more times on the way
#: down.  The run ends when the scale sits at its floor and still nothing
#: improves.
PLATEAU_EVALS = 10
DECAY_FACTOR = 0.5
MIN_LR_SCALE = 1 / 64
WARMUP_STEPS = 2000


def learning_rate(step, base, warmup=WARMUP_STEPS, scale=1.0):
    """The schedule as a pure function: warmup ramp times the plateau scale."""
    return base * scale * min(1.0, (step + 1) / max(warmup, 1))


def sweep_rate(step, low, high, total):
    """The LR range test's exponential ramp: `low` at step 0, `high` at the end.

    One short run under this ramp traces loss against LR; the useful band shows
    as the steepest sustained descent, and the divergence point marks too-high.
    The standard pick is a factor ~3 below the divergence point.
    """
    return low * (high / low) ** (step / max(total - 1, 1))


def display_map(coded, eps=DISPLAY_EPS):
    """Model space -> display space for all six bands, differentiably when eps > 0.

    The torch equivalent of :func:`suvi.metrics.to_display`; with ``eps=0.0`` it *is*
    that mapping, for metrics under no_grad.
    """
    radiance = torch.sinh(coded.clamp(-1, 1) * dataset.ASINH_RANGE) * dataset.ASINH_SCALE
    out = []
    for index, wavelength in enumerate(samples.paths.WAVELENGTHS):
        vmin, vmax, gamma = metrics.DISPLAY_MAPPING[wavelength]
        scaled = ((radiance[:, index] - vmin) / vmax).clamp(0.0, 1.0)
        out.append((scaled + eps) ** gamma - eps**gamma)
    return torch.stack(out, dim=1)


def charbonnier(prediction, truth, bands, eps=CHARBONNIER_EPS):
    """Mean Charbonnier error over the (sample, band)s `bands` (B, 6) marks."""
    error = torch.sqrt((prediction - truth) ** 2 + eps**2).mean(dim=(-2, -1))
    return (error * bands).sum() / bands.sum().clamp(min=1.0)


def gradient_loss(prediction, truth, bands):
    """Match image gradients, so a blurred answer is penalised as well as a wrong one."""
    loss = 0.0
    for axis in (-1, -2):
        loss = loss + charbonnier(prediction.diff(dim=axis), truth.diff(dim=axis), bands)
    return loss / 2


def criterion(prediction, target, lost):
    """The full training loss over the lost bands, in spaces whose gradients are
    bounded everywhere."""
    loss = charbonnier(prediction, target, lost)
    loss = loss + GRADIENT_WEIGHT * gradient_loss(prediction, target, lost)
    loss = loss + DISPLAY_WEIGHT * charbonnier(
        display_map(prediction), display_map(target), lost
    )
    return loss


def host_free_gb():
    """Host memory available to new allocations, in GiB, or -1 if unknown.

    MemAvailable rather than MemFree: page cache is reclaimable, so MemFree
    reads alarmingly low on a healthy box and says nothing about whether the
    next batch will fit.
    """
    try:
        with open("/proc/meminfo") as handle:
            for line in handle:
                if line.startswith("MemAvailable:"):
                    return round(int(line.split()[1]) / 1024**2, 1)
    except OSError:
        pass
    return -1.0


def display_psnr(prediction, target, lost):
    """Mean per-band display PSNR over the lost bands, the bench's definition, in dB.

    Per band then averaged -- PSNR of the pooled MSE reads ~4 dB lower (Jensen) and
    would not be comparable to the published filler table.
    """
    with torch.no_grad():
        shown = display_map(prediction.float(), eps=0.0)
        truth = display_map(target.float(), eps=0.0)
        mse = ((shown - truth) ** 2).mean(dim=(-2, -1)).clamp(min=1e-12)
        psnr = 10.0 * torch.log10(1.0 / mse)
        return float((psnr * lost).sum() / lost.sum().clamp(min=1.0))


# ----------------------------------------------------------------------- data feed


def load_availability(shards_dir, override=None):
    """The archive's availability model, required for training.

    A hard requirement, not a fallback: without it the sampler silently reverts to
    synthetic-only outages, and the whole point of the profiled distribution is
    that the synthetic one under-represented dual-satellite failures.
    """
    path = override or os.path.join(shards_dir, "availability.npz")
    if not os.path.exists(path):
        raise SystemExit(
            f"{path} not found.  Training requires the profiled availability model "
            "(run profile_archive.py); refusing to fall back to synthetic outages."
        )
    return availability_module.Availability.load(path)


class ShardSet(torch.utils.data.Dataset):
    """Samples drawn from the shards of one split.

    Shards hold clean frames; damage is drawn fresh here on every access, so two
    epochs never see the same corruption of the same slot.  Validation pins its seed
    to the slot index instead, because a validation number that moves because the
    *damage* moved says nothing about the model.
    """

    def __init__(self, directory, days, satellites=(16, 18), deterministic=False,
                 availability=None, seed=0, offsets=samples.OFFSETS, segments=None):
        self.deterministic = deterministic
        self.seed = seed
        shards = {}
        for day in days:
            for satellite in satellites:
                path = os.path.join(directory, dataset.shard_name(day, satellite))
                if os.path.exists(path):
                    shards[(day, satellite)] = dataset.Shard(path)
        if not shards:
            raise SystemExit(f"no shards for {len(days)} days under {directory}")
        self.sampler = samples.Sampler(shards, satellites=satellites,
                                       availability=availability, seed=seed,
                                       offsets=offsets, segments=segments)
        self.items = self.sampler.targets()
        if segments is not None and not self.items:
            raise SystemExit("no targets: the resident shards do not cover the "
                             "split's segments (with the span margin)")

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        # Iterative fallback, not recursive: a vetted day can hold *hours* of
        # contiguous unreadable-sentinel slots (present in the archive, dropped by
        # vetting -- the completeness filter counts presence, not cleanliness), and
        # a thousand consecutive absent targets once blew the stack mid-run.
        for offset in range(len(self.items)):
            position = (index + offset) % len(self.items)
            satellite, when = self.items[position]
            rng = (np.random.default_rng(self.seed * 1_000_003 + position)
                   if self.deterministic else self.sampler.rng)
            sample = self.sampler.build(satellite, when, rng)
            if sample is not None:
                return samples.to_tensors(sample, torch)
        raise RuntimeError("no slot in the entire set yields a buildable sample")


class PoolSet(torch.utils.data.IterableDataset):
    """Endless samples from a rotating pool of slices (:mod:`suvi.pool`).

    Each worker visits a random resident slice, draws `per_visit` targets from it
    with the ordinary sampler -- damage fresh on every draw -- and moves on, so the
    model sees whatever the producer has made lately and a retired slice is never
    revisited.  Its shards are opened eagerly, so a slice retired mid-visit stays
    readable through the handles already held.
    """

    def __init__(self, directory, satellites=(16, 18), availability=None, seed=0,
                 offsets=samples.OFFSETS, per_visit=POOL_VISIT, wait=POOL_WAIT):
        self.directory = directory
        self.satellites = tuple(satellites)
        self.availability = availability
        self.seed = seed
        self.offsets = offsets
        self.per_visit = per_visit
        self.wait = wait

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        worker = info.id if info is not None else 0
        rng = np.random.default_rng([self.seed, worker, time.time_ns()])
        waited = 0.0
        while True:
            ready = pool_module.ready_slices(self.directory)
            if not ready:
                if waited >= self.wait:
                    raise RuntimeError(f"no slice in {self.directory} after "
                                       f"{self.wait:.0f} s: is the producer running?")
                time.sleep(5.0)
                waited += 5.0
                continue
            waited = 0.0
            path = ready[int(rng.integers(len(ready)))][0]
            try:
                shards, segments = pool_module.open_slice(path)
            except (OSError, ValueError, KeyError):
                continue                  # retired while we were choosing it
            try:
                sampler = samples.Sampler(shards, satellites=self.satellites,
                                          availability=self.availability,
                                          seed=int(rng.integers(2**31)),
                                          offsets=self.offsets, segments=segments)
                targets = sampler.targets()
                for _ in range(self.per_visit if targets else 0):
                    satellite, when = targets[int(rng.integers(len(targets)))]
                    sample = sampler.build(satellite, when, rng)
                    if sample is not None:
                        yield samples.to_tensors(sample, torch)
            finally:
                for shard in shards.values():
                    shard.close()


class SliceSet(torch.utils.data.Dataset):
    """A fixed set of targets from slice directories -- the validation set written
    by ``pool_producer.py --fixed`` -- with damage pinned to each target's index, so
    a number that moves says something about the model."""

    def __init__(self, directory, satellites=(16, 18), availability=None, seed=1,
                 offsets=samples.OFFSETS):
        self.seed = seed
        self.samplers = []
        self.items = []
        for path in sorted(path for path, _ in pool_module.ready_slices(directory)):
            shards, segments = pool_module.open_slice(path)
            sampler = samples.Sampler(shards, satellites=satellites,
                                      availability=availability, seed=seed,
                                      offsets=offsets, segments=segments)
            self.items.extend((len(self.samplers), satellite, when)
                              for satellite, when in sampler.targets())
            self.samplers.append(sampler)
        if not self.items:
            raise SystemExit(f"no targets in the slices under {directory}")

    def __len__(self):
        return len(self.items)

    def __getitem__(self, index):
        for offset in range(len(self.items)):
            position = (index + offset) % len(self.items)
            which, satellite, when = self.items[position]
            rng = np.random.default_rng(self.seed * 1_000_003 + position)
            sample = self.samplers[which].build(satellite, when, rng)
            if sample is not None:
                return samples.to_tensors(sample, torch)
        raise RuntimeError("no target in the slices yields a buildable sample")


def collate(batch):
    """Pad stacks to a common depth.

    Depth varies because the anchors are the nearest usable frames, per band, and
    near the edge of a shard there may be one or none.  Padding with `missing`
    frames that hold no band is exactly right -- the model gives them zero weight --
    and keeps the batch rectangular.  Padded rows carry the identity transfer so
    alignment ignores them.
    """
    depth = max(item["stack"].shape[0] for item in batch)
    out = {}
    for item in batch:
        short = depth - item["stack"].shape[0]
        if short:
            item = dict(item)
            item["stack"] = torch.cat(
                [item["stack"], torch.zeros(short, *item["stack"].shape[1:])]
            )
            filler = model.frame_conditioning("missing", True, 0.0)
            item["condition"] = torch.cat([item["condition"], filler.repeat(short, 1)])
            item["dts"] = torch.cat([item["dts"], torch.zeros(short)])
            item["gains"] = torch.cat([item["gains"], torch.ones(short, model.BANDS)])
            item["offsets"] = torch.cat([item["offsets"],
                                         torch.zeros(short, model.BANDS)])
            item["present"] = torch.cat([item["present"],
                                         torch.zeros(short, model.BANDS)])
        for key, value in item.items():
            out.setdefault(key, []).append(value)
    return {key: torch.stack(value) for key, value in out.items()}


def aligned_batch(batch, device):
    """Move a batch to the device and run the deterministic alignment, in fp32.

    Alignment is physics, not learning: it runs outside autocast because a bf16
    ``sinh`` at the top of the asinh range keeps only two decimal digits of a
    radiance the gain is about to multiply.
    """
    stack = batch["stack"].to(device, non_blocking=True).float()
    condition = batch["condition"].to(device, non_blocking=True)
    present = batch["present"].to(device, non_blocking=True).float()
    target = batch["target"].to(device, non_blocking=True).float()
    lost = batch["lost"].to(device, non_blocking=True).float()
    valid = condition[..., model.AVAILABLE]
    aligned = align.align_stack(
        stack, batch["dts"].to(device, non_blocking=True), valid,
        batch["gains"].to(device, non_blocking=True),
        batch["offsets"].to(device, non_blocking=True),
        batch["b0"].to(device, non_blocking=True),
        batch["radius"].to(device, non_blocking=True),
    )
    return aligned, condition, present, target, lost


def evaluate(net, loader, device):
    """Loss and bench-comparable PSNR over a whole loader."""
    net.eval()
    totals = {"loss": 0.0, "psnr": 0.0, "n": 0}
    with torch.no_grad():
        for batch in loader:
            aligned, condition, present, target, lost = aligned_batch(batch, device)
            with torch.autocast(device.type, dtype=torch.bfloat16,
                                enabled=device.type == "cuda"):
                prediction = net(aligned, condition, present)
            prediction = prediction.float()
            count = aligned.shape[0]
            totals["loss"] += float(criterion(prediction, target, lost)) * count
            totals["psnr"] += display_psnr(prediction, target, lost) * count
            totals["n"] += count
    net.train()
    n = max(totals["n"], 1)
    return {"loss": totals["loss"] / n, "psnr": totals["psnr"] / n}


# ------------------------------------------------------------------------ training


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--shards", required=True)
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--base", type=int, default=32)
    parser.add_argument("--depth", type=int, default=3)
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--steps-per-epoch", type=int, default=None)
    parser.add_argument("--val-samples", type=int, default=64)
    parser.add_argument("--eval-every", type=int, default=EVAL_EVERY)
    parser.add_argument("--plateau-evals", type=int, default=PLATEAU_EVALS,
                        help="halve the LR after this many evals without a new "
                             "best (0 disables the stepped decay)")
    parser.add_argument("--lr-range", default=None,
                        help="'LO,HI': run the LR range test instead of training "
                             "-- LR ramps exponentially from LO to HI over "
                             "--max-steps while train loss is recorded; no evals")
    parser.add_argument("--max-steps", type=int, default=None,
                        help="stop after this many optimiser steps (gates and probes)")
    parser.add_argument("--overfit", type=int, default=None,
                        help="train and validate on this many frozen held-in samples "
                             "-- the overfit gate; loss must approach zero")
    parser.add_argument("--pool", default=None,
                        help="train from this rotating pool directory (pool_producer.py) "
                             "instead of the corpus's train split; validation and test "
                             "stay on --shards.  Needs --steps-per-epoch.")
    parser.add_argument("--val-slices", default=None,
                        help="validate on the fixed slices in this directory "
                             "(pool_producer.py --fixed) instead of the corpus's val split")
    parser.add_argument("--evaluate", default=None,
                        help="score this checkpoint on the validation set and exit")
    parser.add_argument("--availability", default=None,
                        help="availability.npz (default: beside --shards)")
    parser.add_argument("--warmup", type=int, default=WARMUP_STEPS,
                        help="LR warmup length; short screening runs need a "
                             "short warmup or they compare warmups, not models")
    parser.add_argument("--relative", action="store_true",
                        help="variant: deviation-from-median input channels")
    parser.add_argument("--corr-flow", action="store_true",
                        help="variant: cost-volume input to the flow head")
    parser.add_argument("--attention", action="store_true",
                        help="variant: cross-frame attention on a token grid")
    parser.add_argument("--span", type=int, default=None,
                        help="model input span in slots (a rung of "
                             f"{samples.LADDER}); sets the offset ladder and, "
                             "with a segment manifest, the target margin")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--resume", default=None)
    args = parser.parse_args(argv)

    torch.manual_seed(args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.backends.cudnn.benchmark = True        # worth ~2x on this GPU

    manifest = json.load(open(args.manifest))
    satellites = tuple(manifest.get("satellites", (16, 18)))
    offsets = samples.ladder_for(args.span) if args.span else samples.OFFSETS
    seg_splits = None
    if "segments" in manifest:
        # Pristine-segment manifest: split separation is by contiguous segment,
        # and a calendar day may legitimately host segments from two splits, so
        # the day-level leak check does not apply.  Each split opens only the
        # shards its segments touch, which also bounds open file descriptors.
        seg_splits = {name: [tuple(pair) for pair in chosen]
                      for name, chosen in manifest["segments"].items()}
        splits = {name: [d for d in segments_module.days_touched(chosen)
                         if os.path.exists(os.path.join(
                             args.shards, dataset.shard_name(d, satellites[0])))]
                  for name, chosen in seg_splits.items()}
        leaks = []
        if args.span is None:
            # The corpus's own minimum is the natural default; a larger span
            # simply restricts targets to the segments long enough for it.
            offsets = samples.ladder_for(manifest["span_min"])
    else:
        raise SystemExit(
            f"{args.manifest} has no 'segments': it predates the pristine-segment "
            "corpus.  Rebuild it with dataset_build.py.")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device: {device}"
          + (f" ({torch.cuda.get_device_properties(0).gcnArchName})"
             if device.type == "cuda" else ""))
    print(f"manifest {manifest['digest']}: "
          + ", ".join(f"{k} {len(v)}d" for k, v in sorted(splits.items())))
    print(f"  span +/-{max(offsets)} slots, ladder {offsets}"
          + (f", segment splits "
             + "/".join(str(len(seg_splits[k])) for k in ("train", "val", "test"))
             if seg_splits else ""))
    if leaks:
        print(f"  {len(leaks)} split boundaries on consecutive days: {leaks[:3]}")

    trace = load_availability(args.shards, args.availability)
    print(f"  availability model: {trace.slots:,} slots, satellites {trace.satellites}")

    if args.overfit:
        # The overfit gate: a handful of *frozen* samples (deterministic losses), and
        # validation on those same samples.  A model that cannot drive this loss
        # toward zero has a wiring or conditioning problem that a long run will not
        # cure -- two of the three failed runs would have failed this in minutes.
        held_in = ShardSet(args.shards, splits["train"], satellites,
                           deterministic=True, availability=trace, seed=args.seed,
                           offsets=offsets,
                           segments=seg_splits and seg_splits["train"])
        stride = max(len(held_in) // args.overfit, 1)
        train_set = torch.utils.data.Subset(
            held_in, [i * stride for i in range(args.overfit)]
        )
        val_set = train_set
    else:
        if args.evaluate:
            train_set = None
        elif args.pool:
            if not args.steps_per_epoch:
                raise SystemExit("--pool needs --steps-per-epoch: a pool has no length")
            train_set = PoolSet(args.pool, satellites, availability=trace,
                                seed=args.seed, offsets=offsets)
        else:
            train_set = ShardSet(args.shards, splits["train"], satellites,
                                 availability=trace, seed=args.seed, offsets=offsets,
                                 segments=seg_splits and seg_splits["train"])
        if args.val_slices:
            val_set = SliceSet(args.val_slices, satellites, availability=trace, seed=1,
                               offsets=offsets)
        else:
            val_set = ShardSet(args.shards, splits["val"], satellites,
                               deterministic=True, availability=trace, seed=1,
                               offsets=offsets, segments=seg_splits and seg_splits["val"])
        # An evenly spaced subset, so a small validation pass still spans every val
        # day and both satellites, and is the same subset every time.
        if args.val_samples and args.val_samples < len(val_set):
            stride = len(val_set) / args.val_samples
            val_set = torch.utils.data.Subset(
                val_set, [int(i * stride) for i in range(args.val_samples)]
            )
    if train_set is None:
        described = "none (evaluating)"
    else:
        described = f"pool {args.pool}" if args.pool else f"{len(train_set)} targets"
    print(f"  train {described}, val {len(val_set)}"
          + (f" from {args.val_slices}" if args.val_slices else ""))

    loader_args = dict(batch_size=args.batch, num_workers=args.workers,
                       collate_fn=collate, pin_memory=device.type == "cuda",
                       persistent_workers=args.workers > 0)
    train_loader = None if train_set is None else torch.utils.data.DataLoader(
        train_set, shuffle=not args.pool, drop_last=True, **loader_args)
    # The validation loader gets its own settings rather than the train loader's.
    # Sharing them made it persistent too, so a second full prefetch queue stayed
    # resident for the whole run -- one stack is batch * frames * 6 * 640^2 * 4 B,
    # which at +/-256 is ~0.94 GB, and workers * prefetch of those is ~11 GB held
    # between evaluations that happen once every few hundred steps.  Together with
    # the train queue that was ~23 GB of host RAM, and the +/-256 run exhausted
    # the box's 124 GB and spent its last 101 minutes swap-thrashing.  Validation
    # is a short deterministic pass; paying worker startup each time is far
    # cheaper than keeping its queue alive.
    val_loader = torch.utils.data.DataLoader(
        val_set, shuffle=False, batch_size=args.batch, collate_fn=collate,
        pin_memory=device.type == "cuda",
        num_workers=min(args.workers, 2), persistent_workers=False,
    )

    net = model.build(base=args.base, depth=args.depth, relative=args.relative,
                      corr_flow=args.corr_flow, attention=args.attention).to(device)
    net = net.to(memory_format=torch.channels_last)
    print(f"  {sum(p.numel() for p in net.parameters()) / 1e6:.1f}M parameters")
    if args.evaluate:
        state = torch.load(args.evaluate, map_location=device, weights_only=False)
        net.load_state_dict(state["model"])
        val = evaluate(net, val_loader, device)
        print(f"evaluated {args.evaluate} (step {state.get('step')}): val "
              f"{val['loss']:.5f} ({val['psnr']:.2f} dB) over {len(val_set)} targets",
              flush=True)
        return 0

    optimiser = torch.optim.AdamW(net.parameters(), lr=args.lr, weight_decay=1e-4)
    steps_per_epoch = args.steps_per_epoch or len(train_loader)
    total_steps = args.epochs * steps_per_epoch

    sweep = None
    if args.lr_range:
        if not args.max_steps:
            raise SystemExit("--lr-range needs --max-steps (the ramp's length)")
        sweep = tuple(float(v) for v in args.lr_range.split(","))

    step = 0
    start_epoch = 0
    lr_scale = 1.0
    if args.resume and os.path.exists(args.resume):
        state = torch.load(args.resume, map_location=device)
        net.load_state_dict(state["model"])
        optimiser.load_state_dict(state["optimiser"])
        start_epoch = state["epoch"] + 1
        step = state.get("step", start_epoch * steps_per_epoch)
        lr_scale = state.get("lr_scale", 1.0)
        print(f"  resumed from {args.resume} at epoch {start_epoch}, step {step}")

    os.makedirs(args.out, exist_ok=True)
    history_path = os.path.join(args.out, "history.jsonl")

    def checkpoint(name, epoch):
        state = {"model": net.state_dict(), "optimiser": optimiser.state_dict(),
                 "epoch": epoch, "step": step, "lr_scale": lr_scale,
                 # The conditioning layout rides along so load_learned can refuse
                 # a checkpoint trained under another one.
                 "args": vars(args) | {"cond_layout": list(model.COND_LAYOUT),
                                       "offsets": list(offsets)},
                 "manifest": manifest["digest"]}
        torch.save(state, os.path.join(args.out, name))

    def record(row):
        with open(history_path, "a") as handle:
            handle.write(json.dumps(row) + "\n")

    # The untrained model is the prior-weighted blend of aligned candidates -- a
    # meaningful policy, and the floor every later eval is judged against.
    best = float("-inf")
    if sweep is None:
        baseline = evaluate(net, val_loader, device)
        print(f"step {step:>6}  val {baseline['loss']:.5f} "
              f"({baseline['psnr']:.2f} dB)  [untrained]", flush=True)
        record({"step": step, "val_loss": baseline["loss"],
                "val_psnr": baseline["psnr"], "untrained": True})
        best = baseline["psnr"]
    if not os.path.exists(os.path.join(args.out, "best.pt")):
        checkpoint("best.pt", epoch=-1)     # never clobber a resumed run's best
    drops = 0
    since_best = 0

    window = {"loss": 0.0, "grad": 0.0, "grad_max": 0.0, "n": 0}
    started = time.time()
    for epoch in range(start_epoch, args.epochs):
        epoch_steps = 0
        for batch in train_loader:
            # --steps caps the *loop*, not just the LR schedule.  The first run
            # honoured it only in the schedule, so an "epoch" silently became the
            # whole 4,312-batch split and last.pt was never written before the
            # machine stalled at step ~1,700 -- leaving nothing to resume from.
            if epoch_steps >= steps_per_epoch:
                break
            if step >= total_steps or (args.max_steps and step >= args.max_steps):
                break
            epoch_steps += 1
            aligned, condition, present, target, lost = aligned_batch(batch, device)
            with torch.autocast(device.type, dtype=torch.bfloat16,
                                enabled=device.type == "cuda"):
                prediction = net(aligned, condition, present)
            loss = criterion(prediction.float(), target, lost)

            optimiser.zero_grad(set_to_none=True)
            loss.backward()
            grad = float(torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0))
            lr = (sweep_rate(step, sweep[0], sweep[1], args.max_steps)
                  if sweep else learning_rate(step, args.lr, warmup=args.warmup,
                                              scale=lr_scale))
            for group in optimiser.param_groups:
                group["lr"] = lr
            optimiser.step()
            step += 1

            window["loss"] += float(loss.detach())
            window["grad"] += grad
            window["grad_max"] = max(window["grad_max"], grad)
            window["n"] += 1

            if sweep is not None:
                # The range test only traces loss against LR; evals, checkpoints
                # and the plateau logic are all beside its point.
                if step % 20 == 0:
                    n = max(window["n"], 1)
                    record({"step": step, "lr": lr,
                            "train_loss": window["loss"] / n})
                    print(f"step {step:>6}  lr {lr:.3e}  "
                          f"train {window['loss'] / n:.5f}", flush=True)
                    window = {"loss": 0.0, "grad": 0.0, "grad_max": 0.0, "n": 0}
                continue

            if step % args.eval_every == 0:
                val = evaluate(net, val_loader, device)
                n = max(window["n"], 1)
                row = {
                    "step": step, "epoch": epoch,
                    # Both, so progress can be read in epochs (passes through the
                    # corpus) as well as steps without knowing the corpus size.
                    "steps_per_epoch": steps_per_epoch,
                    "total_steps": min(total_steps, args.max_steps or total_steps),
                    "train_loss": window["loss"] / n,
                    "val_loss": val["loss"], "val_psnr": val["psnr"],
                    "grad_mean": window["grad"] / n, "grad_max": window["grad_max"],
                    "lr": optimiser.param_groups[0]["lr"],
                    "seconds": time.time() - started,
                    # Recorded every eval because the failure it catches is
                    # invisible otherwise: a run that exhausts host memory keeps
                    # its process in R state at high CPU while it swap-thrashes,
                    # so it looks like training and simply stops producing rows.
                    "host_free_gb": host_free_gb(),
                }
                record(row)
                print(f"step {step:>6}  train {row['train_loss']:.5f}  "
                      f"val {val['loss']:.5f} ({val['psnr']:.2f} dB)  "
                      f"grad {row['grad_mean']:.3f}/{row['grad_max']:.2f}  "
                      f"lr {row['lr']:.2e}  {row['seconds'] / 60:.1f} min", flush=True)
                window = {"loss": 0.0, "grad": 0.0, "grad_max": 0.0, "n": 0}
                started = time.time()
                if device.type == "cuda":
                    # The composition holds multi-GB full-resolution tensors, and on
                    # unified GTT memory the caching allocator's fragmentation growth
                    # is host RAM disappearing: the first full run crept to 75 GB
                    # over ~1,700 steps and stalled the machine into swap.  Releasing
                    # the cache every eval costs milliseconds and caps the creep.
                    torch.cuda.empty_cache()

                # A resumable point every eval, not just every epoch: the stall cost
                # 1,700 steps because the only last.pt cadence was the epoch end.
                checkpoint("last.pt", epoch)
                if val["psnr"] > best:
                    best = val["psnr"]
                    checkpoint("best.pt", epoch)
                    drops = 0
                    since_best = 0
                elif val["psnr"] < best - ABORT_DROP_DB:
                    drops += 1
                    since_best += 1
                    if drops >= 2:
                        print(f"ABORT: validation {val['psnr']:.2f} dB has sat "
                              f">{ABORT_DROP_DB} dB below the best ({best:.2f} dB) "
                              "on two consecutive evals.", flush=True)
                        checkpoint("last.pt", epoch)
                        return 1
                else:
                    drops = 0
                    since_best += 1
                if args.plateau_evals and since_best >= args.plateau_evals:
                    if lr_scale <= MIN_LR_SCALE:
                        print(f"FROZEN: lr scale at its floor ({lr_scale:.4f}) "
                              f"and still no new best in {since_best} evals; "
                              f"best {best:.2f} dB stands.  Stopping.", flush=True)
                        checkpoint("last.pt", epoch)
                        return 0
                    lr_scale *= DECAY_FACTOR
                    since_best = 0
                    print(f"PLATEAU: no new best in {args.plateau_evals} evals; "
                          f"halving lr to {args.lr * lr_scale:.2e}.", flush=True)
        checkpoint("last.pt", epoch)
        if step >= total_steps or (args.max_steps and step >= args.max_steps):
            break
    print(f"best validation display PSNR {best:.2f} dB")
    return 0


if __name__ == "__main__":
    sys.exit(main())