code_snapshots/WiSER_public_release_code/scripts/train_phase_cir.py

"""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()