maxpool

test/test_tensor.py

@@ -64,6 +64,7 @@ class TestTinygrad(unittest.TestCase):
     # coarse approx. since a "big" eps and the non-linearities of the model
     self.assertFalse(gradcheck(tiny_func, tiny_x, eps = 0.1))
 
+class TestOps(unittest.TestCase):
   def test_conv2d(self):
     x = torch.randn((5,2,10,7), requires_grad=True)
     w = torch.randn((4,2,3,3), requires_grad=True)
@@ -80,6 +81,22 @@ class TestTinygrad(unittest.TestCase):
     np.testing.assert_allclose(w.grad, wt.grad, atol=1e-5)
     np.testing.assert_allclose(x.grad, xt.grad, atol=1e-5)
 
+  def test_maxpool2x2(self):
+    x = torch.randn((5,2,10,8), requires_grad=True)
+    xt = Tensor(x.detach().numpy())
+
+    # in tinygrad
+    ret = xt.maxpool2x2()
+    assert ret.shape == (5,2,10//2,8//2)
+    ret.mean().backward()
+
+    # in torch
+    out = torch.nn.MaxPool2d((2,2))(x)
+    out.mean().backward()
+
+    # forward and backward the same
+    np.testing.assert_allclose(ret.data, out.detach().numpy(), atol=1e-5)
+    np.testing.assert_allclose(x.grad, xt.grad, atol=1e-5)
 
 if __name__ == '__main__':
   unittest.main()

tinygrad/ops.py

@@ -167,3 +167,27 @@ class FastConv2D(Function):
     return dx, dw
 register('conv2d', FastConv2D)
 
+# TODO: make this parameterizable
+class MaxPool2x2(Function):
+  @staticmethod
+  def forward(ctx, x):
+    stack = []
+    for Y in range(2):
+      for X in range(2):
+        stack.append(x[:, :, Y::2, X::2][None])
+    stack = np.concatenate(stack, axis=0)
+    idxs = np.argmax(stack, axis=0)
+    ctx.save_for_backward(idxs)
+    return np.max(stack, axis=0)
+
+  @staticmethod
+  def backward(ctx, grad_output):
+    idxs, = ctx.saved_tensors
+    s = grad_output.shape
+    ret = np.zeros((s[0], s[1], s[2]*2, s[3]*2), dtype=grad_output.dtype)
+    for Y in range(2):
+      for X in range(2):
+        ret[:, :, Y::2, X::2] = grad_output * (idxs == (Y*2+X))
+    return ret
+register('maxpool2x2', MaxPool2x2)
+