| """torch.export -> coremltools convert (fp16 default) of DenseTrackStep, |
| then CoreML-vs-torch parity over warmup→steady stateful frames. |
| """ |
|
|
| import time |
| import numpy as np |
| import torch |
|
|
| import common |
| import dense_wrapper as dw |
|
|
|
|
| def zero_state(wrapper): |
| wrapper.mem_bank.zero_() |
| wrapper.img_bank.zero_() |
| wrapper.ptr_bank.zero_() |
| wrapper.mem_valid.zero_() |
| wrapper.ptr_valid.zero_() |
|
|
|
|
| def main(): |
| cache = torch.load("eager_cache.pt", weights_only=False) |
| wrapper, model = dw.build_wrapper() |
| inputs = dw.frame_inputs(model, cache, 3) |
|
|
| common.hide_triton_stub() |
| t0 = time.time() |
| with torch.no_grad(): |
| try: |
| ep = torch.export.export(wrapper, inputs) |
| except Exception as e: |
| print(f"strict export failed ({type(e).__name__}: {e}); retrying strict=False") |
| ep = torch.export.export(wrapper, inputs, strict=False) |
| ep = ep.run_decompositions({}) |
| print(f"torch.export + decompositions OK in {time.time()-t0:.1f}s") |
|
|
| import coremltools as ct |
| from coremltools.converters.mil.frontend.torch.torch_op_registry import ( |
| register_torch_op, |
| ) |
| from coremltools.converters.mil.frontend.torch.ops import _get_inputs |
| from coremltools.converters.mil.mil import Builder as mb |
|
|
| @register_torch_op(torch_alias=["where.scalarother"]) |
| def where_scalarother(context, node): |
| cond, a, b = _get_inputs(context=context, node=node, expected=3) |
| context.add(mb.select(cond=cond, a=a, b=b), node.name) |
|
|
| t0 = time.time() |
| mlmodel = ct.convert( |
| ep, |
| minimum_deployment_target=ct.target.iOS18, |
| compute_units=ct.ComputeUnit.CPU_AND_GPU, |
| ) |
| print(f"coremltools convert OK in {time.time()-t0:.1f}s") |
| mlmodel.save("dense_sam3_trackstep.mlpackage") |
| print("saved dense_sam3_trackstep.mlpackage") |
|
|
| zero_state(wrapper) |
| frames = list(range(1, 19)) |
| torch_outs = [] |
| with torch.no_grad(): |
| for f in frames: |
| torch_outs.append([t.clone() for t in wrapper(*dw.frame_inputs(model, cache, f))]) |
|
|
| in_names = [i.name for i in mlmodel.input_description._fd_spec] |
| out_names = [o.name for o in mlmodel.output_description._fd_spec] |
| print("inputs:", in_names) |
| print("outputs:", out_names) |
| print(f"{'frame':>5} {'osl_rel':>9} {'iou_rel':>9} {'sign_agree':>10} {'low_rel':>9}") |
|
|
| state = mlmodel.make_state() |
| worst = 0.0 |
| appear_flips = 0 |
| for i, f in enumerate(frames): |
| feed = {n: v.numpy().astype(np.float32) |
| for n, v in zip(in_names, dw.frame_inputs(model, cache, f))} |
| got = mlmodel.predict(feed, state=state) |
| by_shape = {tuple(np.asarray(v).shape): torch.from_numpy( |
| np.asarray(v)).float() for v in got.values()} |
| t_low, t_high, t_osl, t_ious = torch_outs[i] |
| c_low = by_shape[tuple(t_low.shape)] |
| c_osl = by_shape[tuple(t_osl.shape)] |
| c_ious = by_shape[tuple(t_ious.shape)] |
|
|
| def rel(a, b): |
| return ((a - b).abs().max() / b.abs().max().clamp_min(1e-9)).item() |
|
|
| sign = ((c_low > 0) == (t_low > 0)).float().mean().item() |
| r_osl, r_iou, r_low = rel(c_osl, t_osl), rel(c_ious, t_ious), rel(c_low, t_low) |
| worst = max(worst, r_osl, r_iou, r_low) |
| |
| if not torch.equal((c_osl > 0), (t_osl > 0)): |
| appear_flips += 1 |
| print(f"{f:>5} {r_osl:>9.2e} {r_iou:>9.2e} {sign:>10.4f} {r_low:>9.2e}") |
|
|
| print(f"worst rel: {worst:.3e} appear_flips: {appear_flips}/{len(frames)}") |
| print("FP16 TRAJ:", "PASS" if worst < 5e-2 and appear_flips == 0 else "CHECK") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
