docs(plans): Multi-HMR CoreML conversion plan
This commit is contained in:
L'électron rare
@@ -0,0 +1,247 @@
|
||||
# Multi-HMR → CoreML conversion (ANE-capable, fixed-K detection)
|
||||
|
||||
> **For agentic workers:** REQUIRED SUB-SKILL: Use
|
||||
> `superpowers:subagent-driven-development` or
|
||||
> `superpowers:executing-plans` to implement this plan task-by-task.
|
||||
> Steps use checkbox (`- [ ]`) syntax for tracking.
|
||||
|
||||
**Goal:** convert Multi-HMR ViT-S to a CoreML `.mlpackage` runnable on
|
||||
M5 ANE+GPU, reaching ≥15 fps on M5 (target 20-30 fps) without thermal
|
||||
throttle. Replace the current PyTorch MPS path (3.8 fps measured, thermal
|
||||
drift to ~3 fps) in `data_only_viz/multi_hmr_worker.py`.
|
||||
|
||||
**Non-goals:**
|
||||
- Convert ViT-L or ViT-B variants (S only — already smallest).
|
||||
- Train a new model (we only do model surgery on inference path).
|
||||
- Port the SMPL-X decoder; Multi-HMR already returns `v3d` directly.
|
||||
- Multi-batch inference (B=1 fixed; we run real-time single camera).
|
||||
|
||||
**Architecture:** Multi-HMR forward decomposes into (a) DINOv2 ViT-S
|
||||
backbone (672×672 → tokens), (b) per-token classification head, (c)
|
||||
detection thresholding via `torch.where(scores >= det_thresh)`, (d)
|
||||
per-detection regression heads producing `v3d`, `shape`, `expression`,
|
||||
`transl`. Step (c) is the data-dependent op CoreML refuses. The fix:
|
||||
replace `where()` with `topk(K=4)` returning the K highest-confidence
|
||||
detections; pad/zero remaining slots; emit a score tensor so downstream
|
||||
filtering can be done on the Swift side.
|
||||
|
||||
## Research summary (2026-05-13)
|
||||
|
||||
- **coremltools** (latest stable) supports `RangeDim` /
|
||||
`EnumeratedShapes` for input flexibility but explicitly excludes
|
||||
data-dependent output cardinality for ML Programs (apple.github.io
|
||||
flexible-inputs.html).
|
||||
- Open issues #2189 (ExecuTorch dynamic index), #2040 (`index_put`),
|
||||
#1678 (traced `index_put`) confirm `nonzero`/`where` + indexing
|
||||
patterns still unsupported in 2024-2026.
|
||||
- naver/multi-hmr repo has zero issues/PRs on CoreML, ANE, MLX, ONNX,
|
||||
or Apple Silicon → no prior art; we are first-mover.
|
||||
- `torch.export` (PyTorch 2.5+) + `coremltools.convert` is the modern
|
||||
path, but inherits the same data-dependent restrictions.
|
||||
|
||||
## Risk register
|
||||
|
||||
| Risk | Likelihood | Impact | Mitigation |
|
||||
|---|---|---|---|
|
||||
| ViT-S DINOv2 has CoreML-unsupported op (rare in 2026) | low | high | trace bare backbone first (Task 1) — fail fast |
|
||||
| `transl` head uses `inverse_perspective_projection` w/ matmul on dyn shape | medium | medium | check trace; fallback compute in PyTorch on Python side |
|
||||
| ANE doesn't accept 672×672 input (ANE prefers ≤512 in practice) | medium | medium | benchmark with `compute_units=.cpuAndGPU` first, then `.all`; if ANE rejects, GPU-only still beats MPS path |
|
||||
| Numerical drift after K-fixed swap → false-positive detections | medium | low | Swift-side score gate (already exists via `det_thresh` arg path) |
|
||||
| End-to-end perf doesn't beat MPS baseline | low | high | Phase A bench (headless dummy) is the gate — abandon before integration |
|
||||
|
||||
## Task 1: Bare backbone trace + convert sanity check
|
||||
|
||||
**Context:** before touching the detection head, prove that the
|
||||
DINOv2-S backbone alone (the bulk of the compute) converts and runs.
|
||||
This isolates the bottleneck risk.
|
||||
|
||||
- [ ] **Step 1: Extract backbone-only forward**
|
||||
|
||||
Write a small wrapper module that loads the ViT-S checkpoint, takes
|
||||
the backbone (`model.backbone`), and runs only that on a
|
||||
`(1, 3, 672, 672)` input. Verify it matches the original forward's
|
||||
token output for the first 5 frames of a webcam stream (cosine sim
|
||||
> 0.999 on the token tensor).
|
||||
|
||||
Output: `data_only_viz/_coreml_backbone_probe.py` + numerical proof.
|
||||
|
||||
- [ ] **Step 2: `torch.jit.trace` the wrapper**
|
||||
|
||||
```python
|
||||
example = torch.rand(1, 3, 672, 672)
|
||||
traced = torch.jit.trace(BackboneWrapper(model), example)
|
||||
```
|
||||
|
||||
Escalate if trace fails: identify the op (likely `unpatch` or fourier
|
||||
encoding). For unfixable ops, try `torch.export.export` instead
|
||||
(PyTorch 2.5+ path).
|
||||
|
||||
- [ ] **Step 3: `coremltools.convert` backbone**
|
||||
|
||||
```python
|
||||
import coremltools as ct
|
||||
mlmodel = ct.convert(
|
||||
traced,
|
||||
inputs=[ct.TensorType(shape=(1, 3, 672, 672), name="image")],
|
||||
compute_units=ct.ComputeUnit.ALL,
|
||||
minimum_deployment_target=ct.target.macOS15,
|
||||
convert_to="mlprogram",
|
||||
)
|
||||
mlmodel.save("/tmp/multihmr_backbone.mlpackage")
|
||||
```
|
||||
|
||||
- [ ] **Step 4: Bench backbone alone**
|
||||
|
||||
Time 20 iterations of `mlmodel.predict({"image": np.random.rand(...)})`.
|
||||
Compare to PyTorch MPS backbone-only timing (re-run profiler from
|
||||
Phase A but on backbone only). **Gate: CoreML backbone median latency
|
||||
≤ 60% of MPS backbone latency, otherwise abandon plan.** Backbone
|
||||
alone on MPS is ~50 ms; CoreML target ≤30 ms.
|
||||
|
||||
- [ ] **Step 5: Commit probe**
|
||||
|
||||
```bash
|
||||
git add data_only_viz/_coreml_backbone_probe.py
|
||||
git commit -m "feat(data-only-viz): Multi-HMR backbone CoreML probe"
|
||||
```
|
||||
|
||||
## Task 2: Replace `torch.where` with fixed-K `topk` in head
|
||||
|
||||
**Context:** Multi-HMR detection head emits a score map shape
|
||||
`(1, 48, 48, 1)` (or similar) which is thresholded via `torch.where`
|
||||
to extract K_dynamic candidate persons. Replace with `topk(K=4)`
|
||||
returning the 4 highest-scoring tokens. K=4 matches the worker's
|
||||
`num_persons=4` constructor default — same downstream budget.
|
||||
|
||||
- [ ] **Step 1: Locate `apply_threshold` callsite**
|
||||
|
||||
Confirmed: `~/.cache/av-live-multihmr/multi-hmr/model.py:149` calls
|
||||
`apply_threshold(det_thresh, _scores)` returning a 4-tuple of
|
||||
index tensors used for `scores[idx[0], idx[3], idx[1], idx[2]]`
|
||||
gather at line 154-156.
|
||||
|
||||
- [ ] **Step 2: Write `apply_topk(K, _scores)` replacement**
|
||||
|
||||
Reshape `_scores` from `(B, H, W, C)` to `(B, H*W*C)`, run
|
||||
`topk(K)`, then unravel indices back to 4 components matching the
|
||||
4-tuple signature.
|
||||
|
||||
```python
|
||||
def apply_topk(K, _scores):
|
||||
B, H, W, C = _scores.shape
|
||||
flat = _scores.reshape(B, -1)
|
||||
vals, idx_flat = torch.topk(flat, k=K, dim=1)
|
||||
# Unravel: idx_flat // (H*W*C)? No — already per-batch.
|
||||
hwc = H * W * C
|
||||
idx_b = torch.arange(B, device=_scores.device).unsqueeze(1).expand(-1, K).reshape(-1)
|
||||
idx_local = idx_flat.reshape(-1)
|
||||
idx_h = (idx_local // (W * C))
|
||||
idx_w = (idx_local // C) % W
|
||||
idx_c = idx_local % C
|
||||
return (idx_b, idx_h, idx_w, idx_c), vals.reshape(-1)
|
||||
```
|
||||
|
||||
- [ ] **Step 3: Monkey-patch in a forked Model class**
|
||||
|
||||
Don't edit upstream `~/.cache/...model.py` (re-`setup_multihmr.sh`
|
||||
wipes it). Instead, subclass `Model` in
|
||||
`data_only_viz/multi_hmr_coreml.py` and override the detection head.
|
||||
Inject the patched class via the same `sys.path` mechanism used in
|
||||
`multi_hmr_worker.py`.
|
||||
|
||||
- [ ] **Step 4: Verify numerical equivalence (top-K case)**
|
||||
|
||||
When the original model emits exactly K detections with
|
||||
`det_thresh=0.3`, the new K=4 must return identical `v3d`. Use
|
||||
a recorded webcam frame with 1 person; compare `v3d` cosine sim
|
||||
> 0.9999.
|
||||
|
||||
- [ ] **Step 5: Commit head replacement**
|
||||
|
||||
## Task 3: Full model trace + CoreML conversion
|
||||
|
||||
- [ ] **Step 1: Wrap patched Model for tracing**
|
||||
|
||||
The full model's forward returns a list-of-dicts (`humans`). Tracing
|
||||
refuses dict/list outputs. Wrap to return a fixed-shape tuple:
|
||||
`(v3d[K,10475,3], shape[K,10], expression[K,10], transl[K,3],
|
||||
scores[K])`.
|
||||
|
||||
- [ ] **Step 2: `torch.jit.trace` full model**
|
||||
|
||||
Run on `(1, 3, 672, 672)` dummy. Catch and resolve trace warnings.
|
||||
Likely failures: `inverse_perspective_projection` matmul, fourier
|
||||
encoding (`torch.arange` inside forward). Move these to a
|
||||
pre-computed constant if they don't trace cleanly.
|
||||
|
||||
- [ ] **Step 3: `coremltools.convert` full model**
|
||||
|
||||
Same convert pattern as Task 1 Step 3. Save to
|
||||
`~/.cache/av-live-multihmr/checkpoints/multiHMR_672_S.mlpackage`.
|
||||
|
||||
- [ ] **Step 4: Bench full CoreML model (headless)**
|
||||
|
||||
20 iter, M5, `compute_units=ALL`. **Gate: median ≤80 ms (≥12.5 fps)
|
||||
otherwise revisit Task 1 budget assumptions.**
|
||||
|
||||
## Task 4: Worker integration
|
||||
|
||||
- [ ] **Step 1: Add CoreML backend in `multi_hmr_worker.py`**
|
||||
|
||||
Add `--multi-hmr-backend coreml` flag (default `pytorch`). Load via
|
||||
`import coremltools` (already a transitive dep of macOS Python).
|
||||
Skip the entire PyTorch model setup path when coreml selected.
|
||||
|
||||
- [ ] **Step 2: Translate I/O**
|
||||
|
||||
Input: `np.ndarray (1, 3, 672, 672) float32`. Output: 5 tensors
|
||||
packed back into the same `humans = [...]` list-of-dicts the rest
|
||||
of the worker expects. Filter on `scores >= det_thresh` in numpy.
|
||||
|
||||
- [ ] **Step 3: Camera-live bench**
|
||||
|
||||
Re-run `/tmp/bench_multihmr_camera.py` against the coreml backend
|
||||
for 60 s (longer than the 30s previous run — catch thermal stability).
|
||||
Record: median latency, fps, thermal trend.
|
||||
|
||||
- [ ] **Step 4: Commit + README update**
|
||||
|
||||
## Task 5: ANE compute unit experiments
|
||||
|
||||
- [ ] **Step 1: A/B `.all` vs `.cpuAndGPU` vs `.cpuAndNeuralEngine`**
|
||||
|
||||
ANE may reject ops or not accept the input size; need empirical
|
||||
fallback. Record which compute unit each layer ran on via
|
||||
`mlmodel.compute_unit_used_per_layer` (if available in 2026 API)
|
||||
or system-level `powermetrics`.
|
||||
|
||||
- [ ] **Step 2: If ANE-eligible: re-bench**
|
||||
|
||||
Expected: 25-40 fps on ANE-resident inference (no thermal load
|
||||
on GPU at all). This is the home run.
|
||||
|
||||
## Success criteria
|
||||
|
||||
| Metric | Target | Stretch |
|
||||
|---|---|---|
|
||||
| Median latency on M5 (camera-live, 60s) | ≤ 65 ms | ≤ 35 ms |
|
||||
| FPS sustained over 5 min (no thermal drift) | ≥ 15 | ≥ 25 |
|
||||
| Mesh quality (v3d cosine sim vs PyTorch) | ≥ 0.999 | ≥ 0.9999 |
|
||||
| Wall-time vs PyTorch MPS | ≥ 3× faster | ≥ 6× faster |
|
||||
|
||||
## Abandon criteria (cut losses fast)
|
||||
|
||||
- Task 1 Step 4 backbone-alone CoreML > 60% of MPS → abandon
|
||||
- Task 2 numerical equiv check fails (top-K doesn't approximate
|
||||
top-thresholded) → abandon, propose ONNX Runtime path instead
|
||||
- Task 3 trace fails on >2 ops requiring real model surgery →
|
||||
reframe as multi-week effort, not a perf optimization
|
||||
|
||||
## Fallback if abandoned
|
||||
|
||||
If CoreML proves infeasible within 1-day effort, fall back ranking:
|
||||
1. Frame-skip + Swift interp 60 fps (already in progress, separate file)
|
||||
2. Offload Multi-HMR inference to MacStudio M3 Ultra over Tailscale
|
||||
(worker becomes a TCP client to a remote endpoint)
|
||||
3. SMPLer-X-S port (different model, claims 64 ms M5 — see
|
||||
`data_only_viz/smpler_x_scaffold.md`)
|
||||
Reference in New Issue
Block a user