write forward pass for convolution
MODIFY
test/test.py
test/test.py
import numpy as np
import torch
-from tinygrad.tensor import Tensor
+import unittest
+from tinygrad.tensor import Tensor, Conv2D
x_init = np.random.randn(1,3).astype(np.float32)
W_init = np.random.randn(3,3).astype(np.float32)
m_init = np.random.randn(1,3).astype(np.float32)
+class TestTinygrad(unittest.TestCase):
+ def test_backward_pass(self):
def test_tinygrad():
x = Tensor(x_init)
W = Tensor(W_init)
@@ -27,8 +30,17 @@ def test_pytorch():
return out.detach().numpy(), x.grad, W.grad
for x,y in zip(test_tinygrad(), test_pytorch()):
- print(x,y)
np.testing.assert_allclose(x, y, atol=1e-5)
+ def test_conv2d(self):
+ x = torch.randn((5,2,10,7))
+ w = torch.randn((4,2,3,3))
+ out = torch.nn.functional.conv2d(x,w)
+ ret = Conv2D.apply(Conv2D, Tensor(x.numpy()), Tensor(w.numpy()))
+ np.testing.assert_allclose(ret.data, out.numpy(), atol=1e-5)
+
+
+if __name__ == '__main__':
+ unittest.main()

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -58,8 +58,15 @@ class Function:
# note that due to how partialmethod works, self and arg are switched
def apply(self, arg, *x):
- ctx = arg(self, *x)
- ret = Tensor(arg.forward(ctx, self.data, *[t.data for t in x]))
+ # support the args in both orders
+ if type(arg) == Tensor:
+ op = self
+ x = [arg]+list(x)
+ else:
+ op = arg
+ x = [self]+list(x)
+ ctx = op(*x)
+ ret = Tensor(op.forward(ctx, *[t.data for t in x]))
ret._ctx = ctx
return ret
@@ -147,3 +154,23 @@ class LogSoftmax(Function):
return grad_output - np.exp(output)*grad_output.sum(axis=1).reshape((-1, 1))
register('logsoftmax', LogSoftmax)
+class Conv2D(Function):
+ @staticmethod
+ def forward(ctx, x, w):
+ cout,cin,H,W = w.shape
+ ret = np.zeros((x.shape[0], cout, x.shape[2]-(H-1), x.shape[3]-(W-1)), dtype=w.dtype)
+ for Y in range(ret.shape[2]):
+ for X in range(ret.shape[3]):
+ for j in range(H):
+ for i in range(W):
+ for c in range(cout):
+ tx = x[:, :, Y+j, X+i]
+ tw = w[c, :, j, i]
+ ret[:, c, Y, X] += tx.dot(tw.reshape(-1, 1)).reshape(-1)
+ return ret
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ raise Exception("please write backward pass for Conv2D")
+register('conv2d', Conv2D)
+

write forward pass for convolution