mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 05:23:41 +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.
87 lines
4.5 KiB
Python
87 lines
4.5 KiB
Python
import tempfile, unittest
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
import numpy as np
|
|
from examples.openpilot.load_pickle import make_inputs
|
|
from tinygrad.helpers import Context, DEV, fetch, getenv
|
|
from tinygrad.engine.realize import lower_and_compile
|
|
from tinygrad.uop.ops import Ops
|
|
from examples.openpilot.helpers import load_pickle, make_retargetable
|
|
|
|
|
|
class TestCompile(unittest.TestCase):
|
|
def test_metadata(self):
|
|
import onnx
|
|
from examples.openpilot.compile_onnx import onnx_metadata
|
|
graph = onnx.helper.make_graph([onnx.helper.make_node('Identity', ['input'], ['output'])], 'metadata',
|
|
[onnx.helper.make_tensor_value_info('input', onnx.TensorProto.FLOAT, (1, 3))],
|
|
[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (1, 3))])
|
|
model = onnx.helper.make_model(graph)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / 'metadata.onnx'
|
|
for metadata in ({}, {'model_checkpoint': 'v1.0', 'output_slices': 'dGVzdA=='}):
|
|
onnx.helper.set_model_props(model, metadata)
|
|
onnx.save(model, path)
|
|
properties, output_shapes = onnx_metadata(path)
|
|
self.assertEqual(properties, metadata)
|
|
self.assertEqual(output_shapes, {'output': (1, 3)})
|
|
|
|
def test_warp_layouts(self):
|
|
from tinygrad import Tensor
|
|
from examples.openpilot.compile_warp import NV12Frame, make_luma_warp, make_frame_prepare
|
|
frame = NV12Frame(8, 6, 12, 8, 4, 144)
|
|
data = np.zeros((12, 12), dtype=np.uint8)
|
|
data[:6, :8] = np.arange(48, dtype=np.uint8).reshape(6, 8)
|
|
data[8:, 0:8:2], data[8:, 1:8:2] = 100, 200
|
|
image = Tensor(data.reshape(-1)).realize()
|
|
transform = np.eye(3, dtype=np.float32)
|
|
luma = make_luma_warp(frame, 4, 4, border_fill=16)
|
|
np.testing.assert_array_equal(luma(input_frame=image, M_inv=Tensor(transform, device='NPY')).numpy(), data[:4, :4].reshape(1, 16))
|
|
transform[0, 2] = 1000
|
|
np.testing.assert_array_equal(luma(input_frame=image, M_inv=Tensor(transform, device='NPY')).numpy(), np.full((1, 16), 16, dtype=np.uint8))
|
|
yuv = make_frame_prepare(frame, 4, 4)
|
|
expected = [[[0, 2], [16, 18]], [[8, 10], [24, 26]], [[1, 3], [17, 19]], [[9, 11], [25, 27]], [[100]*2]*2, [[200]*2]*2]
|
|
np.testing.assert_array_equal(yuv(input_frame=image, M_inv=Tensor(np.eye(3, dtype=np.float32), device='NPY')).numpy(), expected)
|
|
|
|
|
|
@unittest.skipUnless(getenv("MODEL_PKL", ""), "requires an artifact from python examples/openpilot/compile_onnx.py")
|
|
class TestCompiledModel(unittest.TestCase):
|
|
def setUp(self):
|
|
self.artifact = load_pickle(getenv("MODEL_PKL", ""), out_of_band=bool(getenv("PICKLE_OOB")))
|
|
self.model = self.artifact['run']
|
|
|
|
def test_retarget(self):
|
|
from tinygrad.runtime.autogen.libc import Elf64_Ehdr, EM_AARCH64
|
|
make_retargetable(self.model)
|
|
with Context(DEV=replace(DEV.target("CPU"), arch="arm64,generic")): linear = lower_and_compile(self.model.captured._linear)
|
|
for p in linear.toposort(enter_calls=True):
|
|
if p.op is Ops.PROGRAM and p.arg.target.device == "CPU":
|
|
self.assertEqual(Elf64_Ehdr.from_buffer_copy(p.src[-1].arg).e_machine, EM_AARCH64)
|
|
|
|
def test_inputs(self):
|
|
def run(seed):
|
|
inputs = make_inputs(self.artifact, seed)
|
|
self.model(**inputs)
|
|
return [t.numpy().copy() for t in inputs['output_buffers'].values()]
|
|
original, changed, repeated = run(100), run(101), run(100)
|
|
for before, after, again in zip(original, changed, repeated, strict=True):
|
|
self.assertTrue(np.isfinite(before).all() and np.isfinite(after).all())
|
|
np.testing.assert_array_equal(before, again)
|
|
self.assertFalse(all(np.array_equal(a, b) for a, b in zip(original, changed, strict=True)))
|
|
|
|
@unittest.skipUnless(getenv("MODEL_ONNX", ""), "requires the source ONNX model")
|
|
def test_onnx(self):
|
|
import onnxruntime as ort
|
|
session = ort.InferenceSession(str(fetch(getenv("MODEL_ONNX", ""))), providers=['CPUExecutionProvider'])
|
|
input_types = {v.name: v.type.removeprefix('tensor(').removesuffix(')') for v in session.get_inputs()}
|
|
for seed in (100, 101):
|
|
inputs = make_inputs(self.artifact, seed)
|
|
reference = session.run(None, {k: v.numpy().astype({'float': 'float32', 'double': 'float64'}.get(input_types[k], input_types[k]))
|
|
for k, v in inputs.items() if k in input_types})
|
|
self.model(**inputs)
|
|
for expected, actual in zip(reference, inputs['output_buffers'].values(), strict=True):
|
|
np.testing.assert_allclose(actual.numpy(), expected, atol=1e-4, rtol=1e-4)
|
|
|
|
|
|
if __name__ == '__main__': unittest.main()
|