apeforest 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_r238951359
##########
File path: src/operator/nn/cudnn/cudnn_pooling-inl.h
##########
@@ -165,55 +170,78 @@ 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.
+ #if CUDNN_VERSION == 7104
+ is_supported = window_height <= 8 && window_width <= 8;
Review comment:
Could you please add a reference to this?
----------------------------------------------------------------
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:
us...@infra.apache.org
With regards,
Apache Git Services
