#!/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())