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