mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-10-28 09:03:46 +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.
83 lines
4.1 KiB
Python
83 lines
4.1 KiB
Python
import unittest
|
|
from extra.export_model import export_model, EXPORT_SUPPORTED_DEVICE
|
|
from tinygrad.tensor import Tensor
|
|
from tinygrad.device import Device
|
|
from tinygrad.nn import Linear
|
|
from tinygrad.nn.state import get_state_dict
|
|
from tinygrad import dtypes
|
|
import json
|
|
|
|
class MockMultiInputModel:
|
|
def forward(self, x1, x2, x3):
|
|
return x1 + x2 + x3
|
|
|
|
class MockMultiOutputModel:
|
|
def __call__(self, x1):
|
|
return x1 + 2.0, x1.pad(((0, 0), (0, 1))) + 1.0
|
|
|
|
# TODO: move compile_efficientnet tests here
|
|
@unittest.skipUnless(Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"Model export is not supported on {Device.DEFAULT}")
|
|
class TextModelExport(unittest.TestCase):
|
|
def test_multi_input_model_export(self):
|
|
model = MockMultiInputModel()
|
|
inputs = [Tensor.rand(2,2), Tensor.rand(2,2), Tensor.rand(2,2)]
|
|
prg, inp_sizes, _, _ = export_model(model, "", *inputs)
|
|
prg = json.loads(prg)
|
|
|
|
assert len(inputs) == len(prg["inputs"]) == len(inp_sizes), f"Model and exported inputs don't match: mdl={len(inputs)}, prg={len(prg['inputs'])}, inp_sizes={len(inp_sizes)}" # noqa: E501
|
|
|
|
for i in range(len(inputs)):
|
|
assert f"input{i}" in inp_sizes, f"input{i} not captured in inp_sizes"
|
|
assert f"input{i}" in prg["buffers"], f"input{i} not captured in exported buffers"
|
|
|
|
for i, exported_input in enumerate(prg["inputs"]):
|
|
assert inputs[i].dtype.name == exported_input["dtype"], f"Model and exported input dtype don't match: mdl={inputs[i].dtype.name}, prg={exported_input['dtype']}" # noqa: E501
|
|
|
|
def test_multi_output_model_export(self):
|
|
model = MockMultiOutputModel()
|
|
input_tensor = Tensor.rand(2,2)
|
|
outputs = model(input_tensor)
|
|
prg, _, out_sizes, _ = export_model(model, "", input_tensor)
|
|
prg = json.loads(prg)
|
|
|
|
assert len(outputs) == len(prg["outputs"]) == len(out_sizes), f"Model and exported outputs don't match: mdl={len(outputs)}, prg={len(prg['outputs'])}, inp_sizes={len(out_sizes)}" # noqa: E501
|
|
|
|
for i in range(len(outputs)):
|
|
assert f"output{i}" in out_sizes, f"output{i} not captured in out_sizes"
|
|
assert f"output{i}" in prg["buffers"], f"output{i} not captured in exported buffers"
|
|
|
|
for i, exported_output in enumerate(prg["outputs"]):
|
|
assert outputs[i].dtype.name == exported_output["dtype"], f"Model and exported output dtype don't match: mdl={outputs[i].dtype.name}, prg={exported_output['dtype']}" # noqa: E501
|
|
|
|
@unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Testing WebGPU specific model export behavior")
|
|
class TextModelExportWebGPU(unittest.TestCase):
|
|
def test_exported_input_output_dtypes(self):
|
|
class MyModel:
|
|
def forward(self, *inputs): return tuple([(inp+2).cast(inp.dtype) for inp in inputs])
|
|
model = MyModel()
|
|
# [:-1] because "ulong" and "long" is not supported
|
|
inputs = [Tensor.randn(2, dtype=dt) for dt in dtypes.uints[:-1] + dtypes.sints[:-1] + (dtypes.bool, dtypes.float)]
|
|
prg, _, _, _ = export_model(model, "webgpu", *inputs)
|
|
expected_buffer_types = ["Uint"]*len(dtypes.uints[:-1]) + ["Int"]*len(dtypes.sints[:-1]) + ["Int", "Float"]
|
|
for i, expected_buffer_type in enumerate(expected_buffer_types):
|
|
dt = inputs[i].dtype
|
|
expected_arr_prefix = f"{expected_buffer_type}{dt.itemsize*8}"
|
|
# test input buffers
|
|
self.assertIn(f"new {expected_arr_prefix}Array(gpuWriteBuffer{i}.getMappedRange()).set(_input{i});", prg)
|
|
# test output buffers
|
|
self.assertIn(f"const resultBuffer{i} = new {expected_arr_prefix}Array(gpuReadBuffer{i}.size/{dt.itemsize});", prg)
|
|
self.assertIn(f"resultBuffer{i}.set(new {expected_arr_prefix}Array(gpuReadBuffer{i}.getMappedRange()));", prg)
|
|
|
|
def test_weights_bound_to_safetensor(self):
|
|
# regression test: every weight ended up as createEmptyBuf (zero-init) instead of createWeightBuf
|
|
class MyModel:
|
|
def __init__(self): self.fc1, self.fc2 = Linear(4, 8), Linear(8, 2)
|
|
def forward(self, x): return self.fc2(self.fc1(x).relu())
|
|
model = MyModel()
|
|
for t in get_state_dict(model).values(): t.realize()
|
|
prg, _, _, _ = export_model(model, "webgpu", Tensor.randn(1, 4))
|
|
self.assertEqual(prg.count("createWeightBuf("), len(get_state_dict(model)))
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|