"""WiSER radiomap V2.2a training entry point: V1.5 arch + frozen D2.2 backbone. Inherits V1.5 hierarchical radiomap head. Adds: * --load-d22-ckpt PATH: load D2.2 best.pt weights into backbone + tx_proj (strict=False, CIR head keys discarded). * --freeze-base: freeze backbone+tx_embed+tx_proj; only radiomap_head trains. The forward uses `torch.no_grad()` around the base when frozen to save activation memory, and the `train()` override keeps the frozen part in `eval()` mode regardless of outer train/eval switches. Usage: python -m torch.distributed.run --standalone --nproc_per_node=8 \ scripts/v2p2a/train_v2p2a.py \ --manifest benchmarks/formal100_sionna_radiomap_100train_10heldoutval.json \ --dataset-kind sionna_radiomap --db-floor -300.0 \ --gradient-weight 1.0 --corridor-budget 192 \ --num-cross-layers 4 --num-self-layers 2 \ --load-d22-ckpt /.../D2p2_.../best.pt --freeze-base \ --run-root --epochs 50 --warmup-epochs 5 --batch-size 4 """ from __future__ import annotations import argparse import csv import json import math import os import platform import sys import time from contextlib import nullcontext from pathlib import Path from typing import Any import numpy as np import torch import torch.distributed as dist import torch.nn.functional as F from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler ROOT = Path(__file__).resolve().parent.parent # WiSER root sys.path.insert(0, str(ROOT)) from wiser.config import ModelConfig as SharedModelConfig from wiser.data.radiomap_dataset import ( RadiomapDataset, radiomap_collate, ) from wiser.alt_models import JointRadiomapCIRModel, FreezeConfig from wiser.alt_engine.checkpoint import ( detect_csi_head_arch_from_ckpt, load_warm_start_ckpt, load_phase_ckpt_into_model, save_phase_ckpt, ) def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--manifest", required=True) p.add_argument("--run-root", required=True, help="absolute output dir for this run") p.add_argument("--epochs", type=int, default=50) p.add_argument("--warmup-epochs", type=int, default=5) p.add_argument("--batch-size", type=int, default=16) p.add_argument("--val-batch-size", type=int, default=8) p.add_argument("--val-every-epochs", type=int, default=2) p.add_argument("--num-workers", type=int, default=4) p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--lr-min", type=float, default=1e-6) p.add_argument("--weight-decay", type=float, default=0.0) p.add_argument("--channels", type=int, default=512) p.add_argument("--downsample-stages", type=int, default=3) p.add_argument("--blocks-per-stage", type=int, default=3) p.add_argument("--spatial-channels", type=int, default=256) p.add_argument("--stem-depth", type=int, default=4) p.add_argument("--output-depth", type=int, default=4) p.add_argument("--grid-h", type=int, default=36) p.add_argument("--grid-w", type=int, default=36) p.add_argument("--db-floor", type=float, default=-200.0) p.add_argument("--dataset-kind", default="csi_path_aggregate", choices=["csi_path_aggregate", "sionna_radiomap"]) p.add_argument("--huber-beta", type=float, default=2.0) p.add_argument("--seed", type=int, default=0) p.add_argument("--backbone-kind", default="trellis2", choices=["trellis2", "dense_fallback"]) # V1.4 addition p.add_argument("--gradient-weight", type=float, default=1.0, help="weight of WiSER spatial decoder-style gradient L1 loss (0 disables, same as V1.3)") # V1.5 additions p.add_argument("--corridor-budget", type=int, default=192, help="RayCorridorGather top-K tokens per query from level 0 (10cm)") p.add_argument("--num-cross-layers", type=int, default=4, help="decoder cross-attn depth (WiSER spatial decoder default 4)") p.add_argument("--num-self-layers", type=int, default=2, help="decoder self-attn depth (WiSER spatial decoder default 2)") # V2.2a additions (kept for backward-compat) p.add_argument("--load-d22-ckpt", default=None, help="path to D2.2 best.pt; loads backbone+tx_proj (ignore CIR head)") p.add_argument("--freeze-base", action="store_true", help="(legacy) freeze backbone+tx_embed+tx_proj; only radiomap_head trains") # WiSER additions p.add_argument("--phase-name", default="P0_warmup_backbone", help="phase name, used for logging / ckpt metadata") p.add_argument("--load-v22a-ckpt", default=None, help="V2.2a best.pt path; provides radiomap head + backbone + tx_proj") p.add_argument("--load-phase-ckpt", default=None, help="Load full prior-phase ckpt (supersedes --load-v22a/--load-d22)") p.add_argument("--freeze-backbone", action="store_true", help="Set backbone.requires_grad=False") p.add_argument("--freeze-tx-proj", action="store_true") p.add_argument("--freeze-radiomap-head", action="store_true") p.add_argument("--freeze-csi-head", action="store_true", help="Set csi_head.requires_grad=False (default for radiomap task!)") return p.parse_args() def gradient_l1_loss(pred: torch.Tensor, gt: torch.Tensor, mask_hw: torch.Tensor) -> torch.Tensor: """WiSER spatial decoder style: L1 of finite-difference gradient diff, masked to pairs of adjacent cells that are BOTH valid. pred/gt: [B, H, W]. mask_hw: bool [B, H, W].""" dx_pred = pred[:, :, 1:] - pred[:, :, :-1] dx_gt = gt[:, :, 1:] - gt[:, :, :-1] dx_mask = mask_hw[:, :, 1:] & mask_hw[:, :, :-1] dy_pred = pred[:, 1:, :] - pred[:, :-1, :] dy_gt = gt[:, 1:, :] - gt[:, :-1, :] dy_mask = mask_hw[:, 1:, :] & mask_hw[:, :-1, :] out = pred.new_zeros(()) if dx_mask.any(): out = out + (dx_pred - dx_gt).abs()[dx_mask].mean() if dy_mask.any(): out = out + (dy_pred - dy_gt).abs()[dy_mask].mean() return out def setup_ddp() -> tuple[int, int, int, torch.device]: dist.init_process_group(backend="nccl") rank = dist.get_rank() world_size = dist.get_world_size() local_rank = int(os.environ.get("LOCAL_RANK", rank)) torch.cuda.set_device(local_rank) device = torch.device(f"cuda:{local_rank}") return rank, local_rank, world_size, device def cleanup_ddp() -> None: if dist.is_initialized(): dist.destroy_process_group() def compute_lr(step: int, total_steps: int, warmup_steps: int, lr_max: float, lr_min: float) -> float: if step < warmup_steps: return lr_max * (step + 1) / max(warmup_steps, 1) progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) progress = min(max(progress, 0.0), 1.0) cos = 0.5 * (1.0 + math.cos(math.pi * progress)) return lr_min + (lr_max - lr_min) * cos # RadiomapModelV15 class + helpers moved to alt_models/joint_model.py # _extract_level_tensors, _pack_levels_to_dense are methods there. def compute_radiomap_metrics(pred_db: torch.Tensor, gt_db: torch.Tensor, mask_hw: torch.Tensor) -> dict[str, float]: """pred_db, gt_db, mask_hw ∈ [B, H, W] (mask bool).""" m = mask_hw.bool() if m.sum() == 0: return {"radiomap_mae_db": 0.0, "radiomap_psnr_db": 0.0, "radiomap_bias_db": 0.0, "valid_cells": 0} d = (pred_db[m] - gt_db[m]).float() mae = d.abs().mean().item() bias = d.mean().item() mse = (d * d).mean().item() span = 120.0 # typical dB span -120..0 → 120 dB dynamic psnr = 20.0 * math.log10(span) - 10.0 * math.log10(max(mse, 1e-8)) return { "radiomap_mae_db": float(mae), "radiomap_psnr_db": float(psnr), "radiomap_bias_db": float(bias), "valid_cells": int(m.sum().item()), } def run_val( model: torch.nn.Module, loader: DataLoader, device: torch.device, amp_dtype: torch.dtype | None, huber_beta: float, world_size: int, gradient_weight: float = 0.0, ) -> dict[str, float]: model.eval() total = {"loss": 0.0, "gain": 0.0, "grad": 0.0, "mae": 0.0, "psnr": 0.0, "bias": 0.0, "cells": 0.0, "n": 0.0} with torch.no_grad(): for batch in loader: voxel_feats_list = [t.to(device, non_blocking=True) for t in batch["voxel_feats_list"]] voxel_coords_list = [t.to(device, non_blocking=True) for t in batch["voxel_coords_list"]] tx_xyz = batch["tx_xyz_norm"].to(device, non_blocking=True) rx_grid = batch["rx_grid_xyz_norm"].to(device, non_blocking=True) gt = batch["gt_radiomap_db"].to(device, non_blocking=True) mask = batch["extent_mask"].to(device, non_blocking=True) # V1.5: metric fields tx_xyz_m = batch["tx_xyz_metric"].to(device, non_blocking=True) rx_grid_m = batch["rx_grid_xyz_metric"].to(device, non_blocking=True) scene_ext_m = batch["scene_extent_xyz_m"].to(device, non_blocking=True) origin_m_list = [t.to(device, non_blocking=True) for t in batch["voxel_origin_m_list"]] cell_size_m = float(batch["voxel_cell_size_m"]) cm = torch.autocast(device_type="cuda", dtype=amp_dtype) if amp_dtype is not None else nullcontext() with cm: pred = model( voxel_feats_list, voxel_coords_list, tx_xyz_norm=tx_xyz, tx_xyz_metric=tx_xyz_m, rx_grid_xyz_norm=rx_grid, rx_grid_xyz_metric=rx_grid_m, scene_extent_xyz_m=scene_ext_m, voxel_origin_m_list=origin_m_list, voxel_cell_size_m=cell_size_m, extent_mask=mask, ) mh = mask.squeeze(1).bool() if mh.sum() > 0: gain = F.smooth_l1_loss(pred[mh], gt[mh], beta=huber_beta).item() else: gain = 0.0 if gradient_weight > 0 and mh.any(): grad = gradient_l1_loss(pred, gt, mh).item() else: grad = 0.0 loss = gain + gradient_weight * grad mets = compute_radiomap_metrics(pred.float(), gt, mask.squeeze(1)) B = gt.shape[0] total["loss"] += loss * B total["gain"] += gain * B total["grad"] += grad * B total["mae"] += mets["radiomap_mae_db"] * B total["psnr"] += mets["radiomap_psnr_db"] * B total["bias"] += mets["radiomap_bias_db"] * B total["cells"] += mets["valid_cells"] total["n"] += B # all-reduce across ranks t = torch.tensor([total[k] for k in ("loss", "gain", "grad", "mae", "psnr", "bias", "cells", "n")], dtype=torch.float64, device=device) dist.all_reduce(t, op=dist.ReduceOp.SUM) loss, gain, grad, mae, psnr, bias, cells, n = t.tolist() model.train() if n == 0: return {"val_loss": 0.0, "val_mae": 0.0, "val_psnr": 0.0, "val_bias": 0.0, "val_cells": 0.0} return { "val_loss": loss / n, "val_gain_loss": gain / n, "val_grad_loss": grad / n, "val_mae_db": mae / n, "val_psnr_db": psnr / n, "val_bias_db": bias / n, "val_cells_mean": cells / n, } def composite_score(val: dict[str, float]) -> float: """Selector for best.pt. Lower MAE is better, higher PSNR is better.""" return -val.get("val_mae_db", 1e9) + 0.1 * val.get("val_psnr_db", 0.0) - 0.1 * abs(val.get("val_bias_db", 0.0)) def save_ckpt( path: Path, model: torch.nn.Module, opt: torch.optim.Optimizer, epoch: int, step: int, metrics: dict[str, float], cli_args: dict[str, Any], phase_name: str = "unknown", ) -> None: raw = model.module if isinstance(model, DDP) else model torch.save( { "phase_name": phase_name, "epoch": int(epoch), "step": int(step), "metrics": metrics, "cli_args": cli_args, "model_state_dict": {k: v.detach().cpu() for k, v in raw.state_dict().items()}, "optimizer_state_dict": opt.state_dict(), }, path, ) def main() -> None: args = parse_args() rank, local_rank, world_size, device = setup_ddp() torch.manual_seed(args.seed + rank) run_root = Path(args.run_root) reports_dir = run_root / "reports" artifacts_dir = run_root / "artifacts" ckpt_dir = run_root / "checkpoints" for d in (reports_dir, reports_dir / "plots" / "best", reports_dir / "plots" / "last", artifacts_dir, ckpt_dir): d.mkdir(parents=True, exist_ok=True) with open(args.manifest) as f: manifest = json.load(f) train_records = manifest["train"] val_records = manifest["val_heldout"] if rank == 0: print("============== WiSER radiomap V1 training ==============") print(f"world_size={world_size} per-rank bs={args.batch_size} global bs={args.batch_size * world_size}") print(f"epochs={args.epochs} warmup={args.warmup_epochs} lr={args.lr}→{args.lr_min}") print(f"manifest: train={len(train_records)} val={len(val_records)}") print(f"run_root: {run_root}") snap = {k: (v if not isinstance(v, Path) else str(v)) for k, v in vars(args).items()} snap["world_size"] = world_size snap["num_train_records"] = len(train_records) snap["num_val_records"] = len(val_records) with open(artifacts_dir / "config_snapshot.json", "w") as f: json.dump(snap, f, indent=2) with open(artifacts_dir / "env.txt", "w") as f: f.write(f"python={platform.python_version()}\n") f.write(f"torch={torch.__version__}\n") f.write(f"cuda_available={torch.cuda.is_available()}\n") if torch.cuda.is_available(): f.write(f"cuda={torch.version.cuda}\n") f.write(f"gpu0={torch.cuda.get_device_name(0)}\n") f.write(f"device_count={torch.cuda.device_count()}\n") dist.barrier() train_ds = RadiomapDataset(train_records, channels=args.channels, grid_h=args.grid_h, grid_w=args.grid_w, db_floor=args.db_floor, dataset_kind=args.dataset_kind) val_ds = RadiomapDataset(val_records, channels=args.channels, grid_h=args.grid_h, grid_w=args.grid_w, db_floor=args.db_floor, dataset_kind=args.dataset_kind) train_sampler = DistributedSampler(train_ds, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True, seed=args.seed) val_sampler = DistributedSampler(val_ds, num_replicas=world_size, rank=rank, shuffle=False, drop_last=False) train_loader = DataLoader( train_ds, batch_size=args.batch_size, sampler=train_sampler, collate_fn=radiomap_collate, num_workers=args.num_workers, pin_memory=True, persistent_workers=args.num_workers > 0, drop_last=True, ) val_loader = DataLoader( val_ds, batch_size=args.val_batch_size, sampler=val_sampler, collate_fn=radiomap_collate, num_workers=max(1, args.num_workers // 2), pin_memory=True, persistent_workers=False, drop_last=False, ) shared = SharedModelConfig( backbone_kind=args.backbone_kind, backbone_channels=args.channels, backbone_downsample_stages=args.downsample_stages, backbone_blocks_per_stage=args.blocks_per_stage, ) # ---- Build JointRadiomapCIRModel ---- # Detect CIR head arch from whatever ckpt we're loading (D2.2 or phase ckpt). # detect_csi_head_arch_from_ckpt handles both "head.*" and "csi_head.*" prefixes. csi_arch_kwargs = {} ckpt_for_arch = args.load_phase_ckpt or args.load_d22_ckpt if ckpt_for_arch: info = detect_csi_head_arch_from_ckpt(ckpt_for_arch) csi_arch_kwargs = {k: v for k, v in info.items() if k.startswith("csi_")} if "backbone_downsample_stages" in info: shared.backbone_downsample_stages = info["backbone_downsample_stages"] if "backbone_blocks_per_stage" in info: shared.backbone_blocks_per_stage = info["backbone_blocks_per_stage"] if rank == 0: print(f"[arch-detect] from {ckpt_for_arch}: {csi_arch_kwargs}") model = JointRadiomapCIRModel( shared=shared, channels=args.channels, spatial_channels=args.spatial_channels, stem_depth=args.stem_depth, output_depth=args.output_depth, grid_h=args.grid_h, grid_w=args.grid_w, corridor_budget=args.corridor_budget, num_cross_layers=args.num_cross_layers, num_self_layers=args.num_self_layers, **csi_arch_kwargs, ).to(device) # ---- Load checkpoints ---- if args.load_phase_ckpt: # Continuing from a previous phase (full model state) meta = load_phase_ckpt_into_model(args.load_phase_ckpt, model) if rank == 0: print(f"[load] phase ckpt from {meta.get('phase_name')} ep={meta.get('epoch')} metrics={meta.get('metrics', {})}") elif args.load_v22a_ckpt and args.load_d22_ckpt: # Warm start: V2.2a (rm head) + D2.2 (CIR head) report = load_warm_start_ckpt(model, args.load_v22a_ckpt, args.load_d22_ckpt, verify=True) if rank == 0: print(f"[load] warm-start: V2.2a + D2.2 merged: {report}") elif args.load_d22_ckpt: # Legacy: only D2.2 (no radiomap head init, like V2.2a) import torch as _torch d22_sd = _torch.load(args.load_d22_ckpt, map_location="cpu", weights_only=False)["model_state_dict"] keep = {k: v for k, v in d22_sd.items() if k.startswith("backbone.") or k.startswith("tx_proj.")} # Also take head.* → csi_head.* csi_remap = {k.replace("head.", "csi_head.", 1): v for k, v in d22_sd.items() if k.startswith("head.")} model.load_state_dict(keep, strict=False) model.load_state_dict(csi_remap, strict=False) if rank == 0: print(f"[load] D2.2 only: backbone+tx_proj ({len(keep)}) + csi_head ({len(csi_remap)})") # ---- Apply freeze config ---- freeze_cfg = FreezeConfig( backbone=args.freeze_backbone or args.freeze_base, tx_proj=args.freeze_tx_proj or args.freeze_base, radiomap_head=args.freeze_radiomap_head, csi_head=args.freeze_csi_head or True, # radiomap phase ALWAYS freezes CIR head (it's not used in forward_radiomap) ) model.set_freeze_state(freeze_cfg) # DDP wraps .forward, so bind .forward = .forward_radiomap for this phase model.forward = model.forward_radiomap n_total = sum(p.numel() for p in model.parameters()) n_train = sum(p.numel() for p in model.parameters() if p.requires_grad) if rank == 0: print(f"[phase={args.phase_name}] params: total={n_total/1e6:.2f}M trainable={n_train/1e6:.2f}M frozen={(n_total-n_train)/1e6:.2f}M") print(f"[phase={args.phase_name}] freeze: backbone={freeze_cfg.backbone} tx_proj={freeze_cfg.tx_proj} rm_head={freeze_cfg.radiomap_head} csi_head={freeze_cfg.csi_head}") # find_unused_parameters when any part frozen any_frozen = any([freeze_cfg.backbone, freeze_cfg.tx_proj, freeze_cfg.radiomap_head, freeze_cfg.csi_head]) ddp = DDP(model, device_ids=[local_rank], find_unused_parameters=any_frozen, gradient_as_bucket_view=True) opt = torch.optim.Adam([p for p in ddp.parameters() if p.requires_grad], lr=args.lr, weight_decay=args.weight_decay) steps_per_epoch = max(1, len(train_loader)) total_steps = args.epochs * steps_per_epoch warmup_steps = args.warmup_epochs * steps_per_epoch amp_dtype = torch.bfloat16 csv_path = reports_dir / "train_curves.csv" json_path = reports_dir / "train_curves.json" if rank == 0: with open(csv_path, "w", newline="") as f: w = csv.writer(f) w.writerow([ "epoch", "step", "lr", "train_loss", "train_gain_loss", "train_grad_loss", "train_mae_db", "train_psnr_db", "train_bias_db", "val_loss", "val_gain_loss", "val_grad_loss", "val_mae_db", "val_psnr_db", "val_bias_db", "val_cells_mean", "secs_train", "secs_val", "secs_per_epoch", "peak_mem_gb", ]) best_score = -1e9 history: list[dict] = [] step = 0 for epoch in range(args.epochs): ddp.train() train_sampler.set_epoch(epoch) epoch_start = time.perf_counter() torch.cuda.reset_peak_memory_stats(device) ep_loss = 0.0; ep_gain = 0.0; ep_grad = 0.0 ep_mae = 0.0; ep_psnr = 0.0; ep_bias = 0.0; ep_n = 0 for batch in train_loader: lr = compute_lr(step, total_steps, warmup_steps, args.lr, args.lr_min) for pg in opt.param_groups: pg["lr"] = lr voxel_feats_list = [t.to(device, non_blocking=True) for t in batch["voxel_feats_list"]] voxel_coords_list = [t.to(device, non_blocking=True) for t in batch["voxel_coords_list"]] tx_xyz = batch["tx_xyz_norm"].to(device, non_blocking=True) rx_grid = batch["rx_grid_xyz_norm"].to(device, non_blocking=True) gt = batch["gt_radiomap_db"].to(device, non_blocking=True) mask = batch["extent_mask"].to(device, non_blocking=True) # V1.5 metric fields tx_xyz_m = batch["tx_xyz_metric"].to(device, non_blocking=True) rx_grid_m = batch["rx_grid_xyz_metric"].to(device, non_blocking=True) scene_ext_m = batch["scene_extent_xyz_m"].to(device, non_blocking=True) origin_m_list = [t.to(device, non_blocking=True) for t in batch["voxel_origin_m_list"]] cell_size_m = float(batch["voxel_cell_size_m"]) opt.zero_grad(set_to_none=True) with torch.autocast(device_type="cuda", dtype=amp_dtype): pred = ddp( voxel_feats_list, voxel_coords_list, tx_xyz_norm=tx_xyz, tx_xyz_metric=tx_xyz_m, rx_grid_xyz_norm=rx_grid, rx_grid_xyz_metric=rx_grid_m, scene_extent_xyz_m=scene_ext_m, voxel_origin_m_list=origin_m_list, voxel_cell_size_m=cell_size_m, extent_mask=mask, ) mh = mask.squeeze(1).bool() if mh.sum() > 0: loss_gain = F.smooth_l1_loss(pred[mh], gt[mh], beta=args.huber_beta) else: loss_gain = pred.sum() * 0.0 # V1.4: WiSER spatial decoder style gradient L1 loss if args.gradient_weight > 0: loss_grad = gradient_l1_loss(pred, gt, mh) else: loss_grad = pred.sum() * 0.0 loss = loss_gain + args.gradient_weight * loss_grad loss.backward() opt.step() with torch.no_grad(): m = compute_radiomap_metrics(pred.float(), gt, mask.squeeze(1)) B = gt.shape[0] ep_loss += loss.item() * B ep_gain += loss_gain.item() * B ep_grad += loss_grad.item() * B ep_mae += m["radiomap_mae_db"] * B ep_psnr += m["radiomap_psnr_db"] * B ep_bias += m["radiomap_bias_db"] * B ep_n += B step += 1 # reduce train metrics across ranks t_train = torch.tensor([ep_loss, ep_gain, ep_grad, ep_mae, ep_psnr, ep_bias, float(ep_n)], dtype=torch.float64, device=device) dist.all_reduce(t_train, op=dist.ReduceOp.SUM) tl, tgain, tgrad, tmae, tpsnr, tbias, tn = t_train.tolist() train_stats = { "train_loss": tl / max(tn, 1), "train_gain_loss": tgain / max(tn, 1), "train_grad_loss": tgrad / max(tn, 1), "train_mae_db": tmae / max(tn, 1), "train_psnr_db": tpsnr / max(tn, 1), "train_bias_db": tbias / max(tn, 1), } train_end = time.perf_counter() secs_train = train_end - epoch_start peak_mem_gb = torch.cuda.max_memory_allocated(device) / 1024**3 val_stats: dict[str, float] = {} secs_val = 0.0 if (epoch + 1) % args.val_every_epochs == 0 or epoch == args.epochs - 1: val_start = time.perf_counter() val_stats = run_val(ddp, val_loader, device, amp_dtype, args.huber_beta, world_size, gradient_weight=args.gradient_weight) dist.barrier() secs_val = time.perf_counter() - val_start secs = time.perf_counter() - epoch_start if rank == 0: ep_entry = { "epoch": int(epoch), "step": int(step), "lr": float(lr), **train_stats, **val_stats, "secs_train": float(secs_train), "secs_val": float(secs_val), "secs_per_epoch": float(secs), "peak_mem_gb": float(peak_mem_gb), } history.append(ep_entry) with open(csv_path, "a", newline="") as f: w = csv.writer(f) w.writerow([ epoch, step, f"{lr:.6e}", f"{train_stats['train_loss']:.4f}", f"{train_stats.get('train_gain_loss', 0):.4f}", f"{train_stats.get('train_grad_loss', 0):.4f}", f"{train_stats['train_mae_db']:.4f}", f"{train_stats['train_psnr_db']:.4f}", f"{train_stats['train_bias_db']:.4f}", f"{val_stats.get('val_loss', ''):.4f}" if val_stats else "", f"{val_stats.get('val_gain_loss', ''):.4f}" if val_stats else "", f"{val_stats.get('val_grad_loss', ''):.4f}" if val_stats else "", f"{val_stats.get('val_mae_db', ''):.4f}" if val_stats else "", f"{val_stats.get('val_psnr_db', ''):.4f}" if val_stats else "", f"{val_stats.get('val_bias_db', ''):.4f}" if val_stats else "", f"{val_stats.get('val_cells_mean', ''):.1f}" if val_stats else "", f"{secs_train:.2f}", f"{secs_val:.2f}", f"{secs:.2f}", f"{peak_mem_gb:.2f}", ]) with open(json_path, "w") as f: json.dump({"history": history}, f, indent=2) msg = ( f"[ep {epoch:3d}/{args.epochs}] " f"lr={lr:.2e} " f"tr(gain={train_stats.get('train_gain_loss',0):.3f},grad={train_stats.get('train_grad_loss',0):.3f}) " f"mae={train_stats['train_mae_db']:.2f}dB " f"psnr={train_stats['train_psnr_db']:.2f} bias={train_stats['train_bias_db']:+.2f}" ) if val_stats: msg += ( f" | val(gain={val_stats.get('val_gain_loss',0):.3f},grad={val_stats.get('val_grad_loss',0):.3f}) " f"mae={val_stats['val_mae_db']:.2f}dB " f"psnr={val_stats['val_psnr_db']:.2f} " f"bias={val_stats['val_bias_db']:+.2f}" ) msg += f" | train={secs_train:.1f}s val={secs_val:.1f}s tot={secs:.1f}s peak={peak_mem_gb:.1f}GB" print(msg, flush=True) # best.pt by composite score on val if val_stats: score = composite_score(val_stats) if score > best_score: best_score = score save_ckpt(ckpt_dir / "best.pt", ddp, opt, epoch, step, {**train_stats, **val_stats, "composite_score": score}, vars(args), phase_name=args.phase_name) print(f" ✓ new best.pt (epoch {epoch}, score {score:.4f}, val_mae_db {val_stats['val_mae_db']:.3f})", flush=True) save_ckpt(ckpt_dir / "last.pt", ddp, opt, epoch, step, {**train_stats, **val_stats}, vars(args), phase_name=args.phase_name) dist.barrier() if rank == 0: save_ckpt(ckpt_dir / "final.pt", ddp, opt, args.epochs - 1, step, history[-1] if history else {}, vars(args), phase_name=args.phase_name) print("============== V1 training DONE ==============") cleanup_ddp() if __name__ == "__main__": main()