File size: 3,685 Bytes
47d8ad6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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)
# is_obj_appearing threshold at 0
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()
|