anirudh2290 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_r239658011
##########
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";
Review comment:
is this needed ?
----------------------------------------------------------------
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:
[email protected]
With regards,
Apache Git Services
