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)); }