(
ckpt: Dict[str, torch.Tensor],
img_feats_hwc: torch.Tensor, # [H*W, B, C] seq-first image features
img_feats_nchw: torch.Tensor, # [B, C, H, W] image features for ROI align
img_pe_hwc: torch.Tensor, # [H*W, B, C] image positional encoding
boxes_cxcywh: torch.Tensor, # [N_boxes, B, 4] normalized CxCyWH [0,1]
box_labels: torch.Tensor, # [N_boxes, B] long (0=pos, 1=neg)
box_mask: torch.Tensor, # [B, N_boxes] bool (True=padded)
outdir: str,
case_name: str,
)
| 154 | # ── Main dump logic ─────────────────────────────────────────────────────── |
| 155 | |
| 156 | def dump_geometry_encoder( |
| 157 | ckpt: Dict[str, torch.Tensor], |
| 158 | img_feats_hwc: torch.Tensor, # [H*W, B, C] seq-first image features |
| 159 | img_feats_nchw: torch.Tensor, # [B, C, H, W] image features for ROI align |
| 160 | img_pe_hwc: torch.Tensor, # [H*W, B, C] image positional encoding |
| 161 | boxes_cxcywh: torch.Tensor, # [N_boxes, B, 4] normalized CxCyWH [0,1] |
| 162 | box_labels: torch.Tensor, # [N_boxes, B] long (0=pos, 1=neg) |
| 163 | box_mask: torch.Tensor, # [B, N_boxes] bool (True=padded) |
| 164 | outdir: str, |
| 165 | case_name: str, |
| 166 | ): |
| 167 | D = 256 |
| 168 | roi_size = 7 |
| 169 | case_dir = os.path.join(outdir, case_name) |
| 170 | os.makedirs(case_dir, exist_ok=True) |
| 171 | |
| 172 | n_boxes, bs = boxes_cxcywh.shape[:2] |
| 173 | print(f"\n=== Geometry Encoder: {case_name} ===") |
| 174 | print(f" boxes: {n_boxes}, batch: {bs}, img_feats: {img_feats_hwc.shape}") |
| 175 | |
| 176 | # Save input coordinates |
| 177 | save_raw(os.path.join(case_dir, "input_boxes_cxcywh"), |
| 178 | boxes_cxcywh.detach().cpu().float().numpy(), |
| 179 | boxes_cxcywh.shape) |
| 180 | save_raw(os.path.join(case_dir, "input_box_labels"), |
| 181 | box_labels.detach().cpu().float().numpy(), |
| 182 | box_labels.shape) |
| 183 | |
| 184 | # ── 1. Direct box projection: Linear(4, 256) ─────────────────────── |
| 185 | box_proj_w = ckpt["detector.geometry_encoder.boxes_direct_project.weight"].float() |
| 186 | box_proj_b = ckpt["detector.geometry_encoder.boxes_direct_project.bias"].float() |
| 187 | boxes_direct = F.linear(boxes_cxcywh, box_proj_w, box_proj_b) # [N, B, D] |
| 188 | save_ggml_sbd(os.path.join(case_dir, "boxes_direct_proj"), boxes_direct) |
| 189 | print(f" boxes_direct_proj: {boxes_direct.shape} " |
| 190 | f"mean={boxes_direct.mean():.6f} std={boxes_direct.std():.6f}") |
| 191 | |
| 192 | # ── 2. ROI Align pooled features ──────────────────────────────────── |
| 193 | # img_pre_norm on image features before pooling |
| 194 | img_pre_norm_w = ckpt["detector.geometry_encoder.img_pre_norm.weight"].float() |
| 195 | img_pre_norm_b = ckpt["detector.geometry_encoder.img_pre_norm.bias"].float() |
| 196 | # Apply LayerNorm to seq-first features for pooling (operate on the NCHW version) |
| 197 | img_normed_hwc = F.layer_norm(img_feats_hwc, [D], img_pre_norm_w, img_pre_norm_b) |
| 198 | H, W = img_feats_nchw.shape[-2:] |
| 199 | img_normed_nchw = img_normed_hwc.permute(1, 2, 0).reshape(bs, D, H, W) |
| 200 | |
| 201 | if n_boxes > 0: |
| 202 | roi_pooled = do_roi_align(img_normed_nchw, boxes_cxcywh, roi_size) |
| 203 | # roi_pooled: [B*N_boxes, D, roi_size, roi_size] |
| 204 | |
| 205 | # boxes_pool_project is Conv2d(D, D, roi_size) |
| 206 | pool_proj_w = ckpt["detector.geometry_encoder.boxes_pool_project.weight"].float() |
| 207 | pool_proj_b = ckpt["detector.geometry_encoder.boxes_pool_project.bias"].float() |
| 208 | pool_proj = F.conv2d(roi_pooled, pool_proj_w, pool_proj_b) # [B*N, D, 1, 1] |
| 209 | pool_proj = pool_proj.view(bs, n_boxes, D).transpose(0, 1) # [N, B, D] |
| 210 | save_ggml_sbd(os.path.join(case_dir, "boxes_pool_proj"), pool_proj) |
| 211 | print(f" boxes_pool_proj: {pool_proj.shape} " |
| 212 | f"mean={pool_proj.mean():.6f} std={pool_proj.std():.6f}") |
| 213 | else: |
no test coverage detected