add padding to conv2d function, other minor things
ADD
test/test_conv.py
+import unittest
+from tinygrad.tensor import Tensor
+
+class TestConv(unittest.TestCase):
+ def test_simple(self):
+ x = Tensor.ones(1,12,128,256)
+ w = Tensor.ones(32,12,3,3)
+ ret = x.conv2d(w, padding=(1,1))
+ # it's not 108 around the padding
+ assert (ret.numpy()[:, :, 1:-1, 1:-1] == 108).all()
+
+if __name__ == '__main__':
+ unittest.main()
\ No newline at end of file

MODIFY
test/test_onnx.py
test/test_onnx.py
@@ -78,9 +78,12 @@ def run_onnx(onnx_model, inputs={}, debug=False):
elif n.op_type == "Conv":
x,w,b = inp if len(inp) == 3 else (inp[0], inp[1], None)
assert 'dilations' not in opt or opt['dilations'] == (1,1)
- # pads are in different order
- pads = (opt['pads'][0], opt['pads'][2], opt['pads'][1], opt['pads'][3])
- ret = x.pad2d(pads).conv2d(w, b, stride=opt['strides'], groups=opt['group'] if 'group' in opt else 1)
+ if opt['pads'][0] == opt['pads'][2] and opt['pads'][1] == opt['pads'][3]:
+ # symmetric padding
+ ret = x.conv2d(w, b, stride=opt['strides'], groups=opt.get('group', 1), padding=opt['pads'][0:2])
+ else:
+ x = x.pad2d((opt['pads'][0], opt['pads'][2], opt['pads'][1], opt['pads'][3]))
+ ret = x.conv2d(w, b, stride=opt['strides'], groups=opt.get('group', 1))
elif n.op_type in ["Add", "Sub", "Mul"]:
# TODO: add this to tinygrad? i don't think it's in torch
if len(inp[0].shape) != len(inp[1].shape) and prod(inp[0].shape) == prod(inp[1].shape):

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -28,6 +28,7 @@ def binary_broadcast(x_shape, y_shape, extra=False):
return (shape_ret, dimlist, complist) if extra else shape_ret
def get_conv_args(x_shape, w_shape, stride, groups):
+ # TODO: https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html#tensor-layout
conv_args = namedtuple('conv_args',
['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs'])
cout,cin,H,W = w_shape

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -5,6 +5,8 @@ from tinygrad.helpers import prod, binary_broadcast, get_conv_args
from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps
cl_ctx, cl_queue = None, None
+def get_cl_ctx(): return cl_ctx
+def get_cl_queue(): return cl_queue
def require_init_gpu():
global cl_ctx, cl_queue
if cl_ctx is None:

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -45,7 +45,7 @@ def log_op(op, ret, inp):
G.nodes[nm(ret)]['fillcolor'] = top_colors[top]
G.nodes[nm(ret)]['style'] = 'filled'
-from tinygrad.helpers import binary_broadcast, reduce_shape
+from tinygrad.helpers import binary_broadcast
class Ops:
def unary_op(ctx, op:UnaryOps, x):
ret = ctx.buffer(x.shape)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -344,8 +344,9 @@ class Tensor:
def max_pool2d(self, kernel_size=(2,2)):
return self._pool2d(*kernel_size).max(axis=(3,5))
- def conv2d(self, weight, bias=None, stride=1, groups=1):
- ret = self._conv2d(weight, stride=stride, groups=groups)
+ def conv2d(self, weight, bias=None, stride=1, groups=1, padding=None):
+ x = self.pad2d((padding[0], padding[0], padding[1], padding[1])) if padding is not None else self
+ ret = x._conv2d(weight, stride=stride, groups=groups)
return ret if bias is None else ret.add(bias.reshape(shape=[1, -1, 1, 1]))
# ***** functional nn ops *****

add padding to conv2d function, other minor things