support dilated convolution in torch
MODIFY
test/test_ops.py
test/test_ops.py
@@ -268,6 +268,20 @@ class TestOps(unittest.TestCase):
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)
+ @unittest.skipUnless(Device.DEFAULT == Device.TORCH, "Not Implemented")
+ def test_dilated_conv2d(self):
+ bs = 4
+ cin = 3
+ H,W = 3,3
+ with self.subTest(dilation := 2):
+ helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
+ lambda x,w: torch.nn.functional.conv2d(x,w,dilation=dilation).relu(),
+ lambda x,w: Tensor.conv2d(x,w,dilation=dilation).relu(), atol=1e-4)
+ with self.subTest(dilation := (2,1)):
+ helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
+ lambda x,w: torch.nn.functional.conv2d(x,w,dilation=dilation).relu(),
+ lambda x,w: Tensor.conv2d(x,w,dilation=dilation).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,16 +6,18 @@ 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=0, groups=1, padding=0):
+def get_conv_args(x_shape, w_shape, stride=1, groups=1, padding=0, dilation=1):
# 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', 'py', 'px'])
+ ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs', 'cout', 'py', 'px', 'dy','dx'])
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
+ dy,dx = (dilation, dilation) if isinstance(dilation, int) else dilation
bs,cin_,iy,ix = x_shape
# 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
+ oy = (iy + 2*py - H - (H-1)*(dy-1))//ys + 1
+ ox = (ix + 2*px - W - (W-1)*(dx-1))//xs + 1
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, py, px)
+ return conv_args(H, W, groups, cout//groups, cin, oy, ox, iy, ix, ys, xs, bs, cout, py, px, dy, dx)

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -32,21 +32,21 @@ def convdw(x,grad_output,dw,C):
grad_output = grad_output.reshape(C.bs * C.groups * C.cin * C.rcout, 1, C.oy, C.ox)
x = x.reshape(1, C.bs * C.groups * C.cin, C.iy, C.ix)
# NOTE: this conv2d always has batch size 1.
- grad_weight = torch.conv2d(x, grad_output, dilation=(C.ys, C.xs), groups=C.bs*C.groups*C.cin)
+ grad_weight = torch.conv2d(x, grad_output, stride=(C.dy, C.dx), dilation=(C.ys, C.xs), groups=C.bs*C.groups*C.cin)
grad_weight = grad_weight.reshape(C.bs, C.groups, C.cin, C.rcout, *grad_weight.shape[2:]).transpose(3, 2).sum(dim=0)
dw[:] = grad_weight.reshape(C.groups*C.rcout, C.cin, *grad_weight.shape[3:])[:, :, :dw.shape[2], :dw.shape[3]]
def processing_op(op,x,w,ret,C):
- stride, groups = (C.ys, C.xs), C.groups
+ stride, groups, dilation = (C.ys, C.xs), C.groups, (C.dy, C.dx)
if op == ProcessingOps.CONV:
- ret[:] = torch.conv2d(x, w, stride=stride, groups=groups)
+ ret[:] = torch.conv2d(x, w, stride=stride, groups=groups, dilation=dilation)
elif op == ProcessingOps.CONVT:
- if stride == 1 or stride == (1,1):
+ if stride == (1,1) and dilation == (1,1):
# strided needs weird "undilation": https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
w = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W).flip(3, 4).transpose(2, 1).reshape(C.groups*C.cin, C.rcout, C.H, C.W)
ret[:] = torch.conv2d(x, w, padding=(C.H-1,C.W-1), groups=groups)
else:
- output_padding = [ret.shape[d+2] - ((x.shape[d+2] - 1) * stride[d] + 1 + (w.shape[d+2] - 1)) for d in range(2)]
- ret[:] = torch.conv_transpose2d(x, w, stride=stride, groups=groups, output_padding=output_padding)
+ output_padding = [ret.shape[d+2] - ((x.shape[d+2] - 1) * stride[d] + 1 + dilation[d] * (w.shape[d+2] - 1)) for d in range(2)]
+ ret[:] = torch.conv_transpose2d(x, w, stride=stride, groups=groups, output_padding=output_padding, dilation=dilation)
elif op == ProcessingOps.CONVDW:
convdw(x,w,ret,C)

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -167,14 +167,13 @@ class Slice(Function):
# ************* processing ops *************
class Conv2D(Function):
- def forward(ctx, x, w, stride=1, groups=1):
- C = get_conv_args(x.shape, w.shape, stride, groups)
- ctx.save_for_backward(x,w,(C.ys,C.xs), C.groups)
+ def forward(ctx, x, w, stride=1, groups=1, dilation=1):
+ C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation)
+ ctx.save_for_backward(x,w,C)
return ctx.processing_op(ProcessingOps.CONV, x, w, (C.bs, C.groups*C.rcout, C.oy, C.ox), C)
def backward(ctx, grad_output):
- x, w, stride, groups = ctx.saved_tensors
- C = get_conv_args(x.shape, w.shape, stride, groups)
+ x, w, C = ctx.saved_tensors
dx = ctx.processing_op(ProcessingOps.CONVT, grad_output, w, x.shape, C) if ctx.needs_input_grad[0] else None
dw = ctx.processing_op(ProcessingOps.CONVDW, x, grad_output, w.shape, C) if ctx.needs_input_grad[1] else None
return dx, dw
\ No newline at end of file

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -344,10 +344,10 @@ 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, padding=None):
+ def conv2d(self, weight, bias=None, padding=None, **kwargs):
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)
+ ret = x._conv2d(weight, **kwargs)
return ret if bias is None else ret.add(bias.reshape(shape=[1, -1, 1, 1]))
# ***** broadcasted binary ops *****

support dilated convolution in torch