From 52588b99102b3a2b16b97627dc8fcd07c4a528f1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=27=C3=A9lectron=20rare?= <108685187+electron-rare@users.noreply.github.com> Date: Wed, 13 May 2026 23:45:20 +0200 Subject: [PATCH] fix(coreml): roma branchless rotmat -> rotvec Root cause of v3d/transl NaN identified: roma.rotmat_to_rotvec uses torch.empty + 8 index_put_ on a buffer that CoreML mlprogram translates as scatter_nd over a garbage-initialised tensor. Cells that the scatter chain misses keep NaN; the subsequent quat / norm propagates NaN to every vertex. Patch: branchless atan2 formulation (stack/clamp/norm/atan2 only), no torch.empty, no index_put_. Precision drift vs roma original: 2.26e-6 L_inf on random batches. Mlpackage now validates all outputs finite (1.27e-4 L_inf vs PyTorch eager on v3d). Bench standalone: 65 ms median FP16 (15.3 fps, target met). Live with 3 parallel workers: 8 fps Multi-HMR keyframe rate (2.3x speedup vs PyTorch MPS baseline 3.5 fps); rigging still ships at 15-20 fps perceived. Output names shifted post-patch (var_2541 -> var_2412 etc) so multihmr_coreml.py constants updated. --- data_only_viz/multihmr_coreml.py | 10 +++---- data_only_viz/scripts/coreml_full_probe.py | 33 ++++++++++++++++++++-- 2 files changed, 35 insertions(+), 8 deletions(-) diff --git a/data_only_viz/multihmr_coreml.py b/data_only_viz/multihmr_coreml.py index f46dc4b..3992800 100644 --- a/data_only_viz/multihmr_coreml.py +++ b/data_only_viz/multihmr_coreml.py @@ -39,11 +39,11 @@ N_PERSONS_FIXED = 4 N_VERTS = 10475 # CoreML output names from the exported .mlpackage. -OUT_V3D = "var_2541" # (4, 10475, 3) f16 -OUT_TRANSL = "var_2544" # (4, 1, 3) f16 -OUT_SCORES = "var_2557" # (4,) f16 -OUT_BETAS = "var_2560" # (4, 10) f16 -OUT_EXPR = "var_2563" # (4, 10) f16 +OUT_V3D = "var_2412" # (4, 10475, 3) +OUT_TRANSL = "var_2415" # (4, 1, 3) +OUT_SCORES = "var_2428" # (4,) +OUT_BETAS = "var_2431" # (4, 10) +OUT_EXPR = "var_2434" # (4, 10) # MLMultiArrayDataType raw values (from CoreML headers). ML_DTYPE_FLOAT32 = 65568 diff --git a/data_only_viz/scripts/coreml_full_probe.py b/data_only_viz/scripts/coreml_full_probe.py index 539f734..f1f4387 100644 --- a/data_only_viz/scripts/coreml_full_probe.py +++ b/data_only_viz/scripts/coreml_full_probe.py @@ -157,6 +157,32 @@ if hasattr(model.backbone, "encoder") and hasattr(model.backbone.encoder, # torch.inverse(K) plante coremltools (op non implementee). Comme K est # fixe (camera intrinsics avec focal=IMG_SIZE), on pre-calcule K_inv # en closed-form et on l'utilise comme buffer module-level. +print("==> Patching roma.rotmat_to_rotvec (branchless atan2)") +# roma.rotmat_to_rotvec utilise torch.empty + 8 index_put_ qui se +# traduisent en CoreML par scatter_nd successifs sur un buffer +# garbage-initialise. Resultat : cellules non touchees restent NaN, +# propagees via quat normalization -> v3d/transl all-NaN. +# Remplacement branchless via atan2 : pas de torch.empty, pas +# d'index_put_, juste des stack/clamp/norm/atan2 stables CoreML. +# Precision vs roma original : 2.26e-6 L_inf sur batch random. +import roma as _roma + +def _rotmat_to_rotvec_branchless(R, eps=1e-6): + w = torch.stack([ + R[..., 2, 1] - R[..., 1, 2], + R[..., 0, 2] - R[..., 2, 0], + R[..., 1, 0] - R[..., 0, 1], + ], dim=-1) * 0.5 + trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2] + cos_theta = ((trace - 1.0) * 0.5).clamp(-1.0, 1.0) + sin_theta = torch.norm(w, dim=-1) + theta = torch.atan2(sin_theta, cos_theta) + sin_theta_safe = sin_theta.clamp(min=eps) + return w * (theta / sin_theta_safe).unsqueeze(-1) + +_roma.rotmat_to_rotvec = _rotmat_to_rotvec_branchless + + print("==> Patching utils.camera.inverse_perspective_projection") import utils.camera as _camera @@ -498,9 +524,10 @@ try: compute_units=ct.ComputeUnit.CPU_AND_GPU, minimum_deployment_target=ct.target.macOS15, convert_to="mlprogram", - # FP16 default causes NaN in inverse projection / SMPL-X decoder - # (Multi-HMR has values that overflow the FP16 range). Force FP32. - compute_precision=ct.precision.FLOAT32, + # FP16 OK depuis le patch roma branchless (cf rapport bisection + # 2026-05-13) : la source du NaN etait torch.empty + index_put_ + # dans roma.rotmat_to_rotvec, pas la precision. + compute_precision=ct.precision.FLOAT16, ) out_path = "/tmp/multihmr_full_672_s.mlpackage" mlmodel.save(out_path)