From 74d98eafb871e56d4323ed2577a86cb07f5902f2 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 24 Mar 2025 12:24:34 +0800 Subject: [PATCH] add onnx frontend stub [pr] (#9558) --- examples/benchmark_onnx.py | 2 +- examples/compile_tensorflow.py | 2 +- examples/openpilot/compile3.py | 2 +- examples/yolov8-onnx.py | 2 +- extra/huggingface_onnx/run_models.py | 3 ++- extra/onnx_helpers.py | 3 ++- test/external/external_benchmark_openpilot.py | 2 +- test/external/external_model_benchmark.py | 2 +- test/external/external_test_onnx_backend.py | 2 +- test/models/test_onnx.py | 2 +- test/test_quantize_onnx.py | 2 +- tinygrad/frontend/onnx.py | 5 +++++ 12 files changed, 18 insertions(+), 11 deletions(-) create mode 100644 tinygrad/frontend/onnx.py diff --git a/examples/benchmark_onnx.py b/examples/benchmark_onnx.py index f0cf7d5d34..2a1425ba6e 100644 --- a/examples/benchmark_onnx.py +++ b/examples/benchmark_onnx.py @@ -1,6 +1,6 @@ import sys, onnx, time from tinygrad import TinyJit, Device, GlobalCounters, fetch, getenv -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner from extra.onnx_helpers import get_example_inputs, validate def load_onnx_model(onnx_file): diff --git a/examples/compile_tensorflow.py b/examples/compile_tensorflow.py index c906592282..3ad6f31195 100644 --- a/examples/compile_tensorflow.py +++ b/examples/compile_tensorflow.py @@ -8,7 +8,7 @@ import numpy as np import subprocess import tensorflow as tf import tf2onnx -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner from tinygrad.tensor import Tensor from extra.export_model import export_model_clang, compile_net, jit_model diff --git a/examples/openpilot/compile3.py b/examples/openpilot/compile3.py index b34072e0bc..476d407250 100644 --- a/examples/openpilot/compile3.py +++ b/examples/openpilot/compile3.py @@ -12,7 +12,7 @@ from tinygrad.engine.realize import CompiledRunner import onnx from onnx.helper import tensor_dtype_to_np_dtype -from extra.onnx import OnnxRunner # TODO: port to main tinygrad +from tinygrad.frontend.onnx import OnnxRunner 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" diff --git a/examples/yolov8-onnx.py b/examples/yolov8-onnx.py index b00938e0dd..9d14024547 100644 --- a/examples/yolov8-onnx.py +++ b/examples/yolov8-onnx.py @@ -3,7 +3,7 @@ import os from ultralytics import YOLO import onnx from pathlib import Path -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner from extra.onnx_helpers import get_example_inputs from tinygrad.tensor import Tensor diff --git a/extra/huggingface_onnx/run_models.py b/extra/huggingface_onnx/run_models.py index a0acc2af72..c81952216d 100644 --- a/extra/huggingface_onnx/run_models.py +++ b/extra/huggingface_onnx/run_models.py @@ -1,6 +1,7 @@ import onnx, yaml, tempfile, time, collections, pprint, argparse, json from pathlib import Path -from extra.onnx import OnnxRunner, get_onnx_ops +from tinygrad.frontend.onnx import OnnxRunner +from extra.onnx import get_onnx_ops from extra.onnx_helpers import validate, get_example_inputs def get_config(root_path: Path): diff --git a/extra/onnx_helpers.py b/extra/onnx_helpers.py index cce053f059..ab020fdb23 100644 --- a/extra/onnx_helpers.py +++ b/extra/onnx_helpers.py @@ -1,6 +1,7 @@ from tinygrad import Tensor from tinygrad.tensor import _to_np_dtype -from extra.onnx import OnnxRunner, OnnxValue +from tinygrad.frontend.onnx import OnnxRunner +from extra.onnx import OnnxValue import onnx import numpy as np import onnxruntime as ort diff --git a/test/external/external_benchmark_openpilot.py b/test/external/external_benchmark_openpilot.py index 780539b1e7..910d8a5ed3 100644 --- a/test/external/external_benchmark_openpilot.py +++ b/test/external/external_benchmark_openpilot.py @@ -2,7 +2,7 @@ import time, sys, hashlib from pathlib import Path import onnx from onnx.helper import tensor_dtype_to_np_dtype -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner from tinygrad import Tensor, dtypes, TinyJit from tinygrad.helpers import IMAGE, GlobalCounters, fetch, colored, getenv, trange from tinygrad.tensor import _from_np_dtype diff --git a/test/external/external_model_benchmark.py b/test/external/external_model_benchmark.py index 6653839ce0..5b1199e1e4 100644 --- a/test/external/external_model_benchmark.py +++ b/test/external/external_model_benchmark.py @@ -6,7 +6,7 @@ import onnx from onnx.helper import tensor_dtype_to_np_dtype import onnxruntime as ort from onnx2torch import convert -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner from tinygrad.helpers import OSX, DEBUG, fetch from tinygrad import Tensor, Device from tinygrad.device import CompileError diff --git a/test/external/external_test_onnx_backend.py b/test/external/external_test_onnx_backend.py index fd9a5896ee..972d7a0bc8 100644 --- a/test/external/external_test_onnx_backend.py +++ b/test/external/external_test_onnx_backend.py @@ -10,7 +10,7 @@ from tinygrad.device import is_dtype_supported # pip3 install tabulate pytest_plugins = 'onnx.backend.test.report', -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner class TinygradModel(BackendRep): def __init__(self, run_onnx, input_names): diff --git a/test/models/test_onnx.py b/test/models/test_onnx.py index e3d1868aed..d1b7def35e 100644 --- a/test/models/test_onnx.py +++ b/test/models/test_onnx.py @@ -7,7 +7,7 @@ try: import onnx except ModuleNotFoundError: raise unittest.SkipTest("onnx not installed, skipping onnx test") -from extra.onnx import OnnxRunner +from tinygrad.frontend.onnx import OnnxRunner from tinygrad.tensor import Tensor from tinygrad.helpers import CI, fetch, temp diff --git a/test/test_quantize_onnx.py b/test/test_quantize_onnx.py index 8e062a11af..15ae65cee8 100644 --- a/test/test_quantize_onnx.py +++ b/test/test_quantize_onnx.py @@ -70,7 +70,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase): import onnx except ImportError: raise unittest.SkipTest() - from extra.onnx import OnnxRunner + from tinygrad.frontend.onnx import OnnxRunner out_file = get_quantized_model(sz) onnx_model = onnx.load(out_file) run_onnx = OnnxRunner(onnx_model) diff --git a/tinygrad/frontend/onnx.py b/tinygrad/frontend/onnx.py new file mode 100644 index 0000000000..5c886163b7 --- /dev/null +++ b/tinygrad/frontend/onnx.py @@ -0,0 +1,5 @@ +# type: ignore +import sys, pathlib +sys.path.append(pathlib.Path(__file__).parent.parent.as_posix()) +try: from extra.onnx import OnnxRunner # noqa: F401 # pylint: disable=unused-import +except ImportError as e: raise ImportError("onnx frontend not in release\nTo fix, install tinygrad from a git checkout with pip install -e .") from e \ No newline at end of file