add test for padding
MODIFY
test/test_ops.py
test/test_ops.py
@@ -255,6 +255,19 @@ class TestOps(unittest.TestCase):
lambda x,w: torch.nn.functional.conv2d(x,w,stride=stride).relu(),
lambda x,w: Tensor.conv2d(x,w,stride=(2,1)).relu(), atol=1e-4)
+ def test_padded_conv2d(self):
+ bs = 4
+ cin = 3
+ H,W = 3,3
+ with self.subTest(padding := 2):
+ helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
+ lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
+ lambda x,w: Tensor.conv2d(x,w,padding=padding).relu(), atol=1e-4)
+ with self.subTest(padding := (2,1)):
+ helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
+ lambda x,w: torch.nn.functional.conv2d(x,w,padding=padding).relu(),
+ lambda x,w: Tensor.conv2d(x,w,padding=padding).relu(), atol=1e-4)
+
def test_maxpool2d(self):
for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]:
with self.subTest(kernel_size=ksz):

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -6,14 +6,16 @@ def prod(x): return int(np.prod(x))
def reduce_shape(shape, axis):
return [1 if i in axis else shape[i] for i in range(len(shape))]
-def get_conv_args(x_shape, w_shape, stride, groups):
+def get_conv_args(x_shape, w_shape, stride=0, groups=1, padding=0):
# 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'])
+ ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs', 'cout', 'py', 'px'])
cout,cin,H,W = w_shape
ys,xs = (stride, stride) if isinstance(stride, int) else stride
+ py,px = (padding, padding) if isinstance(padding, int) else padding
bs,cin_,iy,ix = x_shape
- oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
+ # TODO: should be easy to support asymmetric padding by changing output size
+ oy,ox = (iy+py*2-(H-ys))//ys, (ix+px*2-(W-xs))//xs
if cin*groups != cin_: raise Exception(f"Input Tensor shape {x_shape} does not match the shape of the weights {w_shape}. ({cin*groups} vs. {cin_})")
assert cout % groups == 0
- return conv_args(H, W, groups, cout//groups, cin, oy, ox, iy, ix, ys, xs, bs, cout)
+ return conv_args(H, W, groups, cout//groups, cin, oy, ox, iy, ix, ys, xs, bs, cout, py, px)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -345,7 +345,8 @@ class Tensor:
return self._pool2d(*kernel_size).max(axis=(3,5))
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
+ if isinstance(padding, int): padding = (padding, padding)
+ x = self.pad2d((padding[1], padding[1], padding[0], padding[0])) 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]))

add test for padding