code_snapshots/WiSER_public_release_code/wiser/alt_engine/checkpoint.py

"""WiSER Checkpoint utilities: auto-detect head arch + load V2.2a + D2.2 warm-start."""

from __future__ import annotations

import torch
from pathlib import Path


def detect_csi_head_arch_from_ckpt(ckpt_path: str) -> dict:
    """Inspect a D2.2-style CIR ckpt OR an WiSER phase ckpt to detect CIR head arch.

    Supports both:
      - D2.2 style: keys start with "head.*" (CIR-only model)
      - WiSER phase ckpt style: keys start with "csi_head.*" (joint model)

    Returns kwargs dict suitable for JointRadiomapCIRModel(csi_*=...).
    """
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt["model_state_dict"]
    keys = list(sd.keys())

    # Detect which prefix the ckpt uses for CIR head
    csi_prefix = None
    for candidate in ("csi_head.", "head."):
        if any(k.startswith(candidate) for k in keys):
            csi_prefix = candidate
            break
    if csi_prefix is None:
        return {"_ckpt_model_config": ckpt.get("model_config", {})}

    info = {}

    # num_decoder_layers: count <prefix>decoder.layers.N
    dec_prefix = csi_prefix + "decoder.layers."
    layer_indices = set()
    for k in keys:
        if k.startswith(dec_prefix):
            idx = int(k[len(dec_prefix):].split(".")[0])
            layer_indices.add(idx)
    if layer_indices:
        info["csi_num_decoder_layers"] = max(layer_indices) + 1

    # Detect deep arch for each sub-head
    def _detect_deep(sub_prefix: str) -> tuple[bool, int]:
        pn_key_prefix = f"{csi_prefix}{sub_prefix}.pre_norm"
        om_key_prefix = f"{csi_prefix}{sub_prefix}.out_mlp"
        is_deep = any(
            (pn_key_prefix in k or om_key_prefix in k)
            for k in keys
        )
        hidden = 512
        if is_deep:
            om_0 = f"{csi_prefix}{sub_prefix}.out_mlp.0.weight"
            if om_0 in sd:
                hidden = int(sd[om_0].shape[0])
        return is_deep, hidden

    for prefix_name, arg_base in [
        ("peak_db_head", "csi_peak_db"),
        ("delay_head", "csi_delay"),
        ("exists_head", "csi_exists"),
    ]:
        deep, hid = _detect_deep(prefix_name)
        info[f"{arg_base}_head_arch"] = "deep" if deep else "flat"
        info[f"{arg_base}_hidden"] = hid

    # Backbone config (may be used to override shared config)
    backbone_stage_indices = set()
    backbone_block_indices = set()
    for k in keys:
        if k.startswith("backbone.scene_stages."):
            parts = k.split(".")
            if len(parts) >= 3:
                backbone_stage_indices.add(parts[2])
                if len(parts) >= 4 and parts[2] == "0":
                    backbone_block_indices.add(parts[3])
    if backbone_stage_indices:
        info["backbone_downsample_stages"] = len(backbone_stage_indices)
    if backbone_block_indices:
        info["backbone_blocks_per_stage"] = len(backbone_block_indices)

    info["_ckpt_model_config"] = ckpt.get("model_config", {})
    info["_detected_csi_prefix"] = csi_prefix
    return info


def verify_backbone_identical(v22a_ckpt_path: str, d22_ckpt_path: str) -> int:
    """Verify V2.2a's backbone == D2.2's backbone (since V2.2a was frozen-backbone training).
    Returns number of matching backbone keys. Raises if ANY diff found."""
    v22a_sd = torch.load(v22a_ckpt_path, map_location="cpu", weights_only=False)["model_state_dict"]
    d22_sd = torch.load(d22_ckpt_path, map_location="cpu", weights_only=False)["model_state_dict"]
    v22a_b = {k: v for k, v in v22a_sd.items() if k.startswith("backbone.")}
    d22_b = {k: v for k, v in d22_sd.items() if k.startswith("backbone.")}
    if set(v22a_b.keys()) != set(d22_b.keys()):
        only_v22a = set(v22a_b) - set(d22_b)
        only_d22 = set(d22_b) - set(v22a_b)
        raise RuntimeError(f"backbone keys differ. V22a-only={only_v22a} D22-only={only_d22}")
    diff_keys = []
    for k in v22a_b:
        if not torch.allclose(v22a_b[k].float(), d22_b[k].float(), atol=1e-6):
            diff_keys.append(k)
    if diff_keys:
        raise RuntimeError(f"V2.2a backbone != D2.2 backbone at {len(diff_keys)} keys (V2.2a was supposed to freeze). Examples: {diff_keys[:3]}")
    return len(v22a_b)


def load_warm_start_ckpt(
    model: torch.nn.Module,
    v22a_ckpt_path: str,
    d22_ckpt_path: str,
    verify: bool = True,
) -> dict:
    """Load V2.2a (backbone + tx_proj + radiomap parts) + D2.2 CIR head into model.

    Returns a report dict with load statistics.
    """
    report = {}

    # Sanity: V2.2a backbone should equal D2.2 backbone
    if verify:
        n_match = verify_backbone_identical(v22a_ckpt_path, d22_ckpt_path)
        report["backbone_keys_verified_identical"] = n_match

    # 1. Load V2.2a (backbone + tx_proj + radiomap head components)
    v22a_sd = torch.load(v22a_ckpt_path, map_location="cpu", weights_only=False)["model_state_dict"]
    missing, unexpected = model.load_state_dict(v22a_sd, strict=False)

    # Expected: all csi_head.* keys should be missing (V2.2a has no CIR head)
    allowed_missing = lambda k: k.startswith("csi_head.")
    bad_missing = [k for k in missing if not allowed_missing(k)]
    if bad_missing:
        raise RuntimeError(f"Unexpected missing keys from V2.2a load: {bad_missing[:5]}")
    if unexpected:
        raise RuntimeError(f"Unexpected keys from V2.2a load: {unexpected[:5]}")
    report["v22a_loaded_keys"] = len(v22a_sd)
    report["v22a_missing_csi_head_keys"] = len(missing)

    # 2. Load D2.2 CIR head (head.* → csi_head.*)
    d22_sd = torch.load(d22_ckpt_path, map_location="cpu", weights_only=False)["model_state_dict"]
    csi_head_remapped = {
        k.replace("head.", "csi_head.", 1): v
        for k, v in d22_sd.items()
        if k.startswith("head.")
    }
    missing2, unexpected2 = model.load_state_dict(csi_head_remapped, strict=False)
    # Expected: all radiomap + backbone + tx_proj keys should be "missing" (they're not in this load, but already loaded from V2.2a)
    # The only "bad" case is if some csi_head.* keys are unexpected (shape mismatch or extra)
    unexpected_csi = [k for k in unexpected2 if k.startswith("csi_head.")]
    if unexpected_csi:
        raise RuntimeError(f"Unexpected csi_head keys from D2.2 load: {unexpected_csi[:5]}")
    report["d22_csi_head_loaded_keys"] = len(csi_head_remapped)

    # 3. Sanity: no uninitialized params
    uninit = [n for n, p in model.named_parameters() if p.abs().sum().item() == 0.0]
    # Allow out_mlp.*weight == 0 (deep heads zero-init last linear by design)
    uninit_bad = [n for n in uninit if "out_mlp" not in n]
    if uninit_bad:
        raise RuntimeError(f"Uninitialized params after warm-start load: {uninit_bad[:5]}")
    report["params_uninitialized_excl_out_mlp"] = len(uninit_bad)
    report["params_uninitialized_total"] = len(uninit)

    return report


def save_phase_ckpt(
    ckpt_path: Path,
    model: torch.nn.Module,
    optimizer,
    phase_name: str,
    epoch: int,
    step: int,
    metrics: dict,
    cli_args: dict,
    freeze_cfg_dict: dict,
) -> None:
    """Save a phase-end checkpoint."""
    import torch
    from torch.nn.parallel import DistributedDataParallel as DDP
    raw_model = 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,
        "freeze_cfg": freeze_cfg_dict,
        "model_state_dict": {k: v.detach().cpu() for k, v in raw_model.state_dict().items()},
        "optimizer_state_dict": optimizer.state_dict() if optimizer is not None else {},
    }, ckpt_path)


def load_phase_ckpt_into_model(ckpt_path: str, model: torch.nn.Module) -> dict:
    """Load a previous phase's end checkpoint into model (strict full load).
    Returns the ckpt metadata (epoch, step, metrics)."""
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    sd = ckpt["model_state_dict"]
    # We want STRICT load here because phases save full model state
    missing, unexpected = model.load_state_dict(sd, strict=False)
    if missing or unexpected:
        raise RuntimeError(f"Phase-ckpt load mismatch. missing={missing[:3]} unexpected={unexpected[:3]}")
    return {
        "phase_name": ckpt.get("phase_name"),
        "epoch": ckpt.get("epoch"),
        "step": ckpt.get("step"),
        "metrics": ckpt.get("metrics", {}),
    }


__all__ = [
    "detect_csi_head_arch_from_ckpt",
    "verify_backbone_identical",
    "load_warm_start_ckpt",
    "save_phase_ckpt",
    "load_phase_ckpt_into_model",
]