mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-09 22:13:43 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
147 lines
6.9 KiB
Python
147 lines
6.9 KiB
Python
import unittest, onnx, tempfile, pathlib
|
|
import numpy as np
|
|
from tinygrad import Tensor
|
|
from tinygrad.uop.ops import Ops
|
|
from tinygrad.nn.onnx import OnnxRunner, OnnxDataType
|
|
from hypothesis import given, strategies as st
|
|
|
|
# copied from test_const_folding.py
|
|
def _check_ast_count(desired_count:int, t:Tensor):
|
|
# NOTE: this has side effect because everything can be scheduled only once
|
|
linear = t.schedule_linear()
|
|
asts = [call for call in linear.src if call.src[0].op is Ops.SINK]
|
|
assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
|
|
|
|
def build_onnx(nodes, from_disk:bool=True, opset_imports=None, **kwargs):
|
|
"""Helper to build and return an OnnxRunner from ONNX nodes."""
|
|
graph = onnx.helper.make_graph(nodes, 'test', kwargs.get('inputs', []), kwargs.get('outputs', []), kwargs.get('initializers', []))
|
|
model = onnx.helper.make_model(graph) if opset_imports is None else onnx.helper.make_model(graph, opset_imports=opset_imports)
|
|
if from_disk:
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tmp_path = pathlib.Path(tmpdir)
|
|
model_path = tmp_path / "model.onnx"
|
|
onnx.save(model, model_path)
|
|
runner = OnnxRunner(model_path)
|
|
else:
|
|
# use the in-memory method
|
|
runner = OnnxRunner(Tensor(model.SerializeToString(), device="PYTHON"))
|
|
return runner
|
|
|
|
class TestOnnxRunner(unittest.TestCase):
|
|
def test_tinygrad_contiguous(self):
|
|
runner = build_onnx(
|
|
nodes=[
|
|
onnx.helper.make_node('Add', ['inp', 'one'], ['added']),
|
|
onnx.helper.make_node('Contiguous', ['added'], ['materialized'], domain='org.tinygrad'),
|
|
onnx.helper.make_node('Mul', ['materialized', 'two'], ['output'])
|
|
],
|
|
inputs=[onnx.helper.make_tensor_value_info('inp', onnx.TensorProto.FLOAT, (4,))],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (4,))],
|
|
initializers=[
|
|
onnx.helper.make_tensor('one', onnx.TensorProto.FLOAT, (), [1.0]),
|
|
onnx.helper.make_tensor('two', onnx.TensorProto.FLOAT, (), [2.0])
|
|
],
|
|
opset_imports=[onnx.helper.make_opsetid('', 13), onnx.helper.make_opsetid('org.tinygrad', 1)],
|
|
from_disk=False).to('PYTHON')
|
|
_check_ast_count(2, runner({'inp': Tensor.empty(4, device='PYTHON')})['output'])
|
|
|
|
def _test_const_fold_unary_op(self, from_disk:bool):
|
|
runner = build_onnx(
|
|
nodes=[
|
|
onnx.helper.make_node('Expand', ['inp', 'shape'], ['expanded']),
|
|
onnx.helper.make_node('Exp', ['expanded'], ['output'])
|
|
],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (5,))],
|
|
initializers=[
|
|
onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (), [1.0]),
|
|
onnx.helper.make_tensor('shape', onnx.TensorProto.INT64, (1,), [5])
|
|
],
|
|
from_disk=from_disk)
|
|
output = runner({'inp': Tensor([1.0])})['output']
|
|
_check_ast_count(0, output)
|
|
|
|
def _test_const_fold_binary_op(self, from_disk:bool):
|
|
runner = build_onnx(
|
|
nodes=[onnx.helper.make_node('Add', ['inp', 'const'], ['output'])],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (4,))],
|
|
initializers=[
|
|
onnx.helper.make_tensor('inp', onnx.TensorProto.FLOAT, (4,), [1, 2, 3, 4]),
|
|
onnx.helper.make_tensor('const', onnx.TensorProto.FLOAT, (), [0])
|
|
],
|
|
from_disk=from_disk)
|
|
output = runner({'inp': Tensor([1, 2, 3, 4])})['output']
|
|
_check_ast_count(0, output)
|
|
|
|
@unittest.skip("const folding is removed")
|
|
def test_const_fold_from_disk(self):
|
|
self._test_const_fold_unary_op(True)
|
|
self._test_const_fold_binary_op(True)
|
|
|
|
@unittest.skip("const folding is removed")
|
|
def test_const_fold_from_memory(self):
|
|
self._test_const_fold_unary_op(False)
|
|
# TODO: understand this and fix this, bitcast related
|
|
# self._test_const_fold_binary_op(False)
|
|
|
|
def test_external_data_loading(self):
|
|
weights = np.arange(4, dtype=np.float32)
|
|
tensor_with_data = onnx.helper.make_tensor('weights', onnx.TensorProto.FLOAT, weights.shape, weights.tobytes(), raw=True)
|
|
graph = onnx.helper.make_graph(
|
|
nodes=[onnx.helper.make_node('Add', ['inp', 'weights'], ['output'])],
|
|
name='test_external',
|
|
inputs=[onnx.helper.make_tensor_value_info('inp', onnx.TensorProto.FLOAT, (1,))],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, weights.shape)],
|
|
initializer=[tensor_with_data]
|
|
)
|
|
model = onnx.helper.make_model(graph)
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tmp_path = pathlib.Path(tmpdir)
|
|
model_path = tmp_path / "model.onnx"
|
|
onnx.save_model(model, model_path, save_as_external_data=True, all_tensors_to_one_file=True, size_threshold=0, location="weights.onnx_data")
|
|
runner = OnnxRunner(model_path)
|
|
output = runner({'inp': Tensor([1])})['output']
|
|
np.testing.assert_equal(output.numpy(), weights + 1)
|
|
|
|
all_dtypes = list(OnnxDataType)
|
|
|
|
class TestOnnxRunnerDtypes(unittest.TestCase):
|
|
"""
|
|
Internal tensors (initializers, attributes) fallback to default dtype if unsupported by device.
|
|
External tensors (inputs) preserve their original dtype - user must ensure compatibility with device.
|
|
"""
|
|
def _get_expected_dtype(self, onnx_dtype: int, is_input: bool): return OnnxDataType(onnx_dtype).to_dtype()
|
|
|
|
@given(onnx_dtype=st.sampled_from(all_dtypes))
|
|
def test_input_dtype(self, onnx_dtype: int):
|
|
expected_dtype = self._get_expected_dtype(onnx_dtype, True)
|
|
runner = build_onnx(
|
|
nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
|
|
inputs=[onnx.helper.make_tensor_value_info('input', onnx_dtype, ())],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, ())],
|
|
from_disk=False)
|
|
self.assertEqual(runner.graph_inputs['input'].dtype, expected_dtype)
|
|
|
|
@given(onnx_dtype=st.sampled_from(all_dtypes))
|
|
def test_initializer_dtype(self, onnx_dtype: int):
|
|
expected_dtype = self._get_expected_dtype(onnx_dtype, False)
|
|
runner = build_onnx(
|
|
nodes=[onnx.helper.make_node('Identity', ['initializer'], ['output'])],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
|
|
initializers=[onnx.helper.make_tensor('initializer', onnx_dtype, (2,), [1, 2])],
|
|
from_disk=False)
|
|
self.assertEqual(runner.graph_values['initializer'].dtype, expected_dtype)
|
|
|
|
@given(onnx_dtype=st.sampled_from(all_dtypes))
|
|
def test_node_attribute_dtype(self, onnx_dtype: int):
|
|
expected_dtype = self._get_expected_dtype(onnx_dtype, False)
|
|
value_tensor = onnx.helper.make_tensor('value', onnx_dtype, (2,), [1, 2])
|
|
runner = build_onnx(
|
|
nodes=[onnx.helper.make_node('Constant', [], ['output'], value=value_tensor)],
|
|
outputs=[onnx.helper.make_tensor_value_info('output', onnx_dtype, (2,))],
|
|
from_disk=False)
|
|
self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, expected_dtype)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|