This commit is contained in:
George Hotz
2020-10-25 19:03:49 -07:00
parent 2eebbd32c6
commit 2d37fd686b
2 changed files with 4 additions and 4 deletions

View File

@@ -5,7 +5,7 @@ import timeit
import functools
from tinygrad.tensor import Tensor
def test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-7, grad_atol=1e-7):
def helper_test_op(shps, torch_fxn, tinygrad_fxn, atol=1e-7, grad_atol=1e-7):
ts = [torch.rand(x, requires_grad=True) for x in shps]
tst = [Tensor(x.detach().numpy()) for x in ts]
@@ -36,15 +36,15 @@ class TestOps(unittest.TestCase):
for cin in [1,3]:
for H in [2,5]:
for W in [2,3,5]:
test_op([(bs,cin,11,28), (4,cin,H,W)],
helper_test_op([(bs,cin,11,28), (4,cin,H,W)],
lambda x,w: torch.nn.functional.conv2d(x,w).relu(),
lambda x,w: Tensor.conv2d(x,w).relu(), atol=2e-5, grad_atol=2e-6)
def test_maxpool2x2(self):
test_op([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, (2,2)), Tensor.max_pool2d)
helper_test_op([(32,2,110,28)], lambda x: torch.nn.functional.max_pool2d(x, (2,2)), Tensor.max_pool2d)
def test_avgpool2x2(self):
test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, (2,2)), Tensor.avg_pool2d)
helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, (2,2)), Tensor.avg_pool2d)
if __name__ == '__main__':
unittest.main(verbosity=2)