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_r239265544
##########
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:
Yes, if it does not violate any nvidia license rules, it would be great to
add the link here as comment.
----------------------------------------------------------------
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
