Spatially-grounded depth head (opt-in) + depth-range guards (I1-B/I1-D/I1-G)

[sibocw]

This work was implemented by Claude Code (Opus 4.8) upon user instruction.

Rationale

Audit #48 found 3D keypoints performing worse than body-seg. The root cause (finding I1-B) is the depth head: _build_depth_head does AdaptiveAvgPool2d(1) on the (N, 64, 128, 128) decoder map BEFORE its conv, collapsing it to a single global descriptor. Every keypoint's depth is then regressed from that one global vector with no spatial localization — whereas the x-y heatmap head is fully spatial. Depth is the asymmetric weak link, so this is the primary driver of weak 3D.

Changes

I1-B — spatially-grounded depth head (opt-in)

• New ModelArchitectureConfig.depth_head_type field:
  • "global" (default) — the existing global-pool head; unchanged behavior and checkpoint-compatible.
  • "spatial" — the new head. After the model computes the heatmaps and the soft-argmax keypoint locations xy_px_out (heatmap-pixel coords on the same 128x128 grid as d0), it bilinearly samples d0 at each keypoint's predicted (x, y) via F.grid_sample (coords normalized to [-1, 1], align_corners=False), giving per-keypoint feature vectors (N, n_kp, C). A shared small MLP (Linear -> GroupNorm -> ReLU -> Linear) maps C -> depth_n_bins per keypoint. This grounds each keypoint's depth at its own image location, mirroring the spatial x-y head.
• The head is stored on self.depth_head regardless of type, so the training pipeline's optimizer LR group (model.depth_head.parameters()) and the AMP status check pick it up unchanged — no pipeline.py edit was needed. I verified the pipeline's n_params_optimizer == n_params_model assertion holds for both head types.
• The new head's final layer uses the existing small-std init convention (pose_head_init_std); the hidden layer uses Kaiming (it precedes ReLU).

I1-G — fix disjoint depth-range default + runtime guard

• The old default depth window [-70, -63] is disjoint from the measured label range (~[-102, -99] mm). A run that forgets to override it makes every keypoint OOB and, under oob_treatment="drop", the loss returns 0 (trains nothing). Production configs override it, but it's a landmine. New default: depth_min=-103.0, depth_max=-99.0 (overlaps the data), with comments that it MUST match the labels.
• Added an ERROR-level log in Pose2p5DLoss when >50% of a batch has OOB depth, telling the user the depth window is likely misconfigured. This surfaces I1-G/I1-D at runtime.

I1-D — deferred

Per-keypoint OOB masking (a new "mask" oob_treatment) was not added: whole-sample drop (.any(dim=1)) only loses ~0.18% in practice (not a starvation problem), and adding it cleanly would require threading per-keypoint masks through both depth loss reductions (_compute_depth_ce_loss, _compute_depth_l1_loss), which risks changing existing loss magnitudes. Noted as future work. Existing drop/clamp/ignore behavior is unchanged.

Checkpoint compatibility (important)

• Existing models (global head) load fine with the default depth_head_type="global".
• Switching to "spatial" requires training a new model — the depth-head state_dict keys differ between the two heads (depth_head.0.weight etc. vs depth_head.fc1.weight etc.), so a global checkpoint cannot load into a spatial model (and vice versa). The encoder / decoder / heatmap head are byte-for-byte unaffected.

Tests

New tests/test_depth_head.py (CPU, tiny config; shims the missing private pvio symbol and pins poseforge to this worktree's src before importing):

• (a) both "global" and "spatial" forward on a random batch -> depth_logits of shape (N, n_kp, depth_n_bins), finite;
• (b) grounding: the spatial head's depth output for a keypoint changes when its (x, y) moves across a spatially-varying feature map whose global average is exactly zero (so a global-pool head would be invariant) — plus a full-model variant where moving the heatmap peak changes the depth logits;
• (c) gradients flow (finite, non-zero) to every spatial-head parameter;
• plus a check that the two configs build genuinely different head modules / state-dict keys.

$ .venv/bin/python -m pytest tests/test_depth_head.py -q
.....                                                                    [100%]
5 passed in 4.11s

Refs #48

🤖 Generated with Claude Code