"""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 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", ]