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