From e46d122f653217d265beead3cf20aac4922e8271 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 9 Nov 2020 15:06:58 -0800 Subject: [PATCH] not supporting stride --- test/test_ops.py | 6 +++--- tinygrad/opsgpu.py | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/test/test_ops.py b/test/test_ops.py index 6ef1da3351..b550376e04 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -63,7 +63,7 @@ class TestOps(unittest.TestCase): helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.logsoftmax, atol=1e-7, grad_atol=1e-7, gpu=self.gpu) def test_pad2d(self): - helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,1,1,1)), lambda x: x.pad2d(padding=(1,1,1,1)), gpu=self.gpu) + helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,1,1,1)), lambda x: x.pad2d(padding=(1,1,1,1)), gpu=self.gpu, forward_only=self.gpu) def test_conv2d(self): for bs in [1,8]: @@ -98,11 +98,11 @@ class TestOps(unittest.TestCase): lambda x: Tensor.max_pool2d(x, kernel_size=ksz), gpu=self.gpu, forward_only=self.gpu) def test_avgpool2d(self): - for ksz in [(2,2), ]:#, (3,3), (3,2), (5,5), (5,1)]: + for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]: with self.subTest(kernel_size=ksz): helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz), - lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), gpu=self.gpu) + lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), gpu=self.gpu, forward_only=self.gpu) if GPU: class TestOpsGPU(TestOps): diff --git a/tinygrad/opsgpu.py b/tinygrad/opsgpu.py index bf10a2e93d..194a396a3e 100644 --- a/tinygrad/opsgpu.py +++ b/tinygrad/opsgpu.py @@ -382,9 +382,8 @@ register('sigmoid', Sigmoid, gpu=True) class AvgPool2D(Function): @staticmethod - def forward(ctx, input, kernel_size=(2, 2), stride=None): - if not stride: - ctx.stride = stride = kernel_size + def forward(ctx, input, kernel_size=(2, 2)): + ctx.stride = stride = ctx.kernel_size iter_op = "group_res += input[iid]" result_op = "group_res / (kernel_size.x * kernel_size.y)" ret = subsample_op(ctx, input, kernel_size, stride, iter_op, result_op) @@ -393,6 +392,8 @@ class AvgPool2D(Function): @staticmethod def backward(ctx, grad_output): + raise NotImplementedError("GPU AvgPool2D.backward() is broken") + # TODO implement for stride != kernel_size if ctx.kernel_size != ctx.stride: raise NotImplementedError("GPU AvgPool2D.backward() with stride != kernel_size not implemented") @@ -403,9 +404,8 @@ register('avg_pool2d', AvgPool2D, gpu=True) class MaxPool2D(Function): @staticmethod - def forward(ctx, input, kernel_size=(2, 2), stride=None): - if not stride: - ctx.stride = stride = kernel_size + def forward(ctx, input, kernel_size=(2, 2)): + ctx.stride = stride = ctx.kernel_size init_val = "FLT_MIN" iter_op = "group_res = max(group_res, input[iid])" result_op = "group_res"