"""WiSER Phase trainer for CIR task (Phase 1: restore CIR head, Phase 3: refine CIR). Uses JointRadiomapCIRModel + CsiSetTrainer (from WiSER CIR engine) for loss + matching. Minimal training loop: DDP + bf16 autocast + cosine schedule + best/last ckpt. """ 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 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, TrainConfig, reject_radiomap_artifacts, ) from wiser.data.collate import triple_collate from wiser.data.csi_path_targets import MergedPathTargetConfig from wiser.data.dataset import MultiSceneTripleDataset from wiser.engine.matching import hungarian_match_batch, MatcherConfig from wiser.engine.trainer import CsiSetTrainer, TrainerConfig from wiser.utils.metrics import compute_metric_bundle 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, help="CIR triples manifest (e.g. tiny100_5k_4tx_balanced_manifest.json)") p.add_argument("--val-manifest", default=None, help="Val manifest (e.g. tiny10_triples_manifest.json). Default: same as --manifest.") p.add_argument("--run-root", required=True) p.add_argument("--phase-name", default="P1_restore_csi") p.add_argument("--epochs", type=int, default=30) p.add_argument("--warmup-epochs", type=int, default=2) p.add_argument("--batch-size", type=int, default=2048, help="triples per rank (DataLoader batch_size). D2.2 used 4096.") p.add_argument("--val-batch-size", type=int, default=4096) p.add_argument("--val-every-epochs", type=int, default=5) p.add_argument("--val-max-triples", type=int, default=5000, help="cap val triples to keep val fast (default 5000)") p.add_argument("--num-workers", type=int, default=4) p.add_argument("--lr", type=float, default=5e-5) p.add_argument("--lr-min", type=float, default=1e-7) 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("--seed", type=int, default=0) p.add_argument("--backbone-kind", default="trellis2", choices=["trellis2", "dense_fallback"]) # Radiomap head params (needed to construct full JointModel even for CIR phases) 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("--corridor-budget", type=int, default=192) p.add_argument("--num-cross-layers", type=int, default=4) p.add_argument("--num-self-layers", type=int, default=2) # Checkpoint loading p.add_argument("--load-phase-ckpt", default=None, help="Load prior-phase ckpt (supersedes --load-v22a/--load-d22)") p.add_argument("--load-v22a-ckpt", default=None) p.add_argument("--load-d22-ckpt", default=None, help="For warm start or to provide CIR head init") # Freeze flags p.add_argument("--freeze-backbone", action="store_true") 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") # Loss / matcher config p.add_argument("--matcher-backend", default="scipy", choices=["scipy", "sinkhorn"]) p.add_argument("--no-object-exists-weight", type=float, default=5.0) return p.parse_args() 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 def composite_csi_score(val: dict) -> float: """CIR composite: higher pdp_cosine, lower delay/peak/fpr = better.""" return ( + 1.0 * val.get("pdp_cosine_nonzero", 0.0) - 0.05 * val.get("peak_db_mae_matched", 0.0) - 0.1 * val.get("delay_mae_matched_ns", 0.0) + 1.5 * val.get("count_acc", 0.0) - 0.5 * val.get("zero_path_fpr", 0.0) ) def run_val( model: torch.nn.Module, loader: DataLoader, trainer: CsiSetTrainer, device: torch.device, amp_dtype: torch.dtype | None, world_size: int, max_batches: int = 0, ) -> dict[str, float]: model.eval() total = { "pdp_cosine_nonzero": 0.0, "peak_db_mae_matched": 0.0, "peak_db_bias_matched": 0.0, "delay_mae_matched_ns": 0.0, "count_acc": 0.0, "nonzero_path_count_mae": 0.0, "zero_path_fpr": 0.0, "match_rate_over_all_samples": 0.0, "n": 0, } with torch.no_grad(): for bi, batch in enumerate(loader): if max_batches > 0 and bi >= max_batches: break for k, v in batch.items(): if isinstance(v, torch.Tensor): batch[k] = v.to(device, non_blocking=True) if "scene_voxel_levels" in batch: batch["scene_voxel_levels"] = [ {kk: (tt.to(device, non_blocking=True) if torch.is_tensor(tt) else tt) for kk, tt in sv.items()} if isinstance(sv, dict) else sv for sv in batch.get("scene_voxel_levels", []) ] gt = { "gt_num_paths": batch["gt_num_paths"], "gt_delay_ns": batch["gt_delay_ns"], "gt_peak_db": batch["gt_peak_db"], "gt_path_mask": batch["gt_path_mask"], "gt_truncated": batch["gt_truncated"], } out = trainer.step(batch, gt) mets = compute_metric_bundle( out["predictions"], gt, out["matchings"], query_budget=trainer.cfg.model.query_budget if hasattr(trainer.cfg.model, 'query_budget') else 8, ) B = gt["gt_num_paths"].shape[0] for k in total: if k == "n": continue total[k] += float(mets.get(k, 0.0)) * B total["n"] += B # all-reduce t = torch.tensor([total[k] for k in total], dtype=torch.float64, device=device) dist.all_reduce(t, op=dist.ReduceOp.SUM) vals = t.tolist() model.train() n = vals[-1] if n == 0: return {} keys = list(total.keys()) out = {keys[i]: vals[i] / n for i in range(len(keys) - 1)} out["val_n_samples"] = n return out def main() -> None: args = parse_args() rank, local_rank, world_size, device = setup_ddp() torch.manual_seed(args.seed + rank) np.random.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, artifacts_dir, ckpt_dir): d.mkdir(parents=True, exist_ok=True) if rank == 0: print(f"========== {args.phase_name} (CIR phase) ==========") print(f"world_size={world_size} per-rank bs={args.batch_size} global bs={args.batch_size*world_size}") # Load manifest reject_radiomap_artifacts(args.manifest) manifest = json.loads(Path(args.manifest).read_text()) val_manifest_path = args.val_manifest or args.manifest val_manifest = json.loads(Path(val_manifest_path).read_text()) # Build model 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, ) # Detect CIR head arch from whichever ckpt we load (D2.2 or phase ckpt) 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: 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')}") elif args.load_v22a_ckpt and args.load_d22_ckpt: report = load_warm_start_ckpt(model, args.load_v22a_ckpt, args.load_d22_ckpt, verify=True) if rank == 0: print(f"[load] warm-start: {report}") elif args.load_d22_ckpt: 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.")} 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: {len(keep)} backbone+tx + {len(csi_remap)} csi_head") # Apply freeze (CIR phase defaults: radiomap_head frozen = True, everything else trainable) freeze_cfg = FreezeConfig( backbone=args.freeze_backbone, tx_proj=args.freeze_tx_proj, radiomap_head=args.freeze_radiomap_head or True, # CIR phase always freezes radiomap head csi_head=args.freeze_csi_head, # default: CIR head trains ) model.set_freeze_state(freeze_cfg) 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}") # DDP 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) # CsiSetTrainer: use the raw model (not DDP wrapper) for its internal methods trainer_cfg = TrainerConfig() trainer_cfg.matcher = MatcherConfig(backend=args.matcher_backend) from wiser.engine.losses import LossWeights trainer_cfg.loss_weights = LossWeights( no_object_exists=args.no_object_exists_weight, ) trainer_cfg.model = model.config # shim config (backbone.kind only) # CsiSetTrainer needs model with .head / .scene_encode / .modulate_scene_tx / .encode_tx trainer = CsiSetTrainer(model, trainer_cfg) # Opt + scheduler opt = torch.optim.Adam([p for p in ddp.parameters() if p.requires_grad], lr=args.lr, weight_decay=args.weight_decay) # Dataset + loader target_cfg = MergedPathTargetConfig() train_ds = MultiSceneTripleDataset(manifest, target_config=target_cfg, voxel_channels=args.channels) val_ds_full = MultiSceneTripleDataset(val_manifest, target_config=target_cfg, voxel_channels=args.channels) # Cap val to speed up if args.val_max_triples > 0 and len(val_ds_full.triples) > args.val_max_triples: val_ds_full.triples = val_ds_full.triples[:args.val_max_triples] train_sampler = DistributedSampler(train_ds, num_replicas=world_size, rank=rank, shuffle=True, drop_last=True, seed=args.seed) val_sampler = DistributedSampler(val_ds_full, 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=triple_collate, num_workers=args.num_workers, pin_memory=True, persistent_workers=args.num_workers > 0, drop_last=True, ) val_loader = DataLoader( val_ds_full, batch_size=args.val_batch_size, sampler=val_sampler, collate_fn=triple_collate, num_workers=max(1, args.num_workers // 2), pin_memory=True, persistent_workers=False, drop_last=False, ) if rank == 0: print(f"[phase={args.phase_name}] train: {len(train_ds)} triples, steps/ep/rank={len(train_loader)}") print(f"[phase={args.phase_name}] val: {len(val_ds_full.triples)} triples") 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" 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_n/a", "train_peak_mae", "val_peak_db_mae_matched", "val_delay_mae_matched_ns", "val_count_acc", "val_fpr", "val_pdp_cosine", "secs_train", "secs_val", "peak_mem_gb", ]) # Config snapshot snap = {k: (v if not isinstance(v, Path) else str(v)) for k, v in vars(args).items()} snap["world_size"] = world_size with open(artifacts_dir / "config_snapshot.json", "w") as f: json.dump(snap, f, indent=2) dist.barrier() best_score = -1e9 history = [] 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_peak_mae = 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 # Move to device for k, v in batch.items(): if isinstance(v, torch.Tensor): batch[k] = v.to(device, non_blocking=True) if "scene_voxel_levels" in batch: batch["scene_voxel_levels"] = [ {kk: (tt.to(device, non_blocking=True) if torch.is_tensor(tt) else tt) for kk, tt in sv.items()} if isinstance(sv, dict) else sv for sv in batch.get("scene_voxel_levels", []) ] gt = { "gt_num_paths": batch["gt_num_paths"], "gt_delay_ns": batch["gt_delay_ns"], "gt_peak_db": batch["gt_peak_db"], "gt_path_mask": batch["gt_path_mask"], "gt_truncated": batch["gt_truncated"], } opt.zero_grad(set_to_none=True) out = trainer.step(batch, gt) loss = out["loss_bundle"]["loss_total"] loss.backward() opt.step() # Quick per-batch metric B = gt["gt_num_paths"].shape[0] ep_loss += float(loss.item()) * B ep_n += B step += 1 # Reduce train loss t_train = torch.tensor([ep_loss, float(ep_n)], dtype=torch.float64, device=device) dist.all_reduce(t_train, op=dist.ReduceOp.SUM) tl, tn = t_train.tolist() train_stats = { "train_loss": tl / max(tn, 1), } train_end = time.perf_counter() secs_train = train_end - epoch_start peak_mem = torch.cuda.max_memory_allocated(device) / 1024**3 val_stats = {} 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, trainer, device, amp_dtype, world_size) 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), "peak_mem_gb": float(peak_mem), } 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"{val_stats.get('peak_db_mae_matched', ''):.4f}" if val_stats else "", f"{val_stats.get('peak_db_mae_matched', ''):.4f}" if val_stats else "", f"{val_stats.get('delay_mae_matched_ns', ''):.4f}" if val_stats else "", f"{val_stats.get('count_acc', ''):.4f}" if val_stats else "", f"{val_stats.get('zero_path_fpr', ''):.4f}" if val_stats else "", f"{val_stats.get('pdp_cosine_nonzero', ''):.4f}" if val_stats else "", f"{secs_train:.2f}", f"{secs_val:.2f}", f"{peak_mem:.2f}", ]) msg = f"[ep {epoch:3d}/{args.epochs}] lr={lr:.2e} tr_loss={train_stats['train_loss']:.4f}" if val_stats: msg += f" | val peak_mae={val_stats.get('peak_db_mae_matched',0):.2f} delay_mae={val_stats.get('delay_mae_matched_ns',0):.3f} count_acc={val_stats.get('count_acc',0):.3f} fpr={val_stats.get('zero_path_fpr',0):.3f} pdp={val_stats.get('pdp_cosine_nonzero',0):.4f}" msg += f" | train={secs_train:.1f}s val={secs_val:.1f}s peak={peak_mem:.1f}GB" print(msg, flush=True) if val_stats: score = composite_csi_score(val_stats) if score > best_score: best_score = score save_phase_ckpt(ckpt_dir / "best.pt", ddp.module, opt, args.phase_name, epoch, step, {**train_stats, **val_stats, "composite_score": score}, vars(args), {"backbone": freeze_cfg.backbone, "tx_proj": freeze_cfg.tx_proj, "radiomap_head": freeze_cfg.radiomap_head, "csi_head": freeze_cfg.csi_head}) print(f" ✓ new best.pt (epoch {epoch}, score {score:.4f})", flush=True) save_phase_ckpt(ckpt_dir / "last.pt", ddp.module, opt, args.phase_name, epoch, step, {**train_stats, **val_stats}, vars(args), {"backbone": freeze_cfg.backbone, "tx_proj": freeze_cfg.tx_proj, "radiomap_head": freeze_cfg.radiomap_head, "csi_head": freeze_cfg.csi_head}) dist.barrier() if rank == 0: print(f"========== {args.phase_name} DONE ==========", flush=True) cleanup_ddp() if __name__ == "__main__": main()