extra and test and tuple

extra/onnx.py

@@ -43,6 +43,7 @@ def get_run_onnx(onnx_model):
 
   def run_onnx(inputs={}, debug=False):
     input_tensors = {}
+    intermediate_tensors = {}
 
     # get inputs
     for inp in onnx_model.graph.input:
@@ -63,7 +64,7 @@ def get_run_onnx(onnx_model):
     conv_count = 0
     for num,n in enumerate(onnx_model.graph.node):
       if debug: print(f"{num}: op {n.op_type}")
-      inp = [tensors[x] if x in tensors else input_tensors[x] for x in n.input]
+      inp = [tensors[x] if x in tensors else (intermediate_tensors[x] if x in intermediate_tensors else input_tensors[x]) for x in n.input]
       opt = attribute_to_dict(n.attribute)
 
       # free ones
@@ -110,7 +111,7 @@ def get_run_onnx(onnx_model):
         arg = [(0,x) for x in inp[0].shape]
         for o,s in zip(n.output, opt['split']):
           arg[opt['axis']] = (i,i+s)
-          tensors[o] = inp[0].slice(arg=arg)
+          intermediate_tensors[o] = inp[0].slice(arg=arg)
           i = i+s
         continue
       elif n.op_type == "AveragePool":
@@ -136,8 +137,8 @@ def get_run_onnx(onnx_model):
         raise Exception(f"op_type {n.op_type} not supported")
       assert len(n.output) == 1
       if debug: print(ret.shape)
-      tensors[n.output[0]] = ret
+      intermediate_tensors[n.output[0]] = ret
       #print(ret.numpy().mean())
 
-    return {outp.name:tensors[outp.name] for outp in onnx_model.graph.output}
+    return {outp.name:intermediate_tensors[outp.name] for outp in onnx_model.graph.output}
   return run_onnx

test/test_conv.py

@@ -29,35 +29,42 @@ class TestConv(unittest.TestCase):
     print(ret.numpy())
 
   def test_two_binops_no_rerun(self):
+    Tensor.no_grad = True
     x = Tensor.ones(1,12,128,256)
     w = Tensor.ones(32,12,3,3)
     out = x.conv2d(w, stride=(2,2), padding=(1,1))
     out.relu().numpy(), (out-1).numpy()
+    Tensor.no_grad = False
     # TODO: make this a real test
 
   def test_two_overlapping_binops_no_rerun(self):
+    Tensor.no_grad = True
     x = Tensor.ones(1,12,128,256)
     w = Tensor.ones(32,12,3,3)
     out = x.conv2d(w, stride=(2,2), padding=(1,1))
     out.relu().numpy(), out.elu().numpy()
     # TODO: make this a real test
+    Tensor.no_grad = False
 
   def test_first_three(self):
+    Tensor.no_grad = True
     x = Tensor.ones(1,12,128,256)
 
     w = Tensor.ones(32,12,3,3)
-    x = x.conv2d(w, stride=(2,2), padding=(1,1))
+    x = x.conv2d(w, stride=(2,2), padding=(1,1)).elu()
 
     w = Tensor.ones(32,1,3,3)
-    x = x.conv2d(w, padding=(1,1), groups=32)
+    x = x.conv2d(w, padding=(1,1), groups=32).elu()
 
     w = Tensor.ones(16,32,1,1)
-    x = x.conv2d(w)
+    x = x.conv2d(w).elu()
 
     x = x.numpy()
     print(x.shape)
+    Tensor.no_grad = False
 
   def test_elu(self):
+    Tensor.no_grad = True
     x = Tensor.ones(1,12,128,256)
 
     w = Tensor.ones(32,12,3,3)
@@ -68,16 +75,18 @@ class TestConv(unittest.TestCase):
     w = Tensor.ones(32,1,3,3)
     x = x.conv2d(w, padding=(1,1), groups=32)
     out = x.numpy()
+    Tensor.no_grad = False
 
   def test_bias(self):
+    Tensor.no_grad = True
     from tinygrad.nn import Conv2d
     x = Tensor.ones(1,12,128,256)
     c = Conv2d(12, 32, 3)
-    x = c(x)
-    x = x.relu()
+    x = c(x).relu()
     w = Tensor.uniform(32, 1, 3, 3)
     x = x.conv2d(w, groups=32)
     out = x.numpy()
+    Tensor.no_grad = False
   
   def test_multiadd(self):
     w = Tensor.ones(32)

test/test_onnx.py

@@ -66,6 +66,7 @@ class TestOnnxModel(unittest.TestCase):
     ps.print_stats(30)
 
   def test_openpilot_model(self):
+    Tensor.no_grad = True
     dat = fetch(OPENPILOT_MODEL)
     onnx_model = onnx.load(io.BytesIO(dat))
     run_onnx = get_run_onnx(onnx_model)

tinygrad/mlops.py

@@ -139,8 +139,8 @@ class Permute(Function):
 # TODO: merge Slice and Flip into Stride with the 3 arguments
 class Slice(Function):
   def forward(ctx, x, arg=None):
-    ctx.narg = [(0-p[0], x.shape[i]-p[0]) for i,p in enumerate(arg)]
-    return x.movement_op(MovementOps.SLICE, arg)
+    ctx.narg = tuple((0-p[0], x.shape[i]-p[0]) for i,p in enumerate(arg))
+    return x.movement_op(MovementOps.SLICE, tuple(arg))
 
   def backward(ctx, grad_output):
     return grad_output.movement_op(MovementOps.SLICE, ctx.narg)

tinygrad/tensor.py

@@ -11,7 +11,7 @@ from tinygrad.ops import LazyBuffer
 # **** start with two base classes, Tensor and Function ****
 
 class Tensor:
-  training = False
+  training, no_grad = False, False
 
   def __init__(self, data, device=Device.DEFAULT, requires_grad=True):
     if isinstance(data, list):
@@ -335,7 +335,7 @@ class Function:
     self.device = device
     self.parents = tensors
     self.needs_input_grad = [t.requires_grad for t in tensors]
-    self.requires_grad = any(self.needs_input_grad)
+    self.requires_grad = any(self.needs_input_grad) and not Tensor.no_grad
     self.saved_tensors = []
 
   def forward(self, *args, **kwargs): raise NotImplementedError(f"forward not implemented for {type(self)}")