inline void createPoolingDesc(cudnnPoolingDescriptor_t* pool,
cudnnPoolingMode_t mode, int h, int w, int pad_h, int pad_w,
int stride_h, int stride_w) {
CUDNN_CHECK(cudnnCreatePoolingDescriptor(pool));
CUDNN_CHECK(cudnnSetPooling2dDescriptor(*pool, mode, h, w, pad_h, pad_w,
stride_h, stride_w));
}
THFloatTensor *cudnn_SpatialMaxPooling_updateOutput(struct module *module, THFloatTensor *input)
{
int kW = module->SpatialMaxPooling.kW;
int kH = module->SpatialMaxPooling.kH;
int dW = module->SpatialMaxPooling.dW;
int dH = module->SpatialMaxPooling.dH;
int padW = module->SpatialMaxPooling.padW;
int padH = module->SpatialMaxPooling.padH;
THFloatTensor *output = module->output;
cudnnTensorDescriptor_t dinput, doutput;
cudnnPoolingDescriptor_t dpool;
float one = 1, zero = 0;
int sizes[4];
errcheck(THcudnn_TensorDescriptor(&dinput, input));
errcheck(cudnnCreatePoolingDescriptor(&dpool));
errcheck(cudnnSetPooling2dDescriptor(dpool, CUDNN_POOLING_MAX, kH, kW, padH, padW, dH, dW));
errcheck(cudnnGetPoolingNdForwardOutputDim(dpool, dinput, 4, sizes));
THCudaTensor_resize4d(output, sizes[0], sizes[1], sizes[2], sizes[3]);
errcheck(THcudnn_TensorDescriptor(&doutput, output));
errcheck(cudnnPoolingForward(THcudnn_getHandle(), dpool, &one, dinput, THFloatTensor_data(input), &zero,
doutput, THFloatTensor_data(output)));
cudnnDestroyTensorDescriptor(dinput);
cudnnDestroyTensorDescriptor(doutput);
cudnnDestroyPoolingDescriptor(dpool);
return output;
}
inline void createPoolingDesc(cudnnPoolingDescriptor_t* pool_desc,
PoolingParameter_PoolMethod poolmethod, cudnnPoolingMode_t* mode,
int h, int w, int pad_h, int pad_w, int stride_h, int stride_w) {
switch (poolmethod) {
case PoolingParameter_PoolMethod_MAX:
*mode = CUDNN_POOLING_MAX;
break;
case PoolingParameter_PoolMethod_AVE:
*mode = CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
break;
default:
LOG(FATAL) << "Unknown pooling method.";
}
CUDNN_CHECK(cudnnCreatePoolingDescriptor(pool_desc));
CUDNN_CHECK(cudnnSetPooling2dDescriptor(*pool_desc, *mode, h, w,
pad_h, pad_w, stride_h, stride_w));
}
inline void InitCuDNN() {
init_cudnn_ = false;
dtype_ = CUDNN_DATA_FLOAT;
switch(mode) {
case kMaxPooling: mode_ = CUDNN_POOLING_MAX; break;
// case kAvgPooling: mode_ = CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING; break;
default: utils::Error("This should not happen -,-"); break;
}
CUDA_CHECK(cudnnCreate(&handle_));
CUDA_CHECK(cudnnCreateTensorDescriptor(&in_desc_));
CUDA_CHECK(cudnnCreateTensorDescriptor(&out_desc_));
CUDA_CHECK(cudnnCreatePoolingDescriptor(&pooling_desc_));
CUDA_CHECK(cudnnSetPooling2dDescriptor(pooling_desc_, mode_,
Parent::param_.kernel_height,
Parent::param_.kernel_width,
0, 0,
Parent::param_.stride,
Parent::param_.stride));
}
