void cudnn_util::set_pooling_descriptor(
cudnnPoolingDescriptor_t pooling_desc,
cudnnPoolingMode_t pooling_mode,
const std::vector<unsigned int>& subsampling_sizes)
{
std::vector<int> padding(subsampling_sizes.size(), 0);
std::vector<int> dimensions(subsampling_sizes.rbegin(), subsampling_sizes.rend());
cudnn_safe_call(cudnnSetPoolingNdDescriptor(
pooling_desc,
pooling_mode,
static_cast<int>(subsampling_sizes.size()),
&dimensions[0],
&padding[0],
&dimensions[0]));
}
void PoolBC01CuDNN<T>::fprop(const T *imgs, int *imgs_shape, T *poolout) {
bool new_shape = false;
int n_imgs_dims = n_img_dims + 2;
for (int i = 0; i < n_imgs_dims; ++i) {
if (this->imgs_shape[i] != imgs_shape[i]) {
new_shape = true;
break;
}
}
if (new_shape) {
for (int i = 0; i < n_imgs_dims; ++i) {
this->imgs_shape[i] = imgs_shape[i];
}
int imgs_strides[n_imgs_dims];
array_strides(n_imgs_dims, imgs_shape, imgs_strides);
CUDNN_CHECK(cudnnSetTensorNdDescriptor(
imgs_desc, CUDNN_DATA_FLOAT, n_imgs_dims, imgs_shape, imgs_strides
));
CUDNN_CHECK(cudnnSetPoolingNdDescriptor(
pool_desc, pool_mode, n_img_dims, win_shape, padding, strides
));
int poolout_shape[n_imgs_dims];
poolout_shape[0] = imgs_shape[0];
poolout_shape[1] = imgs_shape[1];
for (int i = 0; i < n_img_dims; ++i) {
poolout_shape[i+2] = (imgs_shape[i+2] + 2*padding[i] - win_shape[i])
/ strides[i] + 1;
}
int poolout_strides[n_imgs_dims];
array_strides(n_imgs_dims, poolout_shape, poolout_strides);
CUDNN_CHECK(cudnnSetTensorNdDescriptor(
poolout_desc, CUDNN_DATA_FLOAT, n_imgs_dims, poolout_shape,
poolout_strides
));
}
CUDNN_CHECK(cudnnPoolingForward(
CUDNN::handle(), pool_desc, &CUDNN::one, imgs_desc, imgs, &CUDNN::zero,
poolout_desc, poolout
));
}
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));
}
