onnx parser (#10435)

* onnx parser

* fix compile, lint

* onnx.load -> onnx_load

* compatible with ModelProto

* fix test external_test_onnx_ops.py

* fix tests

* fix signed int

* reduce to 261 lines

* fix TypeProto.Optional

* debug for _parse_message, add TypeProto.Sequence, cleanup

* onnx_load from Tensor

* remove BufferedReader

* 174 lines and reduce tensor copy

* cleanup

* use onnx_load in external_model_benchmark.py

* fix qcom test

* [onnx] parser support external data

---------

Co-authored-by: b1tg <b1tg@users.noreply.github.com>
Co-authored-by: chenyu <chenyu@fastmail.com>

examples/benchmark_onnx.py

@@ -1,10 +1,10 @@
-import sys, onnx, time, pickle
+import sys, time, pickle
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs, validate
 
 def load_onnx_model(onnx_file):
-  onnx_model = onnx.load(onnx_file)
+  onnx_model = onnx_load(onnx_file)
   run_onnx = OnnxRunner(onnx_model)
   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

examples/openpilot/compile3.py

@@ -12,13 +12,13 @@ from tinygrad.engine.realize import CompiledRunner
 
 import onnx
 from onnx.helper import tensor_dtype_to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
 
 def compile(onnx_file):
-  onnx_model = onnx.load(onnx_file)
+  onnx_model = onnx_load(onnx_file)
   run_onnx = OnnxRunner(onnx_model)
   print("loaded model")
 

examples/yolov8-onnx.py

@@ -1,9 +1,8 @@
 #!/usr/bin/env python3
 import os
 from ultralytics import YOLO
-import onnx
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs
 from tinygrad.tensor import Tensor
 
@@ -11,6 +10,6 @@ os.chdir("/tmp")
 if not Path("yolov8n-seg.onnx").is_file():
   model = YOLO("yolov8n-seg.pt")
   model.export(format="onnx", imgsz=[480,640])
-onnx_model = onnx.load(open("yolov8n-seg.onnx", "rb"))
+onnx_model = onnx_load(open("yolov8n-seg.onnx", "rb"))
 run_onnx = OnnxRunner(onnx_model)
 run_onnx(get_example_inputs(run_onnx.graph_inputs), debug=True)