def test_avgpool2d(self): shape = (32,2,111,28) for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1), shape[2:]]: with self.subTest(kernel_size=ksz): helper_test_op([shape], lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz), lambda x: Tensor.avg_pool2d(x, kernel_size=ksz), gpu=self.gpu)
def test_avgpool2d_strided_fwd(self): for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]: for strd in [(1,1), (2,1), (2,2), (4,2)]: with self.subTest(kernel_size=ksz, stride=strd): helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, stride=strd), lambda x: Tensor.avg_pool2d(x, kernel_size=ksz, stride=strd), gpu=self.gpu, forward_only=True)
def test_avgpool2d(self): for ksz in [(2,2), (3,3), (3,2), (5,5), (5,1)]: for strd in [(1,1), (2,1), (2,2), (4,2)]: # TODO Grad tolerance needs to be slightly relaxed; why? with self.subTest(kernel_size=ksz, stride=strd): helper_test_op([(32,2,111,28)], lambda x: torch.nn.functional.avg_pool2d(x, kernel_size=ksz, stride=strd), lambda x: Tensor.avg_pool2d(x, kernel_size=ksz, stride=strd), gpu=self.gpu, grad_atol=1e-5)
def test_avgpool2d(self): 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, forward_only=self.gpu)