File size: 2,139 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 | """Step 4b: attempt on-device ANE compile of the converted model.
Loads dense_sam3_trackstep.mlpackage with CPU_AND_NE and runs one predict.
The ANECompiler runs at load; failures surface as load errors or CPU fallback.
"""
import sys
import time
import numpy as np
import torch
import common
import dense_wrapper as dw
def main():
import coremltools as ct
cache = torch.load("eager_cache.pt", weights_only=False)
# need model only for input synthesis
model = common.build_model()
inputs = dw.frame_inputs(model, cache, cache["target_frame"])
unit = getattr(ct.ComputeUnit, sys.argv[1] if len(sys.argv) > 1 else "CPU_AND_NE")
print(f"loading with compute_units={unit} ...")
t0 = time.time()
mlmodel = ct.models.MLModel("dense_sam3_trackstep.mlpackage",
compute_units=unit)
print(f"load OK in {time.time()-t0:.1f}s")
in_names = [i.name for i in mlmodel.input_description._fd_spec]
feed = {n: v.numpy().astype(np.float32) for n, v in zip(in_names, inputs)}
state = mlmodel.make_state()
t0 = time.time()
out = mlmodel.predict(feed, state=state)
print(f"predict 1 OK in {time.time()-t0:.1f}s")
t0 = time.time()
out = mlmodel.predict(feed, state=state)
print(f"predict 2 OK in {time.time()-t0:.1f}s")
print("outputs:", {k: np.asarray(v).shape for k, v in out.items()})
# which units actually got used?
try:
plan = ct.models.compute_plan.MLComputePlan.load_from_path(
mlmodel.get_compiled_model_path(), compute_units=unit
)
counts = {}
prog = plan.model_structure.program
fn = prog.functions["main"]
for op in fn.block.operations:
usage = plan.get_compute_device_usage_for_mlprogram_operation(op)
if usage is None:
continue
dev = type(usage.preferred_compute_device).__name__
counts[dev] = counts.get(dev, 0) + 1
print("op placement:", counts)
except Exception as e:
print(f"compute plan inspection failed: {type(e).__name__}: {e}")
if __name__ == "__main__":
main()
|