[mkldnn][pooling] support ceil mode by padding changes

Aapo Kyrola · Aapo Kyrola · commit df0427a70775 · 2019-06-06T10:04:49.000-07:00
diff --git a/aten/src/ATen/native/mkldnn/Pooling.cpp b/aten/src/ATen/native/mkldnn/Pooling.cpp
@@ -73,29 +73,69 @@ static Tensor _mkldnn_pool2d(
     IntArrayRef dilation,
     bool ceil_mode,
     ideep::algorithm algo) {
-  TORCH_CHECK(!ceil_mode, "Currently Mkldnn Pooling operators do not support ceil_mode.");
   auto kernel_size_vec = expand_param_if_needed(kernel_size, "kernel_size", 2);
   auto stride_vec = expand_param_if_needed(stride, "stride", 2);
   auto padding_vec = expand_param_if_needed(padding, "padding", 2);
+  auto padding_vec_l = padding_vec;
+  auto padding_vec_r = padding_vec;
   auto dilation_vec = expand_param_if_needed(dilation, "dilation", 2);
 
   const ideep::tensor& x = itensor_from_mkldnn(input);
-  const std::vector<int64_t> output_sizes = pool_output_sizes(
-      input.sizes(),
-      kernel_size_vec,
-      stride_vec,
-      padding_vec,
-      dilation_vec,
-      ceil_mode);
+  std::vector<int64_t> output_sizes;
+
+  if (ceil_mode) {
+    // MKLDNN does not support ceil mode, so we adjust padding
+    // on the right side to match behavior. Adjust output size
+    // accordingly.
+    const std::vector<int64_t> output_sizes_ceil = pool_output_sizes(
+        input.sizes(),
+        kernel_size_vec,
+        stride_vec,
+        padding_vec_l,
+        padding_vec_r,
+        dilation_vec,
+        true /* ceil_mode */);
+
+    // adjust padding until output sizes agree
+    bool all_equal = false;
+    while (!all_equal) {
+      output_sizes = pool_output_sizes(
+          input.sizes(),
+          kernel_size_vec,
+          stride_vec,
+          padding_vec_l,
+          padding_vec_r,
+          dilation_vec,
+          false /*ceil_mode */);
+
+      all_equal = true;
+      for (size_t i = 2; i < input.sizes().size(); ++i) {
+        if (output_sizes[i] < output_sizes_ceil[i]) {
+           padding_vec_r[i - 2]++;
+           all_equal = false;
+        }
+      }
+    }
+  } else {
+    output_sizes = pool_output_sizes(
+        input.sizes(),
+        kernel_size_vec,
+        stride_vec,
+        padding_vec_l,
+        padding_vec_r,
+        dilation_vec,
+        false /*ceil_mode */);
+  }
+
   ideep::tensor y;
   ideep::pooling_forward::compute<AllocForMKLDNN>(
       x,
       {output_sizes.cbegin(), output_sizes.cend()},
       y,
       {stride_vec.cbegin(), stride_vec.cend()},
       {kernel_size_vec.cbegin(), kernel_size_vec.cend()},
-      {padding_vec.cbegin(), padding_vec.cend()},
-      {padding_vec.cbegin(), padding_vec.cend()},
+      {padding_vec_l.cbegin(), padding_vec_l.cend()},
+      {padding_vec_r.cbegin(), padding_vec_r.cend()},
       algo,
       ideep::prop_kind::forward);
 
diff --git a/aten/src/ATen/native/mkldnn/Utils.cpp b/aten/src/ATen/native/mkldnn/Utils.cpp
@@ -25,19 +25,21 @@ std::vector<int64_t> pool_output_sizes(
     IntArrayRef input_size,
     IntArrayRef kernel_size,
     IntArrayRef stride,
-    IntArrayRef padding,
+    IntArrayRef padding_l,
+    IntArrayRef padding_r,
     IntArrayRef dilation,
     bool ceil_mode) {
   std::vector<int64_t> output_size(input_size.size());
   // copy N and C
   output_size[0] = input_size[0];
   output_size[1] = input_size[1];
 
-  for (int i = 2; i < input_size.size(); ++i) {
-    output_size[i] = pooling_output_shape<int64_t>(
+  for (size_t i = 2; i < input_size.size(); ++i) {
+    output_size[i] = pooling_output_shape_pad_lr<int64_t>(
       input_size[i],
       kernel_size[i - 2],
-      padding[i - 2],
+      padding_l[i - 2],
+      padding_r[i - 2],
       stride[i - 2],
       dilation[i - 2],
       ceil_mode
diff --git a/aten/src/ATen/native/mkldnn/Utils.h b/aten/src/ATen/native/mkldnn/Utils.h
@@ -16,7 +16,8 @@ std::vector<int64_t> pool_output_sizes(
     IntArrayRef input_size,
     IntArrayRef kernel_size,
     IntArrayRef stride,
-    IntArrayRef padding,
+    IntArrayRef padding_l,
+    IntArrayRef padding_r,
     IntArrayRef dilation,
     bool ceil_mode);
 }}
diff --git a/aten/src/THNN/generic/pooling_shape.h b/aten/src/THNN/generic/pooling_shape.h
@@ -2,16 +2,26 @@
 #define THNN_POOLING_SHAPE_H
 
 template<typename T>
-static inline T pooling_output_shape(
-        T inputSize, T kernelSize, T pad, T stride, T dilation, bool ceil_mode) {
-    T outputSize = ((inputSize + 2 * pad - dilation * (kernelSize - 1) - 1 + (ceil_mode ? stride - 1 : 0)) / stride + 1);
-    if (pad) {
+static inline T pooling_output_shape_pad_lr(
+    T inputSize, T kernelSize, T pad_l, T pad_r, T stride, T dilation,
+    bool ceil_mode
+  ) {
+    T outputSize = ((inputSize + pad_l + pad_r - dilation * (kernelSize - 1)
+        - 1 + (ceil_mode ? stride - 1 : 0)) / stride + 1);
+    if (pad_l) {
         // ensure that the last pooling starts inside the image
         // needed to avoid problems in ceil mode
-        if ((outputSize - 1) * stride >= inputSize + pad)
+        if ((outputSize - 1) * stride >= inputSize + pad_l)
           --outputSize;
     }
     return outputSize;
 }
 
+template<typename T>
+static inline T pooling_output_shape(
+      T inputSize, T kernelSize, T pad, T stride, T dilation, bool ceil_mode) {
+    return pooling_output_shape_pad_lr(
+        inputSize, kernelSize, pad, pad, stride, dilation, ceil_mode);
+}
+
 #endif
diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -126,16 +126,21 @@ def test_relu_(self):
     def test_max_pool2d(self):
         N = torch.randint(3, 10, (1,)).item()
         C = torch.randint(3, 10, (1,)).item()
-        x = torch.randn(N, C, 64, 64, dtype=torch.float32) * 10
 
-        max_pool2d = torch.nn.MaxPool2d(
-            kernel_size=3,
-            stride=2,
-            padding=1)
+        for stride in [1, 2, 3]:
+            for H, W in [(64, 64), (35, 39), (16, 19), [7, 8]]:
+                x = torch.randn(N, C, H, W, dtype=torch.float32) * 10
 
-        self.assertEqual(
-            max_pool2d(x),
-            max_pool2d(x.to_mkldnn()).to_dense())
+                for ceil_mode in [False, True]:
+                    max_pool2d = torch.nn.MaxPool2d(
+                        kernel_size=3 if not ceil_mode else 7,
+                        stride=stride,
+                        padding=1,
+                        ceil_mode=ceil_mode)
+
+                    self.assertEqual(
+                        max_pool2d(x),
+                        max_pool2d(x.to_mkldnn()).to_dense())
 
     def test_avg_pool2d(self):
         N = torch.randint(3, 10, (1,)).item()
diff --git a/third_party/fbgemm b/third_party/fbgemm
@@ -1 +1 @@
-Subproject commit 8197494f3ae280941639c72bc1342a9faa8e2ad6
+Subproject commit 85f4105cebfe538905dd167db1e2355f9637d8a1
diff --git a/third_party/onnx b/third_party/onnx
@@ -1 +1 @@
-Subproject commit 5160f3ac3380302224998f1c95e111cd961c4bc5
+Subproject commit 9005291283e943f1a91da5f0acf218bc4e8eb2ca
