753 lines
36 KiB
Python
753 lines
36 KiB
Python
#!/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())
|