From 2720ef49ca892d99bae381505bb7fc808bf664c5 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 7 Jul 2022 10:01:33 -0700 Subject: [PATCH] extra and test and tuple --- extra/onnx.py | 9 +++++---- test/test_conv.py | 19 ++++++++++++++----- test/test_onnx.py | 1 + tinygrad/mlops.py | 4 ++-- tinygrad/tensor.py | 4 ++-- 5 files changed, 24 insertions(+), 13 deletions(-) diff --git a/extra/onnx.py b/extra/onnx.py index bf0a559d54..82f2997316 100644 --- a/extra/onnx.py +++ b/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 diff --git a/test/test_conv.py b/test/test_conv.py index 7f791a5b88..62ec110a2c 100644 --- a/test/test_conv.py +++ b/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) diff --git a/test/test_onnx.py b/test/test_onnx.py index 1ef908be27..e617b2fe8f 100644 --- a/test/test_onnx.py +++ b/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) diff --git a/tinygrad/mlops.py b/tinygrad/mlops.py index 2dfbc127ff..8535eeda0e 100644 --- a/tinygrad/mlops.py +++ b/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) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 2d744a7df8..f53a838e3c 100644 --- a/tinygrad/tensor.py +++ b/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)}")