()
| 151 | # ═══════════════════════════════════════════════════════════════════════════ |
| 152 | |
| 153 | def main(): |
| 154 | parser = argparse.ArgumentParser() |
| 155 | parser.add_argument("--checkpoint", required=True) |
| 156 | parser.add_argument("--image", required=True) |
| 157 | parser.add_argument("--outdir", default="tests/ref_phase3") |
| 158 | parser.add_argument("--blocks", default="0,1,7,15,23,31", |
| 159 | help="Comma-separated block indices to dump detailed intermediates") |
| 160 | args = parser.parse_args() |
| 161 | |
| 162 | os.makedirs(args.outdir, exist_ok=True) |
| 163 | device = "cpu" |
| 164 | dump_blocks = set(int(b) for b in args.blocks.split(",")) |
| 165 | |
| 166 | # ── Load checkpoint ───────────────────────────────────────────────── |
| 167 | print("Loading checkpoint...") |
| 168 | ckpt = torch.load(args.checkpoint, map_location=device, weights_only=False) |
| 169 | if "model" in ckpt: |
| 170 | ckpt = ckpt["model"] |
| 171 | |
| 172 | vit_prefix = "detector.backbone.vision_backbone.trunk." |
| 173 | vit_sd = {k[len(vit_prefix):]: v for k, v in ckpt.items() if k.startswith(vit_prefix)} |
| 174 | print(f" Found {len(vit_sd)} ViT keys") |
| 175 | |
| 176 | # ── Step 3.4: RoPE ───────────────────────────────────────────────── |
| 177 | print("\n=== RoPE Frequencies ===") |
| 178 | |
| 179 | # Window RoPE (24x24, scale_pos=1.0) |
| 180 | window_freqs = compute_axial_cis(dim=64, end_x=24, end_y=24, theta=10000.0, scale_pos=1.0) |
| 181 | # Store as real: [576, 32, 2] (cos, sin) |
| 182 | window_real = torch.view_as_real(window_freqs) |
| 183 | save_tensor(os.path.join(args.outdir, "rope_window_real"), window_real) |
| 184 | |
| 185 | # Global RoPE (72x72, scale_pos=24/72=1/3) |
| 186 | global_freqs = compute_axial_cis(dim=64, end_x=72, end_y=72, theta=10000.0, scale_pos=24.0/72.0) |
| 187 | global_real = torch.view_as_real(global_freqs) |
| 188 | save_tensor(os.path.join(args.outdir, "rope_global_real"), global_real) |
| 189 | |
| 190 | # Also dump the checkpoint's freqs_cis for block 0 (window) and block 7 (global) |
| 191 | for bi in [0, 7]: |
| 192 | key = f"blocks.{bi}.attn.freqs_cis" |
| 193 | if key in vit_sd: |
| 194 | fc = vit_sd[key] |
| 195 | fc_real = torch.view_as_real(fc) if fc.is_complex() else fc |
| 196 | save_tensor(os.path.join(args.outdir, f"ckpt_freqs_cis_block{bi}"), fc_real) |
| 197 | |
| 198 | # ── Step 3.9: Sinusoidal PE ───────────────────────────────────────── |
| 199 | print("\n=== Sinusoidal PE ===") |
| 200 | for S, name in [(288, "pe_288"), (144, "pe_144"), (72, "pe_72"), (36, "pe_36")]: |
| 201 | pe = sinusoidal_pe_2d(S, S, num_pos_feats=256) |
| 202 | save_tensor(os.path.join(args.outdir, name), pe) |
| 203 | |
| 204 | # ── Load and preprocess image ───────────────────────────────────── |
| 205 | print(f"\n=== Preprocessing: {args.image} ===") |
| 206 | img = Image.open(args.image).convert("RGB") |
| 207 | transform = v2.Compose([ |
| 208 | v2.ToDtype(torch.uint8, scale=True), |
| 209 | v2.Resize(size=(1008, 1008)), |
| 210 | v2.ToDtype(torch.float32, scale=True), |
no test coverage detected