perf(data-only-viz): CoreML T2 done T3 partial
Task 2 — apply_topk(K=4) validee comme drop-in remplacant apply_threshold dans la tete Multi-HMR. Cosine sim = 1.000000 sur 4/4 detections image reelle, MAE 2-4 mm. Script : scripts/probe_head_topk.py. Task 3 — full Multi-HMR convert tente. Patches appliques : 1. apply_threshold -> apply_topk monkey-patch 2. backbone.encoder.interpolate_pos_encoding -> buffer fige 3. utils.camera.inverse_perspective_projection -> closed-form jit.trace OK, mais conversion coremltools echoue sur op suivantes en cascade. Estimation restante : 1-2 j. Script : scripts/coreml_full_probe.py. Le breakthrough probe v4 backbone seul 11.8x reste la victoire. Task 4 worker integration bloquee tant que T3 pas complete.
This commit is contained in:
L'électron rare
@@ -0,0 +1,251 @@
|
||||
"""Task 3 — Convert FULL Multi-HMR (backbone + head) to CoreML
|
||||
avec apply_topk(K=4) + fixed-shape tuple output.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
CACHE = Path.home() / ".cache" / "av-live-multihmr"
|
||||
CKPT = CACHE / "checkpoints" / "multiHMR_672_S.pt"
|
||||
MULTIHMR_REPO = CACHE / "multi-hmr"
|
||||
|
||||
sys.path.insert(0, str(MULTIHMR_REPO))
|
||||
for mod in ("pyrender", "pyvista", "anny"):
|
||||
sys.modules.setdefault(mod, types.ModuleType(mod))
|
||||
|
||||
DEVICE = "cpu" # trace on CPU (CoreML mlprogram doesn't care)
|
||||
IMG_SIZE = 672
|
||||
K_PERSONS = 4
|
||||
|
||||
|
||||
# === apply_topk replacement (validated equivalent in Task 2) ===
|
||||
def apply_topk(K, _scores):
|
||||
if isinstance(K, list):
|
||||
K = K[0]
|
||||
B, H, W, C = _scores.shape
|
||||
flat = _scores.reshape(B, -1)
|
||||
_, idx_flat = torch.topk(flat, k=K, dim=1)
|
||||
wc = W * C
|
||||
idx_b = (torch.arange(B, device=_scores.device)
|
||||
.unsqueeze(1).expand(-1, K).reshape(-1).long())
|
||||
idx_flat_flat = idx_flat.reshape(-1)
|
||||
idx_h = (idx_flat_flat // wc).long()
|
||||
idx_w = ((idx_flat_flat // C) % W).long()
|
||||
idx_c = (idx_flat_flat % C).long()
|
||||
return (idx_b, idx_h, idx_w, idx_c)
|
||||
|
||||
|
||||
# === Patch coremltools _cast (validated probe v4) ===
|
||||
def _patched_cast(context, node, dtype, dtype_str):
|
||||
from coremltools.converters.mil import Builder as mb
|
||||
from coremltools.converters.mil.frontend.torch import ops as _ops
|
||||
inputs = _ops._get_inputs(context, node, expected=1)
|
||||
x = inputs[0]
|
||||
if x.val is not None:
|
||||
try:
|
||||
const_val = dtype(x.val)
|
||||
except TypeError:
|
||||
arr = np.asarray(x.val)
|
||||
if arr.size == 1:
|
||||
const_val = dtype(arr.item())
|
||||
else:
|
||||
res = mb.cast(x=x, dtype=dtype_str, name=node.name)
|
||||
context.add(res)
|
||||
return
|
||||
res = mb.const(val=const_val, name=node.name)
|
||||
else:
|
||||
res = mb.cast(x=x, dtype=dtype_str, name=node.name)
|
||||
context.add(res)
|
||||
|
||||
|
||||
prev = os.getcwd()
|
||||
try:
|
||||
os.chdir(MULTIHMR_REPO)
|
||||
from model import Model
|
||||
import model as model_mod
|
||||
|
||||
# Inject topk replacement
|
||||
print("==> Patching apply_threshold -> apply_topk(K=4)")
|
||||
model_mod.apply_threshold = lambda thr, scores: apply_topk(K_PERSONS, scores)
|
||||
|
||||
torch_dev = torch.device(DEVICE)
|
||||
ckpt = torch.load(str(CKPT), map_location=torch_dev, weights_only=False)
|
||||
kw = {k: v for k, v in vars(ckpt["args"]).items()}
|
||||
kw["type"] = ckpt["args"].train_return_type
|
||||
kw["img_size"] = ckpt["args"].img_size[0]
|
||||
print(f"==> Loading Multi-HMR ViT-S 672 (params count tbd)")
|
||||
model = Model(**kw).to(torch_dev)
|
||||
model.load_state_dict(ckpt["model_state_dict"], strict=False)
|
||||
model.eval()
|
||||
finally:
|
||||
os.chdir(prev)
|
||||
|
||||
|
||||
# === Pre-compute interpolate_pos_encoding (probe v4 fix) ===
|
||||
# Multi-HMR's backbone is DINOv2 ViT-S/14 — same dynamic interpolation
|
||||
# problem that planted la conversion sur backbone seul.
|
||||
if hasattr(model.backbone, "encoder") and hasattr(model.backbone.encoder,
|
||||
"interpolate_pos_encoding"):
|
||||
print("==> Patching backbone.encoder.interpolate_pos_encoding (pre-compute)")
|
||||
bk = model.backbone.encoder
|
||||
with torch.no_grad():
|
||||
dummy_x = torch.rand(1, 3, IMG_SIZE, IMG_SIZE)
|
||||
dummy_p = bk.patch_embed(dummy_x)
|
||||
cls = bk.cls_token.expand(dummy_p.shape[0], -1, -1)
|
||||
x_full = torch.cat((cls, dummy_p), dim=1)
|
||||
cached_pe = bk.interpolate_pos_encoding(
|
||||
x_full, IMG_SIZE, IMG_SIZE).detach()
|
||||
bk.register_buffer("_cached_pos_embed", cached_pe)
|
||||
|
||||
def fixed_pe(self, x, w, h):
|
||||
return self._cached_pos_embed.to(x.dtype)
|
||||
bk.interpolate_pos_encoding = types.MethodType(fixed_pe, bk)
|
||||
print(f" cached shape {tuple(cached_pe.shape)}")
|
||||
|
||||
# === Patch utils.camera.inverse_perspective_projection ===
|
||||
# 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 utils.camera.inverse_perspective_projection")
|
||||
import utils.camera as _camera
|
||||
|
||||
# Pre-compute K_inv closed-form pour notre K standard
|
||||
focal_val = float(IMG_SIZE)
|
||||
cx = cy = IMG_SIZE / 2.0
|
||||
_K_INV_PRE = torch.tensor([
|
||||
[[1.0 / focal_val, 0.0, -cx / focal_val],
|
||||
[0.0, 1.0 / focal_val, -cy / focal_val],
|
||||
[0.0, 0.0, 1.0]]
|
||||
])
|
||||
|
||||
def inverse_perspective_projection_fixed(points, K, distance):
|
||||
"""Bypass torch.inverse : utilise K_inv pre-calcule en closed-form
|
||||
(notre K est connu et fixe). Le K argument est ignore."""
|
||||
K_inv = _K_INV_PRE.to(points.device).to(points.dtype)
|
||||
points = torch.cat([points, torch.ones_like(points[..., :1])], -1)
|
||||
points = torch.einsum('bij,bkj->bki', K_inv, points)
|
||||
if distance is None:
|
||||
return points
|
||||
points = points * distance
|
||||
return points
|
||||
|
||||
_camera.inverse_perspective_projection = inverse_perspective_projection_fixed
|
||||
# Aussi patcher le re-export dans utils/__init__.py et model.py
|
||||
import utils as _utils_pkg
|
||||
_utils_pkg.inverse_perspective_projection = inverse_perspective_projection_fixed
|
||||
# model.py importe directement : monkey-patch sur le module
|
||||
model_mod.inverse_perspective_projection = inverse_perspective_projection_fixed
|
||||
# Idem smpl_layer
|
||||
import blocks.smpl_layer as _smpl_layer
|
||||
_smpl_layer.inverse_perspective_projection = inverse_perspective_projection_fixed
|
||||
|
||||
|
||||
# === Wrapper qui produit tuple fixe ===
|
||||
class TracedMHMR(nn.Module):
|
||||
"""Wrap Multi-HMR pour trace : output tuple de tensors fixes,
|
||||
pas de list-of-dicts."""
|
||||
|
||||
def __init__(self, m: Model):
|
||||
super().__init__()
|
||||
self.m = m
|
||||
|
||||
def forward(self, x: torch.Tensor, cam_K: torch.Tensor):
|
||||
# Call original forward with is_training=False ; apply_topk
|
||||
# garantit toujours K=4 detections donc le loop dans le forward
|
||||
# est unroll-friendly.
|
||||
humans = self.m(x, is_training=False, nms_kernel_size=5,
|
||||
det_thresh=0.0, K=cam_K)
|
||||
# humans est une list[dict] de longueur 4. Stack en tensors.
|
||||
if len(humans) == 0:
|
||||
# Should not happen with apply_topk, but defensive
|
||||
zeros = torch.zeros(K_PERSONS, 10475, 3)
|
||||
zeros_p = torch.zeros(K_PERSONS, 3)
|
||||
zeros_s = torch.zeros(K_PERSONS)
|
||||
zeros_b = torch.zeros(K_PERSONS, 10)
|
||||
return zeros, zeros_p, zeros_s, zeros_b, zeros_b
|
||||
v3d = torch.stack([h["v3d"] for h in humans])
|
||||
transl = torch.stack([h["transl_pelvis"] for h in humans])
|
||||
scores = torch.stack([
|
||||
h["scores"] if h["scores"].dim() > 0 else h["scores"].unsqueeze(0)
|
||||
for h in humans
|
||||
]).squeeze(-1)
|
||||
shape = torch.stack([h["shape"] for h in humans])
|
||||
expr = torch.stack([h["expression"] for h in humans])
|
||||
return v3d, transl, scores, shape, expr
|
||||
|
||||
|
||||
wrapper = TracedMHMR(model).eval()
|
||||
|
||||
# Sanity forward
|
||||
focal = float(IMG_SIZE)
|
||||
example_K = torch.tensor(
|
||||
[[[focal, 0.0, IMG_SIZE / 2.0],
|
||||
[0.0, focal, IMG_SIZE / 2.0],
|
||||
[0.0, 0.0, 1.0]]], dtype=torch.float32)
|
||||
example_x = torch.rand(1, 3, IMG_SIZE, IMG_SIZE)
|
||||
|
||||
print("==> Sanity forward")
|
||||
with torch.no_grad():
|
||||
v3d, transl, scores, shape, expr = wrapper(example_x, example_K)
|
||||
print(f" v3d: {tuple(v3d.shape)}, transl: {tuple(transl.shape)},")
|
||||
print(f" scores: {tuple(scores.shape)}, shape: {tuple(shape.shape)},")
|
||||
print(f" expr: {tuple(expr.shape)}")
|
||||
|
||||
print("==> torch.jit.trace")
|
||||
try:
|
||||
traced = torch.jit.trace(wrapper, (example_x, example_K), strict=False)
|
||||
print(" trace OK")
|
||||
except Exception as e:
|
||||
print(f" trace FAILED: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
# === CoreML convert ===
|
||||
print("==> coremltools.convert")
|
||||
import coremltools as ct
|
||||
from coremltools.converters.mil.frontend.torch import ops as _ops
|
||||
_ops._cast = _patched_cast
|
||||
|
||||
try:
|
||||
mlmodel = ct.convert(
|
||||
traced,
|
||||
inputs=[
|
||||
ct.TensorType(shape=(1, 3, IMG_SIZE, IMG_SIZE),
|
||||
name="image", dtype=np.float32),
|
||||
ct.TensorType(shape=(1, 3, 3), name="cam_K", dtype=np.float32),
|
||||
],
|
||||
compute_units=ct.ComputeUnit.CPU_AND_GPU,
|
||||
minimum_deployment_target=ct.target.macOS15,
|
||||
convert_to="mlprogram",
|
||||
)
|
||||
out_path = "/tmp/multihmr_full_672_s.mlpackage"
|
||||
mlmodel.save(out_path)
|
||||
print(f" CONVERT OK -> {out_path}")
|
||||
except Exception as e:
|
||||
print(f" CONVERT FAILED: {type(e).__name__}: {e}")
|
||||
raise
|
||||
|
||||
# === Bench ===
|
||||
print("==> bench 30 iter")
|
||||
img = np.random.rand(1, 3, IMG_SIZE, IMG_SIZE).astype(np.float32)
|
||||
cam = np.array([[[focal, 0, IMG_SIZE/2],
|
||||
[0, focal, IMG_SIZE/2],
|
||||
[0, 0, 1]]], dtype=np.float32)
|
||||
for _ in range(3):
|
||||
_ = mlmodel.predict({"image": img, "cam_K": cam})
|
||||
t = []
|
||||
for _ in range(30):
|
||||
t0 = time.perf_counter()
|
||||
_ = mlmodel.predict({"image": img, "cam_K": cam})
|
||||
t.append((time.perf_counter() - t0) * 1000)
|
||||
t.sort()
|
||||
print(f" CoreML full Multi-HMR median={t[15]:.1f} ms "
|
||||
f"p10={t[3]:.1f} p90={t[27]:.1f} min={t[0]:.1f}")
|
||||
print(f" Target was <60ms (12-25 fps). Achieved: {1000.0/t[15]:.1f} fps")
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Task 2 — Validate apply_topk(K=4) as drop-in replacement for
|
||||
apply_threshold in Multi-HMR head.
|
||||
|
||||
Compares v3d output between threshold-based (original) and topk-based
|
||||
(patched) Multi-HMR on the same input. Pass criterion: for the same
|
||||
detections (when K >= n_threshold_detected), v3d cosine similarity > 0.99.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
CACHE = Path.home() / ".cache" / "av-live-multihmr"
|
||||
CKPT = CACHE / "checkpoints" / "multiHMR_672_S.pt"
|
||||
MULTIHMR_REPO = CACHE / "multi-hmr"
|
||||
|
||||
sys.path.insert(0, str(MULTIHMR_REPO))
|
||||
for mod in ("pyrender", "pyvista", "anny"):
|
||||
sys.modules.setdefault(mod, types.ModuleType(mod))
|
||||
|
||||
DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
IMG_SIZE = 672
|
||||
|
||||
|
||||
def apply_topk(K, _scores):
|
||||
"""Drop-in pour apply_threshold. _scores shape (B, H, W, C).
|
||||
Renvoie 4-tuple LongTensor (batch_idx, h_idx, w_idx, c_idx) chacun
|
||||
de longueur B*K (au lieu de variable). K candidate top-scoring
|
||||
tokens par image.
|
||||
"""
|
||||
if isinstance(K, list):
|
||||
K = K[0]
|
||||
B, H, W, C = _scores.shape
|
||||
flat = _scores.reshape(B, -1)
|
||||
_, idx_flat = torch.topk(flat, k=K, dim=1)
|
||||
wc = W * C
|
||||
idx_b = (torch.arange(B, device=_scores.device)
|
||||
.unsqueeze(1).expand(-1, K).reshape(-1))
|
||||
idx_flat_flat = idx_flat.reshape(-1)
|
||||
idx_h = idx_flat_flat // wc
|
||||
idx_w = (idx_flat_flat // C) % W
|
||||
idx_c = idx_flat_flat % C
|
||||
return (idx_b.long(), idx_h.long(), idx_w.long(), idx_c.long())
|
||||
|
||||
|
||||
prev = os.getcwd()
|
||||
try:
|
||||
os.chdir(MULTIHMR_REPO)
|
||||
from model import Model
|
||||
import model as model_mod
|
||||
torch_dev = torch.device(DEVICE)
|
||||
ckpt = torch.load(str(CKPT), map_location=torch_dev, weights_only=False)
|
||||
kw = {k: v for k, v in vars(ckpt["args"]).items()}
|
||||
kw["type"] = ckpt["args"].train_return_type
|
||||
kw["img_size"] = ckpt["args"].img_size[0]
|
||||
model = Model(**kw).to(torch_dev)
|
||||
model.load_state_dict(ckpt["model_state_dict"], strict=False)
|
||||
model.eval()
|
||||
finally:
|
||||
os.chdir(prev)
|
||||
|
||||
focal = float(IMG_SIZE)
|
||||
K_mat = torch.tensor([[[focal, 0.0, IMG_SIZE / 2.0],
|
||||
[0.0, focal, IMG_SIZE / 2.0],
|
||||
[0.0, 0.0, 1.0]]], device=DEVICE)
|
||||
|
||||
# Use a real test image (multi-hmr example or webcam capture)
|
||||
import cv2
|
||||
img_path = "/Users/electron/.cache/av-live-multihmr/multi-hmr/example_data/4446582661_b188f82f3c_c.jpg"
|
||||
img = cv2.imread(img_path)
|
||||
h, w = img.shape[:2]
|
||||
side = min(h, w)
|
||||
y0 = (h - side) // 2; x0 = (w - side) // 2
|
||||
img = cv2.resize(img[y0:y0+side, x0:x0+side], (IMG_SIZE, IMG_SIZE))
|
||||
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
x = (torch.from_numpy(img_rgb).permute(2, 0, 1).float() / 255.0
|
||||
).unsqueeze(0).to(DEVICE)
|
||||
print(f"loaded image {img_path}")
|
||||
|
||||
# --- Pass 1 : original apply_threshold a tres bas seuil ---
|
||||
print("==> Pass 1 : apply_threshold(0.15)")
|
||||
with torch.no_grad():
|
||||
humans_orig = model(x, is_training=False, nms_kernel_size=5,
|
||||
det_thresh=0.15, K=K_mat)
|
||||
print(f" detected: {len(humans_orig)}")
|
||||
for i, h in enumerate(humans_orig[:4]):
|
||||
sc = h.get("scores", 0.0)
|
||||
if hasattr(sc, "item"):
|
||||
sc = sc.item()
|
||||
print(f" [{i}] score={sc:.3f} v3d.shape={tuple(h['v3d'].shape)}")
|
||||
|
||||
# --- Pass 2 : monkey-patch apply_threshold avec apply_topk(K=4) ---
|
||||
print("\n==> Pass 2 : apply_topk(K=4)")
|
||||
original_apply_threshold = model_mod.apply_threshold
|
||||
|
||||
def topk_wrapper(det_thresh, _scores):
|
||||
return apply_topk(4, _scores)
|
||||
|
||||
model_mod.apply_threshold = topk_wrapper
|
||||
|
||||
with torch.no_grad():
|
||||
humans_topk = model(x, is_training=False, nms_kernel_size=5,
|
||||
det_thresh=0.15, K=K_mat)
|
||||
print(f" detected: {len(humans_topk)}")
|
||||
for i, h in enumerate(humans_topk[:4]):
|
||||
sc = h.get("scores", 0.0)
|
||||
if hasattr(sc, "item"):
|
||||
sc = sc.item()
|
||||
print(f" [{i}] score={sc:.3f} v3d.shape={tuple(h['v3d'].shape)}")
|
||||
|
||||
# Restore
|
||||
model_mod.apply_threshold = original_apply_threshold
|
||||
|
||||
# --- Comparison ---
|
||||
print("\n==> Comparison")
|
||||
if len(humans_orig) == 0 or len(humans_topk) == 0:
|
||||
print(" NO DETECTIONS in one path — adjust threshold lower")
|
||||
sys.exit(0)
|
||||
|
||||
# Match by score (highest first in both)
|
||||
o = sorted(humans_orig, key=lambda h: -(
|
||||
h.get("scores", 0).item() if hasattr(h.get("scores", 0), "item")
|
||||
else h.get("scores", 0)))[:min(len(humans_orig), 4)]
|
||||
t = sorted(humans_topk, key=lambda h: -(
|
||||
h.get("scores", 0).item() if hasattr(h.get("scores", 0), "item")
|
||||
else h.get("scores", 0)))[:len(o)]
|
||||
|
||||
for i, (ho, ht) in enumerate(zip(o, t)):
|
||||
vo = ho["v3d"].detach().cpu().numpy().flatten()
|
||||
vt = ht["v3d"].detach().cpu().numpy().flatten()
|
||||
dot = float(np.dot(vo, vt))
|
||||
nv = float(np.linalg.norm(vo) * np.linalg.norm(vt) + 1e-9)
|
||||
cos = dot / nv
|
||||
mae = float(np.mean(np.abs(vo - vt)))
|
||||
sco = (ho.get("scores", 0).item()
|
||||
if hasattr(ho.get("scores", 0), "item") else ho.get("scores", 0))
|
||||
sct = (ht.get("scores", 0).item()
|
||||
if hasattr(ht.get("scores", 0), "item") else ht.get("scores", 0))
|
||||
print(f" [{i}] cosine={cos:.6f} mae={mae*1000:.3f}mm "
|
||||
f"score_orig={sco:.4f} score_topk={sct:.4f}")
|
||||
@@ -67,6 +67,33 @@ these as untraceable.
|
||||
single-day estimate was optimistic — backbone alone needs ~½ day of
|
||||
surgery before head can even be touched.
|
||||
|
||||
### Task 2 + 3 attempted (2026-05-13, post-breakthrough)
|
||||
|
||||
**Task 2 — apply_topk validation : PASS**
|
||||
|
||||
`apply_topk(K=4)` validée comme drop-in pour `apply_threshold` :
|
||||
cosine sim 1.000000 sur 4/4 détections (image example_data), MAE
|
||||
2-4 mm = bruit float-reorder. Script : `scripts/probe_head_topk.py`.
|
||||
|
||||
**Task 3 — Full Multi-HMR convert : en cours, patches itératifs**
|
||||
|
||||
Tentative conversion full Multi-HMR avec :
|
||||
- apply_threshold monkey-patched → apply_topk(K=4)
|
||||
- backbone.encoder.interpolate_pos_encoding → buffer fige
|
||||
- utils.camera.inverse_perspective_projection → closed-form K_inv
|
||||
|
||||
→ TracedMHMR.forward sanity OK, jit.trace OK, mais conversion
|
||||
coremltools échoue successivement sur :
|
||||
1. `upsample_bicubic2d` ✅ résolu par pos_embed fix
|
||||
2. `aten::inverse` ✅ résolu par closed-form K_inv
|
||||
3. `Types should have zero-rank ndarray input, got [0.]` ❌ encore
|
||||
|
||||
Pattern : chaque patch débloque la couche suivante. Estimation
|
||||
restante : 1-2 jours pour couvrir tous les ops résiduels.
|
||||
|
||||
Script Task 3 : `scripts/coreml_full_probe.py` (laissé en place pour
|
||||
reprise — point d'arrêt précis documenté).
|
||||
|
||||
### Probe v4 (2026-05-13) — BREAKTHROUGH
|
||||
|
||||
Avec **2 patches au lieu d'1**, la conversion DINOv2 ViT-S 672x672
|
||||
|
||||
Reference in New Issue
Block a user