write forward pass for convolution

test/test.py

@@ -1,34 +1,46 @@
 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)
 
-def test_tinygrad():
-  x = Tensor(x_init)
-  W = Tensor(W_init)
-  m = Tensor(m_init)
-  out = x.dot(W).relu()
-  out = out.logsoftmax()
-  out = out.mul(m).add(m).sum()
-  out.backward()
-  return out.data, x.grad, W.grad
+class TestTinygrad(unittest.TestCase):
+  def test_backward_pass(self):
+    def test_tinygrad():
+      x = Tensor(x_init)
+      W = Tensor(W_init)
+      m = Tensor(m_init)
+      out = x.dot(W).relu()
+      out = out.logsoftmax()
+      out = out.mul(m).add(m).sum()
+      out.backward()
+      return out.data, x.grad, W.grad
 
-def test_pytorch():
-  x = torch.tensor(x_init, requires_grad=True)
-  W = torch.tensor(W_init, requires_grad=True)
-  m = torch.tensor(m_init)
-  out = x.matmul(W).relu()
-  out = torch.nn.functional.log_softmax(out, dim=1)
-  out = out.mul(m).add(m).sum()
-  out.backward()
-  return out.detach().numpy(), x.grad, W.grad
+    def test_pytorch():
+      x = torch.tensor(x_init, requires_grad=True)
+      W = torch.tensor(W_init, requires_grad=True)
+      m = torch.tensor(m_init)
+      out = x.matmul(W).relu()
+      out = torch.nn.functional.log_softmax(out, dim=1)
+      out = out.mul(m).add(m).sum()
+      out.backward()
+      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)
-  
+    for x,y in zip(test_tinygrad(), test_pytorch()):
+      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()
 

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)
+