File size: 3,241 Bytes
fcf4209 c6e4872 fcf4209 c6e4872 fcf4209 c6e4872 fcf4209 c6e4872 fcf4209 c6e4872 fcf4209 | 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 | """Selected fixed-bucket end-to-end Core ML extraction latency on this Mac."""
import argparse
import json
import platform
import statistics
import time
from pathlib import Path
import coremltools as ct
import psutil
from gliner2 import Schema
from extraction_runtime import CoreMLAdaptiveBoundaryExtractor, CoreMLBoundaryExtractor
UNITS = {
"cpu_only": ct.ComputeUnit.CPU_ONLY,
"cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU,
"cpu_and_neural_engine": ct.ComputeUnit.CPU_AND_NE,
"all": ct.ComputeUnit.ALL,
}
def percentile(values, fraction):
ordered = sorted(values)
return ordered[min(round(fraction * (len(ordered) - 1)), len(ordered) - 1)]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", required=True)
parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp32")
parser.add_argument("--feature-package")
parser.add_argument("--adaptive-w8", action="store_true")
parser.add_argument("--units", choices=list(UNITS), default="all")
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--iterations", type=int, default=200)
args = parser.parse_args()
text = "Alice founded Acme in Toronto in 2020."
schema = Schema().entities(["person", "organization", "location"])
started = time.perf_counter()
if args.adaptive_w8:
runtime = CoreMLAdaptiveBoundaryExtractor(args.model_dir, fp16_compute_units=UNITS[args.units])
else:
runtime = CoreMLBoundaryExtractor(
args.model_dir,
precision=args.precision,
compute_units=UNITS[args.units],
feature_package=args.feature_package,
)
load_ms = (time.perf_counter() - started) * 1000
for _ in range(args.warmup):
runtime.extract(text, schema)
process = psutil.Process()
latencies = []
peak_rss = process.memory_info().rss
for _ in range(args.iterations):
start = time.perf_counter()
runtime.extract(text, schema)
latencies.append((time.perf_counter() - start) * 1000)
peak_rss = max(peak_rss, process.memory_info().rss)
report = {
"purpose": "selected end-to-end entity extraction latency, no benchmark scoring",
"fixture": text,
"shape": "L128/W64/Q8/C192",
"precision": args.precision,
"compute_units": args.units,
"feature_package": args.feature_package,
"adaptive_w8": args.adaptive_w8,
"warmup": args.warmup,
"iterations": args.iterations,
"load_ms": load_ms,
"p50_ms": statistics.median(latencies),
"p95_ms": percentile(latencies, 0.95),
"mean_ms": statistics.mean(latencies),
"peak_process_rss_bytes": peak_rss,
"macos": platform.mac_ver()[0],
"machine": platform.machine(),
"coremltools": ct.__version__,
}
folder = Path(args.model_dir)
variant = "adaptive-w8" if args.adaptive_w8 else "w8-embedding" if args.feature_package else "baseline"
path = folder / f"benchmark-{args.precision}-{args.units}-{variant}.json"
path.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|