You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/12/07 00:40:16 UTC

[GitHub] yuxihu commented on a change in pull request #13362: Add NHWC layout support to Pooling (cuDNN only)

yuxihu commented on a change in pull request #13362: Add NHWC layout support to Pooling (cuDNN only)
URL: https://github.com/apache/incubator-mxnet/pull/13362#discussion_r239659470
 
 

 ##########
 File path: src/operator/nn/cudnn/cudnn_pooling-inl.h
 ##########
 @@ -165,55 +170,80 @@ class CuDNNPoolingOp {
     } else {
       LOG(FATAL) << "Only support 2D or 3D pooling";
     }
+    return true;
   }
 
  private:
-  inline void Init(mshadow::Stream<gpu> *s, const TBlob &in_data,
+  // Return boolean saying whether pooling configuration is supported
+  inline bool Init(mshadow::Stream<gpu> *s, const TBlob &in_data,
       const TBlob &out_data) {
     using namespace mshadow;
+    bool is_supported = true;
     #if CUDNN_MAJOR >= 5
     nan_prop_ = CUDNN_NOT_PROPAGATE_NAN;
     #endif
     if (param_.kernel.ndim() == 2) {
       // 2d conv
+      CHECK(param_.layout.value() == mshadow::kNCHW ||
+            param_.layout.value() == mshadow::kNHWC) << "Need 2D layout";
+      cudnnTensorFormat_t cudnn_layout =
+          (param_.layout.value() == mshadow::kNCHW) ? CUDNN_TENSOR_NCHW
+                                                    : CUDNN_TENSOR_NHWC;
       Tensor<gpu, 4, DType> data = in_data.get<gpu, 4, DType>(s);
       Tensor<gpu, 4, DType> out = out_data.get<gpu, 4, DType>(s);
-      mshadow::Shape<4> dshape = data.shape_;
+      // Perform shape calculations in a standard (NCHW) layout space
+      mshadow::Shape<4> dshape_nchw = (param_.layout.value() == mshadow::kNHWC) ?
+                                      ConvertLayout(data.shape_, mshadow::kNHWC, mshadow::kNCHW) :
+                                      data.shape_;
+      mshadow::Shape<4> oshape_nchw = (param_.layout.value() == mshadow::kNHWC) ?
+                                      ConvertLayout(out.shape_, mshadow::kNHWC, mshadow::kNCHW) :
+                                      out.shape_;
       CUDNN_CALL(cudnnSetTensor4dDescriptor(in_desc_,
-                                            CUDNN_TENSOR_NCHW,
+                                            cudnn_layout,
                                             dtype_,
-                                            data.shape_[0],
-                                            data.shape_[1],
-                                            data.shape_[2],
-                                            data.shape_[3]));
+                                            dshape_nchw[0],
+                                            dshape_nchw[1],
+                                            dshape_nchw[2],
+                                            dshape_nchw[3]));
       CUDNN_CALL(cudnnSetTensor4dDescriptor(out_desc_,
-                                            CUDNN_TENSOR_NCHW,
+                                            cudnn_layout,
                                             dtype_,
-                                            out.shape_[0],
-                                            out.shape_[1],
-                                            out.shape_[2],
-                                            out.shape_[3]));
+                                            oshape_nchw[0],
+                                            oshape_nchw[1],
+                                            oshape_nchw[2],
+                                            oshape_nchw[3]));
+      int window_height = param_.global_pool ? dshape_nchw[2] : param_.kernel[0];
+      int window_width = param_.global_pool ? dshape_nchw[3] : param_.kernel[1];
+      // CuDNN v7.1.4 backprop kernel doesn't support window sizes 9 and above.
+      // For reference see Fixed Issues section in
+      // https://docs.nvidia.com/deeplearning/sdk/cudnn-release-notes/rel_721.html#rel_721
+      #if CUDNN_VERSION == 7104
+      is_supported = window_height <= 8 && window_width <= 8;
+      #endif
       #if CUDNN_MAJOR >= 5
       CUDNN_CALL(cudnnSetPooling2dDescriptor(pooling_desc_,
                                              mode_,
                                              nan_prop_,
-                                             param_.global_pool ? dshape[2] : param_.kernel[0],
-                                             param_.global_pool ? dshape[3] : param_.kernel[1],
+                                             window_height,
+                                             window_width,
                                              param_.global_pool ? 0 : param_.pad[0],
                                              param_.global_pool ? 0 : param_.pad[1],
                                              param_.global_pool ? 1 : param_.stride[0],
-                                             param_.global_pool ? 1 :param_.stride[1]));
+                                             param_.global_pool ? 1 : param_.stride[1]));
       #else
       CUDNN_CALL(cudnnSetPooling2dDescriptor(pooling_desc_,
                                              mode_,
-                                             param_.global_pool ? dshape[2] : param_.kernel[0],
-                                             param_.global_pool ? dshape[3] : param_.kernel[1],
+                                             window_height,
+                                             window_width,
                                              param_.global_pool ? 0 : param_.pad[0],
-                                             param_.global_ppol ? 0 : param_.pad[1],
+                                             param_.global_pool ? 0 : param_.pad[1],
                                              param_.global_pool ? 1 : param_.stride[0],
                                              param_.global_pool ? 1 : param_.stride[1]));
       #endif
     } else {
+      CHECK(param_.layout.value() == mshadow::kNCDHW ||
+            param_.layout.value() == mshadow::kNDHWC) << "Need 3D layout";
+      CHECK(param_.layout.value() == mshadow::kNCDHW) << "Only the NCDHW layout is supported.";
 
 Review comment:
   If only NCDHW layout is supported, it seems me that there is no need to have the above "Need 3D layout" check.
   
   Besides that, I am also a little bit confused here. In conv_layers.py, it is mentioned that "Dimension ordering of data and weight. Only supports 'NCDHW' and 'NDHWC' (only with cuDNN) layouts for now." for all 3D pool classes. Why here we only support NCDHW? This is CuDNNPoolingOp class, right? Am I missing something?

----------------------------------------------------------------
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