[GitHub] szha closed pull request #11047: [MXNET-33] Enhance mkldnn pooling to support full convention

[GitBox <gi...@apache.org>]

szha closed pull request #11047: [MXNET-33] Enhance mkldnn pooling to support full convention
URL: https://github.com/apache/incubator-mxnet/pull/11047
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
index 66679613d3a..f548778c761 100644
--- a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
+++ b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
@@ -113,20 +113,12 @@ inline bool SupportMKLDNNPooling(const PoolingParam &param,
   if (!ret)
     return false;
 
-  if (param.pooling_convention == pool_enum::kValid)
+  if (param.pooling_convention == pool_enum::kValid) {
     return true;
-  else
-    return false;
-
-// need to support pooling convention full
-// https://issues.apache.org/jira/browse/MXNET-33
-#if 0
-  if (((dshape[2] + 2 * param.pad[0] - param.kernel[0]) % param.stride[0] == 0) &&
-      ((dshape[3] + 2 * param.pad[1] - param.kernel[1]) % param.stride[1] == 0))
-    return true;
-  else
-    return false;
-#endif
+  } else {
+    // currently, only max-pooling is supported for full convention
+    return param.pool_type == pool_enum::kMaxPooling;
+  }
 }
 
 inline bool MKLDNNRequireWorkspace(const PoolingParam &param) {
diff --git a/src/operator/nn/mkldnn/mkldnn_pooling.cc b/src/operator/nn/mkldnn/mkldnn_pooling.cc
index 1610944304e..18dc835c0d0 100644
--- a/src/operator/nn/mkldnn/mkldnn_pooling.cc
+++ b/src/operator/nn/mkldnn/mkldnn_pooling.cc
@@ -134,6 +134,14 @@ mkldnn::algorithm GetMKLDNNPoolAlgo(const PoolingParam &param) {
   }
 }
 
+static inline int GetPaddingSizeFull(int x, int padl, int padr, int k, int s) {
+  if ((x + padl + padr - k) % s != 0) {
+    return (padr + s - ((x + padl + padr - k) % s));
+  } else {
+    return padr;
+  }
+}
+
 mkldnn::pooling_forward::primitive_desc GetPoolingFwdPdesc(
     const PoolingParam &param, const bool is_train, const memory::desc &data_md,
     const memory::desc &out_md) {
@@ -154,11 +162,17 @@ mkldnn::pooling_forward::primitive_desc GetPoolingFwdPdesc(
   int pad_l_ = param.pad[1], pad_r_ = param.pad[1];
   int stride_h_ = param.stride[0], stride_w_ = param.stride[1];
 
+  if (param.pooling_convention == pool_enum::kFull) {
+    pad_b_ = GetPaddingSizeFull(data_md.data.dims[2], pad_t_, pad_b_, kernel_h_, stride_h_);
+    pad_r_ = GetPaddingSizeFull(data_md.data.dims[3], pad_l_, pad_r_, kernel_w_, stride_w_);
+  }
+
   const mkldnn::engine engine = CpuEngine::Get()->get_engine();
   if (param.global_pool) {
     pad_t_ = pad_b_ = pad_l_ = pad_r_ = 0;
     stride_h_ = stride_w_ = 1;
   }
+
   if (pad_t_ != 0 || pad_l_ != 0) {
     CHECK(param.pool_type == pool_enum::kAvgPooling ||
           param.pool_type == pool_enum::kMaxPooling)
@@ -167,7 +181,6 @@ mkldnn::pooling_forward::primitive_desc GetPoolingFwdPdesc(
     CHECK_LT(pad_t_, kernel_h_);
   }
 
-
   const mkldnn::algorithm alg = GetMKLDNNPoolAlgo(param);
   mkldnn::prop_kind kind = mkldnn::prop_kind::forward_scoring;
   if (is_train && alg != algorithm::pooling_avg) {
@@ -227,17 +240,22 @@ MKLDNNPoolingFwd &GetPoolingFwd(const PoolingParam &param,
     int pad_l_ = param.pad[1], pad_r_ = param.pad[1];
     int stride_h_ = param.stride[0], stride_w_ = param.stride[1];
 
+    if (param.pooling_convention == pool_enum::kFull) {
+      pad_b_ = GetPaddingSizeFull(data_md.data.dims[2], pad_t_, pad_b_, kernel_h_, stride_h_);
+      pad_r_ = GetPaddingSizeFull(data_md.data.dims[3], pad_l_, pad_r_, kernel_w_, stride_w_);
+    }
+
     if (param.global_pool) {
-        pad_t_ = pad_b_ = pad_l_ = pad_r_ = 0;
-        stride_h_ = stride_w_ = 1;
+      pad_t_ = pad_b_ = pad_l_ = pad_r_ = 0;
+      stride_h_ = stride_w_ = 1;
     }
 
     if (pad_t_ != 0 || pad_l_ != 0) {
-        CHECK(param.pool_type == pool_enum::kAvgPooling ||
-              param.pool_type == pool_enum::kMaxPooling)
-              << "Padding implemented only for average and max pooling.";
-        CHECK_LT(pad_l_, kernel_w_);
-        CHECK_LT(pad_t_, kernel_h_);
+      CHECK(param.pool_type == pool_enum::kAvgPooling ||
+            param.pool_type == pool_enum::kMaxPooling)
+            << "Padding implemented only for average and max pooling.";
+      CHECK_LT(pad_l_, kernel_w_);
+      CHECK_LT(pad_t_, kernel_h_);
     }
 
     const mkldnn::algorithm alg = GetMKLDNNPoolAlgo(param);
@@ -353,6 +371,12 @@ MKLDNNPoolingBwd &GetPoolingBwd(const PoolingParam &param,
     int pad_t_ = param.pad[0], pad_b_ = param.pad[0];
     int pad_l_ = param.pad[1], pad_r_ = param.pad[1];
     int stride_h_ = param.stride[0], stride_w_ = param.stride[1];
+
+    if (param.pooling_convention == pool_enum::kFull) {
+      pad_b_ = GetPaddingSizeFull(data_md.data.dims[2], pad_t_, pad_b_, kernel_h_, stride_h_);
+      pad_r_ = GetPaddingSizeFull(data_md.data.dims[3], pad_l_, pad_r_, kernel_w_, stride_w_);
+    }
+
     if (param.global_pool) {
       pad_t_ = pad_b_ = pad_l_ = pad_r_ = 0;
       stride_h_ = stride_w_ = 1;
diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
index e3299681814..8054937a84c 100644
--- a/tests/python/gpu/test_operator_gpu.py
+++ b/tests/python/gpu/test_operator_gpu.py
@@ -996,6 +996,35 @@ def test_3d_pooling(pool_type, p_value=2, count_include_pad=True):
     test_3d_pooling('lp', p_value=3)
 
 
+@with_seed()
+def test_pooling_full_2d():
+    def test_pooling_full_2d_type(pool_type):
+        data = (2, 2, 10, 10)
+        kernel = (4, 5)
+        pad = (1, 2)
+        stride = (3, 4)
+
+        convention = 'full'
+        ctx_list = []
+        sym_list = []
+
+        # o_h = ceil((10 + 1 + 1 - 4) / 3) + 1 = 4
+        # o_w = ceil((10 + 2 + 2 - 5) / 4) + 1 = 4
+        ctx_list.append({'ctx': mx.cpu(0), 'pool_data': data, 'type_dict': {'pool_data': np.float32}})
+        sym_list.append(mx.sym.Pooling(kernel=kernel, pad=pad, stride=stride, pool_type=pool_type,
+                                       pooling_convention=convention, global_pool=False, name='pool'))
+
+        ctx_list.append({'ctx': mx.gpu(0), 'pool_data': data, 'type_dict': {'pool_data': np.float32}})
+        sym_list.append(mx.sym.Pooling(kernel=kernel, pad=pad, stride=stride, pool_type=pool_type,
+                                       pooling_convention=convention, global_pool=False, name='pool'))
+
+        check_consistency(sym_list, ctx_list)
+
+    test_pooling_full_2d_type('max')
+    test_pooling_full_2d_type('avg')
+    test_pooling_full_2d_type('sum')
+
+
 @with_seed()
 def test_global_pooling():
     def test_1d_pooling(pool_type, p_value=2):


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services