| 16 | */ |
| 17 | template <typename Dtype> |
| 18 | void CuDNNConvolutionLayer<Dtype>::LayerSetUp( |
| 19 | const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) { |
| 20 | ConvolutionLayer<Dtype>::LayerSetUp(bottom, top); |
| 21 | // Initialize CUDA streams and cuDNN. |
| 22 | stream_ = new cudaStream_t[this->group_ * CUDNN_STREAMS_PER_GROUP]; |
| 23 | handle_ = new cudnnHandle_t[this->group_ * CUDNN_STREAMS_PER_GROUP]; |
| 24 | |
| 25 | // Initialize algorithm arrays |
| 26 | fwd_algo_ = new cudnnConvolutionFwdAlgo_t[bottom.size()]; |
| 27 | bwd_filter_algo_= new cudnnConvolutionBwdFilterAlgo_t[bottom.size()]; |
| 28 | bwd_data_algo_ = new cudnnConvolutionBwdDataAlgo_t[bottom.size()]; |
| 29 | |
| 30 | // initialize size arrays |
| 31 | workspace_fwd_sizes_ = new size_t[bottom.size()]; |
| 32 | workspace_bwd_filter_sizes_ = new size_t[bottom.size()]; |
| 33 | workspace_bwd_data_sizes_ = new size_t[bottom.size()]; |
| 34 | |
| 35 | // workspace data |
| 36 | workspaceSizeInBytes = 0; |
| 37 | workspaceData = NULL; |
| 38 | workspace = new void*[this->group_ * CUDNN_STREAMS_PER_GROUP]; |
| 39 | |
| 40 | for (size_t i = 0; i < bottom.size(); ++i) { |
| 41 | // initialize all to default algorithms |
| 42 | fwd_algo_[i] = (cudnnConvolutionFwdAlgo_t)0; |
| 43 | bwd_filter_algo_[i] = (cudnnConvolutionBwdFilterAlgo_t)0; |
| 44 | bwd_data_algo_[i] = (cudnnConvolutionBwdDataAlgo_t)0; |
| 45 | // default algorithms don't require workspace |
| 46 | workspace_fwd_sizes_[i] = 0; |
| 47 | workspace_bwd_data_sizes_[i] = 0; |
| 48 | workspace_bwd_filter_sizes_[i] = 0; |
| 49 | } |
| 50 | |
| 51 | for (int g = 0; g < this->group_ * CUDNN_STREAMS_PER_GROUP; g++) { |
| 52 | CUDA_CHECK(cudaStreamCreate(&stream_[g])); |
| 53 | CUDNN_CHECK(cudnnCreate(&handle_[g])); |
| 54 | CUDNN_CHECK(cudnnSetStream(handle_[g], stream_[g])); |
| 55 | workspace[g] = NULL; |
| 56 | } |
| 57 | |
| 58 | // Set the indexing parameters. |
| 59 | bias_offset_ = (this->num_output_ / this->group_); |
| 60 | |
| 61 | // Create filter descriptor. |
| 62 | const int* kernel_shape_data = this->kernel_shape_.cpu_data(); |
| 63 | const int kernel_h = kernel_shape_data[0]; |
| 64 | const int kernel_w = kernel_shape_data[1]; |
| 65 | cudnn::createFilterDesc<Dtype>(&filter_desc_, |
| 66 | this->num_output_ / this->group_, this->channels_ / this->group_, |
| 67 | kernel_h, kernel_w); |
| 68 | |
| 69 | // Create tensor descriptor(s) for data and corresponding convolution(s). |
| 70 | for (int i = 0; i < bottom.size(); i++) { |
| 71 | cudnnTensorDescriptor_t bottom_desc; |
| 72 | cudnn::createTensor4dDesc<Dtype>(&bottom_desc); |
| 73 | bottom_descs_.push_back(bottom_desc); |
| 74 | cudnnTensorDescriptor_t top_desc; |
| 75 | cudnn::createTensor4dDesc<Dtype>(&top_desc); |