| """Report where Core ML would place each op of a package (CPU / GPU / Neural Engine). |
| |
| Uses MLComputePlan on this Mac. Placement on iPhone can differ; confirm with a device run. |
| """ |
| import argparse |
| import json |
| import shutil |
| import tempfile |
| from collections import Counter |
| from pathlib import Path |
|
|
|
|
| def plan(package, units): |
| import coremltools as ct |
| from coremltools.models.compute_plan import MLComputePlan |
| from coremltools.models.compute_device import MLCPUComputeDevice, MLGPUComputeDevice, MLNeuralEngineComputeDevice |
|
|
| names = {MLCPUComputeDevice: "cpu", MLGPUComputeDevice: "gpu", MLNeuralEngineComputeDevice: "ane"} |
| temporary = Path(tempfile.mkdtemp()) |
| try: |
| compiled = ct.utils.compile_model(str(package), str(temporary / "model.mlmodelc")) |
| compute_plan = MLComputePlan.load_from_path(path=str(compiled), compute_units=getattr(ct.ComputeUnit, units)) |
| program = compute_plan.model_structure.program |
| placement, cost, details = Counter(), Counter(), [] |
| for function in program.functions.values(): |
| for operation in function.block.operations: |
| usage = compute_plan.get_compute_device_usage_for_mlprogram_operation(operation) |
| if usage is None: |
| continue |
| device = names.get(type(usage.preferred_compute_device), "other") |
| estimate = compute_plan.get_estimated_cost_for_mlprogram_operation(operation) |
| weight = estimate.weight if estimate else 0.0 |
| placement[device] += 1 |
| cost[device] += weight |
| details.append({"op": operation.operator_name, "device": device, "cost": weight, |
| "supported": sorted({names.get(type(d), "other") |
| for d in usage.supported_compute_devices})}) |
| return {"package": str(package), "compute_units": units, "ops_by_device": dict(placement), |
| "estimated_cost_by_device": {k: round(v, 4) for k, v in cost.items()}, "ops": details} |
| finally: |
| shutil.rmtree(temporary, ignore_errors=True) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model", type=Path, required=True, help="Experiment directory with model.mlpackage") |
| parser.add_argument("--compute-units", default="CPU_AND_NE", choices=("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL")) |
| parser.add_argument("--details", action="store_true", help="Print every op") |
| args = parser.parse_args() |
| report = plan(args.model / "model.mlpackage", args.compute_units) |
| if not args.details: |
| report.pop("ops") |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
