Skip to content

Add mkldnn support for adaptive_avg_pool2d #19818

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
bddppq wants to merge 1 commit into from
Closed

Add mkldnn support for adaptive_avg_pool2d #19818

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions aten/src/ATen/native/AdaptiveAveragePooling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
51 changes: 50 additions & 1 deletion aten/src/ATen/native/mkldnn/Pooling.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -118,7 +131,7 @@ Tensor mkldnn_avg_pool2d(
kernel_size,
stride,
padding,
std::vector<int64_t>{1, 1},
/*dilation*/ std::vector<int64_t>{1, 1},
ceil_mode,
count_include_pad ? ideep::algorithm::pooling_avg_include_padding
: ideep::algorithm::pooling_avg_exclude_padding);
Expand All @@ -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<int64_t> 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

Expand Down
6 changes: 6 additions & 0 deletions aten/src/ATen/native/native_functions.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator

@wanchaol wanchaol Apr 30, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I might miss some context here. What's the meaning of this option?

Copy link
Contributor Author

@bddppq bddppq Apr 30, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes the codegen script to insert code to unwrap a variable to a tensor before passing the input to the kernel:
https://github.com/pytorch/pytorch/blob/de19eee/tools/autograd/gen_variable_type.py#L869-L870
https://github.com/pytorch/pytorch/blob/de19eee/tools/autograd/gen_variable_type.py#L954


- func: _adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor
dispatch:
CPU: adaptive_avg_pool2d_cpu
Expand Down
11 changes: 11 additions & 0 deletions test/test_mkldnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Copy link
Collaborator

@wanchaol wanchaol Apr 30, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it mean that everytime we call mkldnn ops we are required to call to_dense ?

Copy link
Contributor Author

@bddppq bddppq Apr 30, 2019

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not every op, but every chunk of ops (in this case the unittest is only testing one op so yeah it's calling to_dense after each op to convert the mkldnn tensor to dense tensor so that we can compare with the regular dense kernel's output).
basically one needs to do:

to_mkldnn -> mkldnn_op1 -> mkldnn_op2 ... -> mkldnn_opn -> to_dense


def test_batch_norm2d(self):
N = torch.randint(3, 10, (1,)).item()
C = torch.randint(3, 100, (1,)).item()
Expand Down