| """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) |
| |
| 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()}) |
|
|
| |
| 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() |
|
|
