fix torch convdw
MODIFY
test/test_ops.py
test/test_ops.py
@@ -160,6 +160,11 @@ class TestOps(unittest.TestCase):
def test_detach(self):
helper_test_op([(4,3,6,6)], lambda x: x.detach(), lambda x: x.detach(), forward_only=True)
+ def test_simple_conv2d(self):
+ helper_test_op([(1,1,9,9), (1,1,3,3)],
+ lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
+ lambda x,w: Tensor.conv2d(x,w).relu(), atol=1e-4, grad_rtol=1e-5)
+
def test_conv2d(self):
for bs in [1,8]:
for cin in [1,3]:
@@ -181,6 +186,24 @@ class TestOps(unittest.TestCase):
lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5)
+ def test_simple_grouped_conv2d(self):
+ bs = 1
+ groups = 2
+ rcout = 1
+ cin = 2
+ helper_test_op([(bs,groups*cin,1,1), (groups*rcout,cin,1,1)],
+ lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
+ lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5)
+
+ def test_medium_grouped_conv2d(self):
+ bs = 1
+ groups = 2
+ rcout = 2
+ cin = 2
+ helper_test_op([(bs,groups*cin,1,1), (groups*rcout,cin,1,1)],
+ lambda x,w: torch.nn.functional.conv2d(x,w,groups=groups).relu(),
+ lambda x,w: Tensor.conv2d(x,w,groups=groups).relu(), atol=1e-4, grad_rtol=1e-5)
+
def test_grouped_conv2d(self):
bs = 4
groups = 5

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
import numpy as np
from collections import namedtuple
+def prod(x): return int(np.prod(x))
+
def binary_broadcast(x_shape, y_shape, extra=False):
n_dims = max(len(x_shape), len(y_shape))
shape_x, shape_y = np.ones(n_dims, dtype=np.int32), np.ones(n_dims, dtype=np.int32)

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
import torch
import numpy as np
+from tinygrad.helpers import get_conv_args
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class TorchBuffer(torch.Tensor):
@@ -27,8 +28,16 @@ def conv(x,w,ret,stride,groups):
ret[:] = torch.nn.functional.conv2d(x, w, stride=stride, groups=groups)
return ret
-def convdw(x,grad_output,dw,stride,groups):
- dw[:] = torch.nn.grad.conv2d_weight(x, dw.shape, grad_output, stride=stride, groups=groups)
+def convdw(input,grad_output,dw,stride,groups):
+ # NOTE: torch.nn.grad.conv2d_weight is wrong for groups in pytorch, wonder who it affects
+ C = get_conv_args(input.shape, dw.shape, stride, groups)
+ grad_output = grad_output.reshape(C.bs, C.groups, C.rcout, C.oy, C.ox).repeat(1, 1, C.cin, 1, 1)
+ grad_output = grad_output.reshape(C.bs * C.groups * C.rcout * C.cin, 1, C.oy, C.ox)
+ input = input.reshape(1, C.bs * C.groups * C.cin, C.iy, C.ix)
+ grad_weight = torch.nn.functional.conv2d(input, grad_output, dilation=stride, groups=C.bs*C.groups*C.cin)
+ grad_weight = grad_weight.reshape(C.bs,-1).sum(dim=0)
+ grad_weight = grad_weight.view(C.groups, C.cin, C.rcout, C.H, C.W).transpose(2, 1)
+ dw[:] = grad_weight.contiguous().view(C.groups*C.rcout, C.cin, C.H, C.W)
return dw
def convdx(w,grad_output,dx,stride,groups):

fix torch convdw