mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-09 15:08:02 -05:00
5ce278b245
* file path as input and have parse be in OnnxRunner.__init__ * modelproto_to_onnxrunner -> modelproto_to_runner * whoops, fix import * oh flakiness again, is it because it's getting gc-ed? * small changes * CI flaky so just move compile4 fix in * copy typing of onnx_load * actually can just import onnx_load instead of onnx.load * fix external_benchmark_openpilot * fix onnx_runner test to use onnx_helper * rerun CI * try run_modelproto * spam CI a few times * revert run_modelproto since that's flaky also * no external onnx_load usage except onnx.py * cursor tab complete is evil. Snuck a darn sorted in. But does order change result? Why? * model_benchmark 193s -> 80s, add OnnxRunner.to()... * minimize diff and clean up * device can be None, weird but eh --------- Co-authored-by: chenyu <chenyu@fastmail.com>
36 lines
1.2 KiB
Python
36 lines
1.2 KiB
Python
import sys, time
|
|
from tinygrad import TinyJit, GlobalCounters, fetch, getenv
|
|
from tinygrad.frontend.onnx import OnnxRunner
|
|
from extra.onnx_helpers import get_example_inputs, validate
|
|
|
|
def load_onnx_model(onnx_file):
|
|
run_onnx = OnnxRunner(onnx_file)
|
|
run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
|
|
return run_onnx_jit, run_onnx.graph_inputs
|
|
|
|
if __name__ == "__main__":
|
|
onnx_file = fetch(sys.argv[1])
|
|
run_onnx_jit, input_specs = load_onnx_model(onnx_file)
|
|
print("loaded model")
|
|
|
|
for i in range(3):
|
|
new_inputs = get_example_inputs(input_specs)
|
|
GlobalCounters.reset()
|
|
print(f"run {i}")
|
|
run_onnx_jit(**new_inputs)
|
|
|
|
# run 20 times
|
|
for _ in range(20):
|
|
new_inputs = get_example_inputs(input_specs)
|
|
GlobalCounters.reset()
|
|
st = time.perf_counter()
|
|
out = run_onnx_jit(**new_inputs)
|
|
mt = time.perf_counter()
|
|
val = out.numpy()
|
|
et = time.perf_counter()
|
|
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
|
|
|
|
if getenv("ORT"):
|
|
validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
|
|
print("model validated")
|