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.  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, dataset, metrics, model, samples

#: 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


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, eps=CHARBONNIER_EPS):
    return torch.sqrt((prediction - truth) ** 2 + eps**2).mean()


def gradient_loss(prediction, truth):
    """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))
    return loss / 2


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


def display_psnr(prediction, target):
    """Mean per-band display PSNR, 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)
        return float((10.0 * torch.log10(1.0 / mse)).mean())


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


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,
                 seed=0):
        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, seed=seed)
        self.items = self.sampler.targets()

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

    def __getitem__(self, index):
        satellite, when = self.items[index]
        rng = (np.random.default_rng(self.seed * 1_000_003 + index)
               if self.deterministic else self.sampler.rng)
        sample = self.sampler.build(satellite, when, rng)
        if sample is None:                       # target absent; fall back to a neighbour
            return self[(index + 1) % len(self.items)]
        return samples.to_tensors(sample, torch)


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

    Depth varies because the anchors are the nearest usable frames, and near the
    edge of a shard there may be one or none.  Padding with `missing` frames 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)])
        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)
    target = batch["target"].to(device, non_blocking=True).float()
    valid = (condition[..., 0] + condition[..., 2]).clamp(0, 1)
    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, target


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, target = aligned_batch(batch, device)
            with torch.autocast(device.type, dtype=torch.bfloat16,
                                enabled=device.type == "cuda"):
                prediction = net(aligned, condition)
            prediction = prediction.float()
            count = aligned.shape[0]
            totals["loss"] += float(criterion(prediction, target)) * count
            totals["psnr"] += display_psnr(prediction, target) * 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("--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("--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))
    splits = manifest["splits"]
    leaks = dataset.check_split(splits)          # raises if a day is in two splits
    satellites = tuple(manifest.get("satellites", (16, 18)))

    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())))
    if leaks:
        print(f"  {len(leaks)} split boundaries on consecutive days: {leaks[:3]}")

    if args.overfit:
        # The overfit gate: a handful of *frozen* samples (deterministic damage), 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, seed=args.seed)
        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:
        train_set = ShardSet(args.shards, splits["train"], satellites, seed=args.seed)
        val_set = ShardSet(args.shards, splits["val"], satellites, deterministic=True,
                           seed=1)
        # 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)]
            )
    print(f"  train {len(train_set)} targets, val {len(val_set)}")

    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 = torch.utils.data.DataLoader(train_set, shuffle=True, drop_last=True,
                                               **loader_args)
    val_loader = torch.utils.data.DataLoader(val_set, shuffle=False, **loader_args)

    net = model.build(base=args.base, depth=args.depth).to(device)
    net = net.to(memory_format=torch.channels_last)
    print(f"  {sum(p.numel() for p in net.parameters()) / 1e6:.1f}M parameters")

    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
    schedule = torch.optim.lr_scheduler.OneCycleLR(
        optimiser, max_lr=args.lr, total_steps=total_steps, pct_start=0.2
    )

    step = 0
    start_epoch = 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)
        for _ in range(min(step, total_steps - 1)):
            schedule.step()
        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, "args": vars(args),
                 "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.
    baseline = evaluate(net, val_loader, device)
    print(f"step {step:>6}  val {baseline['loss']:.5f} ({baseline['psnr']:.2f} dB)  "
          f"[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

    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, target = aligned_batch(batch, device)
            with torch.autocast(device.type, dtype=torch.bfloat16,
                                enabled=device.type == "cuda"):
                prediction = net(aligned, condition)
            loss = criterion(prediction.float(), target)

            optimiser.zero_grad(set_to_none=True)
            loss.backward()
            grad = float(torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0))
            optimiser.step()
            schedule.step()
            step += 1

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

            if step % args.eval_every == 0:
                val = evaluate(net, val_loader, device)
                n = max(window["n"], 1)
                row = {
                    "step": step, "epoch": epoch,
                    "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,
                }
                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
                elif val["psnr"] < best - ABORT_DROP_DB:
                    drops += 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
        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())