File size: 737 Bytes
7bf8323 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | import sys, time
import numpy as np, coremltools as ct
for path in sys.argv[1:]:
m = ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_GPU)
spec = {i.name: tuple(i.type.multiArrayType.shape) for i in m.get_spec().description.input}
x = {}
for k, shp in spec.items():
dt = np.int32 if k in ("input_ids", "cand_idx", "query_idx") else np.float16
x[k] = np.zeros(shp, dtype=dt)
x["valid"][:] = 1; x["segment"][:] = np.eye(spec["segment"][0])
if "seg_chunks" in x: x["seg_chunks"][:] = np.eye(64)
m.predict(x); ts = []
for _ in range(10):
t = time.time(); m.predict(x); ts.append((time.time() - t) * 1000)
print(f"{path.split('/')[-1]:40s} p50 {np.median(ts):6.1f} ms")
|