diff --git a/aten/src/ATen/native/AdaptiveAveragePooling.cpp b/aten/src/ATen/native/AdaptiveAveragePooling.cpp index 94314f5c0ec8a..f632cf5f781ee 100644 --- a/aten/src/ATen/native/AdaptiveAveragePooling.cpp +++ b/aten/src/ATen/native/AdaptiveAveragePooling.cpp @@ -262,9 +262,11 @@ namespace { return output; } - Tensor adaptive_avg_pool2d( - at::Tensor const& input, - IntArrayRef output_size){ + Tensor adaptive_avg_pool2d(at::Tensor const& input, IntArrayRef output_size) { + if (input.is_mkldnn()) { + return at::mkldnn_adaptive_avg_pool2d(input, output_size); + } + if (output_size[0] == 1 && output_size[1] == 1) { //in this case, adaptive pooling is just computing mean over hw dimensions, which can be done more efficiently int64_t mean_size = input.size(-1) * input.size(-2); diff --git a/aten/src/ATen/native/mkldnn/Pooling.cpp b/aten/src/ATen/native/mkldnn/Pooling.cpp index f6ddaac8f9cdf..0e1b351bdde16 100644 --- a/aten/src/ATen/native/mkldnn/Pooling.cpp +++ b/aten/src/ATen/native/mkldnn/Pooling.cpp @@ -41,6 +41,19 @@ Tensor& mkldnn_avg_pool2d_out( bool count_include_pad) { AT_ERROR("mkldnn_avg_pool2d_out: ATen not compiled with MKLDNN support"); } + +Tensor mkldnn_adaptive_avg_pool2d(Tensor const& input, IntArrayRef output_size) { + AT_ERROR("mkldnn_adaptive_avg_pool2d: ATen not compiled with MKLDNN support"); +} + +Tensor& mkldnn_adaptive_avg_pool2d_out( + Tensor& output, + const Tensor& input, + IntArrayRef output_size) { + AT_ERROR( + "mkldnn_adaptive_avg_pool2d_out: ATen not compiled with MKLDNN support"); +} + } // namespace native } // namespace at @@ -118,7 +131,7 @@ Tensor mkldnn_avg_pool2d( kernel_size, stride, padding, - std::vector{1, 1}, + /*dilation*/ std::vector{1, 1}, ceil_mode, count_include_pad ? ideep::algorithm::pooling_avg_include_padding : ideep::algorithm::pooling_avg_exclude_padding); @@ -136,6 +149,42 @@ Tensor& mkldnn_avg_pool2d_out( "mkldnn_avg_pool2d_out: in-place mkldnn operations are not supported yet"); } +Tensor mkldnn_adaptive_avg_pool2d( + Tensor const& input, + IntArrayRef output_size) { + AT_ASSERTM(input.dim() == 4, "mkldnn_adaptive_avg_pool2d: Expect 2D input"); + + auto output_size_vec = + expand_param_if_needed(output_size, "output_size", input.dim() - 2); + std::vector kernel_size(input.dim() - 2); + for (size_t i = 2; i < input.dim(); ++i) { + auto s1 = input.size(i); + auto s2 = output_size_vec[i - 2]; + AT_ASSERTM(s2 != 0, "output size can not be zero"); + AT_ASSERTM( + s1 % s2 == 0, + "input size is not divisible by the output size is not supported yet"); + kernel_size[i - 2] = s1 / s2; + } + return _mkldnn_pool2d( + input, + kernel_size, + /*stride*/ kernel_size, + /*padding*/ {0, 0}, + /*dilation*/ {1, 1}, + /*ceil_mode*/ false, + /*algo*/ ideep::algorithm::pooling_avg); +} + +Tensor& mkldnn_adaptive_avg_pool2d_out( + Tensor& output, + const Tensor& input, + IntArrayRef output_size) { + AT_ERROR( + "mkldnn_adaptive_avg_pool2d_out: in-place mkldnn operations are not supported yet"); +} + + } // namespace native } // namespace at diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index d58acd0cc7427..2fec270eabbfb 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -3634,10 +3634,16 @@ dispatch: CPU: adaptive_avg_pool2d_out_cpu CUDA: adaptive_avg_pool2d_out_cuda + MkldnnCPU: mkldnn_adaptive_avg_pool2d_out - func: adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor python_module: nn +- func: mkldnn_adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor + dispatch: + MkldnnCPU: mkldnn_adaptive_avg_pool2d + requires_tensor: True + - func: _adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor dispatch: CPU: adaptive_avg_pool2d_cpu diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py index 0157bf95aa1a6..ea3ae4228aa98 100644 --- a/test/test_mkldnn.py +++ b/test/test_mkldnn.py @@ -148,6 +148,17 @@ def test_avg_pool2d(self): avg_pool2d(x), avg_pool2d(x.to_mkldnn()).to_dense()) + def test_adaptive_avg_pool2d(self): + N = torch.randint(3, 10, (1,)).item() + C = torch.randint(3, 10, (1,)).item() + x = torch.randn(N, C, 224, 224, dtype=torch.float32) * 100 + + adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d(7) + + self.assertEqual( + adaptive_avg_pool2d(x), + adaptive_avg_pool2d(x.to_mkldnn()).to_dense()) + def test_batch_norm2d(self): N = torch.randint(3, 10, (1,)).item() C = torch.randint(3, 100, (1,)).item()