From bb98cdfef78d52be96e21671a3d47f45ec1bd687 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sun, 25 Oct 2020 12:46:04 -0700 Subject: [PATCH] improve conv testing --- test/test_mnist.py | 1 - test/test_tensor.py | 25 ++++++++++++++----------- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/test/test_mnist.py b/test/test_mnist.py index a94e18ceed..e6ccf2f7e4 100644 --- a/test/test_mnist.py +++ b/test/test_mnist.py @@ -74,7 +74,6 @@ def evaluate(model): assert accuracy > 0.95 class TestMNIST(unittest.TestCase): - @unittest.skip(reason="mad slow") def test_conv(self): np.random.seed(1337) model = TinyConvNet() diff --git a/test/test_tensor.py b/test/test_tensor.py index 5f284bbf53..9d7bd845e2 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -66,20 +66,23 @@ class TestTinygrad(unittest.TestCase): class TestOps(unittest.TestCase): def test_conv2d(self): - x = torch.randn((5,2,10,7), requires_grad=True) - w = torch.randn((4,2,3,2), requires_grad=True) - xt = Tensor(x.detach().numpy()) - wt = Tensor(w.detach().numpy()) + for cin in [1,2,3]: + for H in [2,3,5]: + for W in [2,3,5]: + x = torch.randn((5,cin,10,7), requires_grad=True) + w = torch.randn((4,cin,H,W), requires_grad=True) + xt = Tensor(x.detach().numpy()) + wt = Tensor(w.detach().numpy()) - out = torch.nn.functional.conv2d(x,w) - ret = Tensor.conv2d(xt, wt) - np.testing.assert_allclose(ret.data, out.detach().numpy(), atol=1e-5) + out = torch.nn.functional.conv2d(x,w) + ret = Tensor.conv2d(xt, wt) + np.testing.assert_allclose(ret.data, out.detach().numpy(), atol=1e-5) - out.mean().backward() - ret.mean().backward() + out.mean().backward() + ret.mean().backward() - np.testing.assert_allclose(w.grad, wt.grad, atol=1e-7) - np.testing.assert_allclose(x.grad, xt.grad, atol=1e-7) + np.testing.assert_allclose(w.grad, wt.grad, atol=1e-7) + np.testing.assert_allclose(x.grad, xt.grad, atol=1e-7) def test_maxpool2x2(self): x = torch.randn((5,2,10,8), requires_grad=True)