File size: 2,759 Bytes
29f25be | 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 | """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()
|