multicat + test onnx is generic onnx
RENAME
test/test_openpilot_model.py
test/test_onnx.py
@@ -49,13 +49,8 @@ def run_onnx(onnx_model, inputs={}):
elif n.op_type == "Add": ret = inp[0] + inp[1]
elif n.op_type == "Sub": ret = inp[0] - inp[1]
elif n.op_type == "Mul": ret = inp[0] * inp[1]
- elif n.op_type == "Flatten":
- ret = inp[0].flatten(opt['axis'])
- elif n.op_type == "Concat":
- # TODO: add multicat to tinygrad
- ret = inp[0]
- for x in inp[1:]:
- ret = ret.cat(x, dim=opt['axis'])
+ elif n.op_type == "Flatten": ret = inp[0].flatten(opt['axis'])
+ elif n.op_type == "Concat": ret = inp[0].cat(*inp[1:], dim=opt['axis'])
elif n.op_type == "Split":
i = 0
arg = [(0,x) for x in inp[0].shape]
@@ -86,8 +81,8 @@ def run_onnx_torch(onnx_model, inputs):
torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
return torch_out
-class TestOpenpilotModel(unittest.TestCase):
- def test(self):
+class TestOnnxModel(unittest.TestCase):
+ def test_openpilot_model(self):
dat = fetch("https://github.com/commaai/openpilot/raw/7da48ebdba5e3cf4c0b8078c934bee9a199f0280/selfdrive/modeld/models/supercombo.onnx")
onnx_model = onnx.load(io.BytesIO(dat))
inputs = {
@@ -100,6 +95,7 @@ class TestOpenpilotModel(unittest.TestCase):
inputs = {k:v.astype(np.float32) for k,v in inputs.items()}
tinygrad_out = run_onnx(onnx_model, inputs)['outputs'].numpy()
torch_out = run_onnx_torch(onnx_model, inputs).numpy()
+ print(tinygrad_out, torch_out)
np.testing.assert_allclose(torch_out, tinygrad_out, atol=1e-4, rtol=1e-2)
if __name__ == "__main__":

MODIFY
test/test_ops.py
test/test_ops.py
@@ -260,7 +260,11 @@ class TestOps(unittest.TestCase):
def test_cat(self):
for dim in range(-1, 2):
- helper_test_op([(45,65), (45,65)], lambda x,y: torch.cat((x,y), dim), lambda x,y: x.cat(y, dim))
+ helper_test_op([(45,65), (45,65)], lambda x,y: torch.cat((x,y), dim), lambda x,y: x.cat(y, dim=dim))
+
+ def test_multicat(self):
+ for dim in range(-1, 2):
+ helper_test_op([(45,65), (45,65), (45,65)], lambda x,y,z: torch.cat((x,y,z), dim), lambda x,y,z: x.cat(y, z, dim=dim))
if __name__ == '__main__':
np.random.seed(1337)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -227,19 +227,27 @@ class Tensor:
ret = self.slice(arg = arg + [(0,self.shape[i]) for i in range(len(arg), len(self.shape))])
return ret.reshape(shape=new_shape) if tuple(ret.shape) != tuple(new_shape) else ret
- def cat(self, y, dim=0):
- assert len(self.shape) == len(y.shape)
+ def cat(self, *args, dim=0):
dim = (dim + len(self.shape)) if dim < 0 else dim
- s1, s2 = [], []
+ for y in args: assert len(self.shape) == len(y.shape)
+ args = [self] + list(args)
+ s = [[] for _ in range(len(args))]
for i in range(len(self.shape)):
if i != dim:
assert self.shape[i] == y.shape[i]
- s1.append((0, self.shape[i]))
- s2.append((0, self.shape[i]))
+ for j in range(len(args)):
+ s[j].append((0, self.shape[i]))
else:
- s1.append((0, self.shape[i]+y.shape[i]))
- s2.append((-self.shape[i], y.shape[i]))
- return self.slice(arg=s1) + y.slice(arg=s2)
+ shape_sum = 0
+ for y in args: shape_sum += y.shape[i]
+ k = 0
+ for j,y in enumerate(args):
+ s[j].append((-k, shape_sum-k))
+ k += y.shape[i]
+ ret = self.slice(arg=s[0])
+ for ts,y in zip(s[1:], args[1:]):
+ ret = ret + y.slice(arg=ts)
+ return ret
def pad2d(self, padding):
return self[:, :, -padding[2]:self.shape[2]+padding[3], -padding[0]:self.shape[3]+padding[1]]
@@ -428,7 +436,8 @@ def register(name, fxn):
setattr(Tensor, "_"+name if (getattr(Tensor, name, None) is not None) else name, dispatch)
if name in ['add', 'sub', 'mul', 'pow', 'matmul']:
setattr(Tensor, f"__{name}__", dispatch)
- setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(dispatch(self,x)))
+ # TODO: += is broken!
+ #setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(dispatch(self,x)))
setattr(Tensor, f"__r{name}__", lambda self,x: dispatch(x,self))
# register functions to move between devices

multicat + test onnx is generic onnx