extra and test and tuple
MODIFY
extra/onnx.py
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
\ No newline at end of file

MODIFY
test/test_conv.py
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)

MODIFY
test/test_onnx.py
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)

MODIFY
tinygrad/mlops.py
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)

MODIFY
tinygrad/tensor.py
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)}")

extra and test and tuple