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,
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));
}
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 createPoolingDesc(cudnnPoolingDescriptor_t* conv,
PoolingParameter_PoolMethod poolmethod, cudnnPoolingMode_t* mode,
int h, int 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;
break;
default:
LOG(FATAL) << "Unknown pooling method.";
}
CUDNN_CHECK(cudnnCreatePoolingDescriptor(conv));
CUDNN_CHECK(cudnnSetPoolingDescriptor(*conv, *mode, h, 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));
}
PoolBC01CuDNN<T>::PoolBC01CuDNN(int n_img_dims, int *win_shape, int *padding,
int *strides, PoolMode pool_mode) : n_img_dims(n_img_dims) {
if (n_img_dims > MAX_IMG_DIMS + 2) {
throw std::runtime_error("More than 3 image dimensions.");
}
for (int i = 0; i < n_img_dims; ++i) {
this->win_shape[i] = win_shape[i];
this->padding[i] = padding[i];
this->strides[i] = strides[i];
}
for (int i = 0; i < n_img_dims + 2; ++i) {
imgs_shape[i] = -1;
}
this->pool_mode = pool_mode == POOL_MAX ? CUDNN_POOLING_MAX :
CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING;
CUDNN_CHECK(cudnnCreateTensorDescriptor(&imgs_desc));
CUDNN_CHECK(cudnnCreateTensorDescriptor(&poolout_desc));
CUDNN_CHECK(cudnnCreatePoolingDescriptor(&pool_desc));
}
inline void createPoolingDesc(cudnnPoolingDescriptor_t* pool_desc,
PoolingParameter_PoolMethod poolmethod, cudnnPoolingMode_t* mode,
const int_tp num_spatial_dims,
const int_tp* shape,
const int_tp* pad, const int_tp* stride) {
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));
std::vector<int> shape_int(num_spatial_dims);
std::vector<int> pad_int(num_spatial_dims);
std::vector<int> stride_int(num_spatial_dims);
for (int_tp i = 0; i < num_spatial_dims; ++i) {
shape_int[i] = shape[i];
pad_int[i] = pad[i];
stride_int[i] = stride[i];
}
const int* shape_ptr = &shape_int[0];
const int* pad_ptr = &pad_int[0];
const int* stride_ptr = &stride_int[0];
CUDNN_CHECK(cudnnSetPoolingNdDescriptor(*pool_desc, *mode,
num_spatial_dims,
shape_ptr,
pad_ptr,
stride_ptr));
}
