add flatten to tinygrad
MODIFY
requirements.txt
requirements.txt
@@ -9,3 +9,4 @@ torch==1.6.0+cpu
torchvision==0.7.0+cpu
pytest
tqdm
+onnx

MODIFY
test/test_openpilot_model.py
test/test_openpilot_model.py
@@ -27,7 +27,7 @@ def run_onnx(dat):
for inp in onnx_model.graph.initializer:
assert inp.data_type == 1
#print(inp.name, inp.dims, inp.data_type, len(inp.raw_data))
- tensors[inp.name] = Tensor(np.frombuffer(inp.raw_data, dtype=np.float32).reshape(inp.dims))
+ tensors[inp.name] = Tensor(np.frombuffer(inp.raw_data, dtype=np.float32).reshape(inp.dims).copy())
for num,n in enumerate(onnx_model.graph.node):
#print(f"{num}: op {n.op_type}")
@@ -45,10 +45,9 @@ def run_onnx(dat):
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":
- # TODO: add to tinygrad?
- ret = inp[0].reshape(shape=tuple(list(inp[0].shape[0:opt['axis']]) + [-1]))
+ ret = inp[0].flatten(opt['axis'])
elif n.op_type == "Concat":
- # terribly slow
+ # TODO: add multicat to tinygrad
ret = inp[0]
for x in inp[1:]:
ret = ret.cat(x, dim=opt['axis'])
@@ -72,13 +71,13 @@ def run_onnx(dat):
assert len(n.output) == 1
tensors[n.output[0]] = ret
- #for k,v in tensors.items():
- # print(k, v.shape)
+ return {outp.name:tensors[outp.name] for outp in onnx_model.graph.output}
class TestOpenpilotModel(unittest.TestCase):
def test(self):
dat = fetch("https://github.com/commaai/openpilot/raw/7da48ebdba5e3cf4c0b8078c934bee9a199f0280/selfdrive/modeld/models/supercombo.onnx")
- onnx_model = run_onnx(io.BytesIO(dat))
+ out = run_onnx(io.BytesIO(dat))
+ print(out)
if __name__ == "__main__":
unittest.main()

MODIFY
test/test_ops.py
test/test_ops.py
@@ -159,6 +159,10 @@ class TestOps(unittest.TestCase):
helper_test_op([(4,3,6,6)], lambda x: torch.reshape(x, (-1,3,6,6)), lambda x: x.reshape(shape=(-1,3,6,6)))
helper_test_op([(4,3,6,6)], lambda x: torch.reshape(x, (-1,1,6,6)), lambda x: x.reshape(shape=(-1,1,6,6)))
+ def test_flatten(self):
+ for axis in range(3):
+ helper_test_op([(4,3,6,6)], lambda x: torch.flatten(x, start_dim=axis), lambda x: x.flatten(axis))
+
def test_detach(self):
helper_test_op([(4,3,6,6)], lambda x: x.detach(), lambda x: x.detach(), forward_only=True)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -260,6 +260,9 @@ class Tensor:
def transpose(self, order=(1,0)):
return self.permute(order=order)
+ def flatten(self, dim=0):
+ return self.reshape(shape=tuple(list(self.shape[0:dim]) + [-1]))
+
def _canonicalize_reduce_axis(self, axis):
if axis is None: axis = range(len(self.shape))
if isinstance(axis, int): axis = [axis]

add flatten to tinygrad