fix(models): guard keypoint loss and its shared ref_wh composition against non-finite predictions - #1336

What does this PR do?

compute_l1_keypoint_loss has four loss terms. Three of them (location_loss, findable_loss, visible_loss) excluded invalid keypoints by multiplying the per-keypoint loss by a float mask. Under IEEE 754, 0.0 * nan == nan, not 0.0, so a masked-out position with a non-finite prediction still poisoned the sum. location_loss had an isfinite check that looked like it should catch this, but the multiply defeated it the same way.

Switching to masked_fill (mirroring the pattern this function already used correctly for nll_keypoints) fixes the forward value, but that alone isn't always enough for the gradient. Every loss op's own backward formula still evaluates its original input locally. For binary_cross_entropy_with_logits that's sigmoid(x) - y, so a non-finite x gives a non-finite local gradient regardless of the output mask, and the correctly-zeroed upstream gradient times that non-finite local gradient is 0.0 * nan == nan again, one level down — a single non-finite findable/visible prediction still poisoned every parameter gradient reachable from the keypoint head, even though the loss value came out finite. (F.l1_loss's local gradient is sign(pred - target), and torch.sign(nan) happens to return 0.0 in the PyTorch version this was tested on, so location_loss's own gradient was already clean with masked_fill alone — verified directly, not assumed. Applied the same input-guard there anyway for consistency across all four terms rather than relying on that implementation detail holding across PyTorch versions/devices.)

The Gaussian NLL term has the same hazard through dx/dy/the Cholesky params: it already called nan_to_num on its final nll_raw, but that only cleans the exact node it's applied to, not the chain of ops (u0, u1, maha2) leading up to it, where the non-finite value re-enters the backward formula the same way.

Fixed by swapping in a safe placeholder (torch.where) for the non-finite input before each loss op, on top of the existing masked_fill/nan_to_num on the output — verified with a direct backward pass: 7 prediction channels × 4 loss terms, 28 combinations, all give a finite gradient with the fix and (for the combinations that actually break) a non-finite one without it.

Three more gaps in the same spirit, found and fixed in review:

• The torch.where guards above only protect ops inside compute_l1_keypoint_loss. In the real GroupPose training path (use_grouppose_keypoints=True), lwdetr.py builds pred_keypoints[..., :2] as outputs_keypoints_delta[..., :2] * ref_wh + ref_xy before calling this function — ref_wh/ref_xy come from ref_unsigmoid, shared with the box head. A non-finite outputs_keypoints_delta still poisons ref_wh's gradient the same 0.0 * nan == nan way, one level earlier, because that multiply's own backward (d(a*b)/db == a) evaluates the original (non-finite) a regardless of how clean the gradient flowing out of this function is. Fixed by nan_to_num-ing outputs_keypoints_delta at the source in lwdetr.py, before the multiply — verified with a full LWDETR.forward() pass (mocked backbone/transformer, real keypoint_embed) injecting a NaN and checking ref_unsigmoid.grad.
• A finite log-precision input can still overflow one op later. finite_uncertainty only checks the raw (pre-exp) Cholesky params, so e.g. raw_log_l11 = 100.0 passes it, then exp(100) overflows to inf — several ops before the existing isfinite(u0)/isfinite(maha2) check that excludes it from the forward sum. By then u0/maha2 are already built from the overflowed value, so their own backward (d(u0**2)/d(u0) == 2*u0 == inf) poisons x/y/Cholesky gradients the same way. Fixed by refining the mask after computing u0/u1/maha2 once, then redoing that chain with torch.where-sanitized inputs before exp/multiply/square see them.
• A non-finite target_areas still produced nan in location_loss/nll_loss. valid_area correctly drops those keypoints from location_loss_mask/gaussian_loss_mask, and masked_fill correctly zeroes their numerator — but both terms then divide by area.clamp_min(area_eps) directly, and torch.clamp_min(nan, eps) leaves nan unchanged (comparisons against nan are always False). 0.0 / nan == nan survives the mask. Fixed by sanitizing area itself (torch.where(valid_area, area, 1.0)) before either division.

All three follow TDD (failing test confirmed against the pre-fix code via git stash, passing after).

Related Issue(s): None, no existing report found.

Type of Change

• Bug fix (non-breaking change that fixes an issue)

Testing

• I have tested this change locally
• I have added/updated tests for this change

Test details:

tests/models/heads/test_keypoint_heads.py: two parametrized tests from the first pass (forward value, 3 cases; gradient, 7 cases — one per prediction channel, summing all four losses and calling backward()), plus two new ones for the gaps above (exp() overflow of a finite log-precision input; non-finite target_areas). All four written/confirmed failing against the pre-fix code first (git stash). Full file now 30/30 (18 pre-existing + 12 new: 3 forward-value cases + 7 gradient cases + the exp()-overflow case + the non-finite-target_areas case), including the pre-existing NLL-gradient-vs-reference-formula test, confirming no behaviour change in the normal finite case. tests/models/test_criterion.py (12/12, same package, not a direct consumer) also passes unchanged.

tests/models/test_lwdetr_keypoints.py: one new test building a real LWDETR (mocked backbone/transformer, real keypoint_embed), injecting a NaN via a forward hook on keypoint_embed, and checking ref_unsigmoid.grad after backward(). Confirmed failing against the pre-fix lwdetr.py (git stash), passing with the fix. Full file 7/7.

Broader sanity check after all three fixes (not the full 3854-item suite — no need to re-pay that cost for an unrelated part of the tree): tests/models/ + tests/models/heads/ together, 194 passed, 3 skipped (gpu-marked).

Full CPU suite per ci-tests-cpu.yml, from the prior review pass: ran to completion (3780 passed + 72 skipped + 2 failed = 3854 items, single worker, 188s). Both failures are in tests/models/test_evaluate.py, a file that doesn't import or exercise keypoints.py. Re-ran both in isolation; the first failed again from an external SIGTERM — confirmed against journalctl -u claude-watchdog.service (machine-level systemd log, independent of this session): a real entry at 22:43:32 kills pytest tests/models/test_evaluate.py::test_train_then_from_checkpoint_then_evaluate at cpu=2281%, matching the number cited in repro/NOTAS-verificacion.md exactly. The second failure (test_sequential_evaluate_calls_do_not_leak_state) has no equivalent isolated confirmation — it only ran as part of the original full-suite pass.

Checklist

• My code follows the style guidelines of this project
• I have performed a self-review of my own code
• I have commented my code where necessary, particularly in hard-to-understand areas
• My changes generate no new warnings or errors

Not applicable for docs: no public API changed.