| /************************************************************ |
| * |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, |
| * software distributed under the License is distributed on an |
| * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| * KIND, either express or implied. See the License for the |
| * specific language governing permissions and limitations |
| * under the License. |
| * |
| *************************************************************/ |
| |
| #include "singa/neuralnet/neuron_layer.h" |
| |
| namespace singa { |
| |
| CudnnConvLayer::~CudnnConvLayer() { |
| if (has_init_cudnn_) { |
| CHECK_CUDNN(cudnnDestroyTensorDescriptor(bias_desc_)); |
| CHECK_CUDNN(cudnnDestroyFilterDescriptor(filter_desc_)); |
| CHECK_CUDNN(cudnnDestroyConvolutionDescriptor(conv_desc_)); |
| } |
| } |
| |
| void CudnnConvLayer::InitCudnn() { |
| CudnnBase::InitCudnn(); |
| // convert MB to bytes |
| workspace_byte_limit_ |
| = layer_conf_.convolution_conf().workspace_byte_limit() << 20; |
| |
| CHECK_CUDNN(cudnnCreateTensorDescriptor(&bias_desc_)); |
| CHECK_CUDNN(cudnnCreateFilterDescriptor(&filter_desc_)); |
| CHECK_CUDNN(cudnnCreateConvolutionDescriptor(&conv_desc_)); |
| |
| CHECK_CUDNN(cudnnSetConvolution2dDescriptor(conv_desc_, |
| pad_y_, |
| pad_x_, |
| stride_y_, |
| stride_x_, |
| 1, |
| 1, |
| CUDNN_CROSS_CORRELATION)); |
| CHECK_CUDNN(cudnnSetFilter4dDescriptor(filter_desc_, |
| CUDNN_DATA_FLOAT, |
| num_filters_, |
| channels_, |
| kernel_y_, |
| kernel_x_)); |
| if (bias_) { |
| CHECK_CUDNN(cudnnSetTensor4dDescriptor(bias_desc_, |
| CUDNN_TENSOR_NCHW, |
| CUDNN_DATA_FLOAT, |
| 1, |
| num_filters_, |
| 1, |
| 1)); |
| } |
| CHECK_CUDNN(cudnnSetTensor4dDescriptor(src_desc_, |
| CUDNN_TENSOR_NCHW, |
| CUDNN_DATA_FLOAT, |
| batchsize_, |
| channels_, |
| height_, |
| width_)); |
| CHECK_CUDNN(cudnnSetTensor4dDescriptor(my_desc_, |
| CUDNN_TENSOR_NCHW, |
| CUDNN_DATA_FLOAT, |
| batchsize_, |
| num_filters_, |
| conv_height_, |
| conv_width_)); |
| |
| CHECK_CUDNN(cudnnGetConvolutionForwardAlgorithm(handle_, |
| src_desc_, |
| filter_desc_, |
| conv_desc_, |
| my_desc_, |
| CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, |
| workspace_byte_limit_, |
| &fp_alg_)); |
| |
| CHECK_CUDNN(cudnnGetConvolutionBackwardFilterAlgorithm(handle_, |
| src_desc_, |
| my_desc_, |
| conv_desc_, |
| filter_desc_, |
| CUDNN_CONVOLUTION_BWD_FILTER_PREFER_FASTEST, |
| workspace_byte_limit_, |
| &bp_filter_alg_)); |
| CHECK_CUDNN(cudnnGetConvolutionBackwardDataAlgorithm(handle_, |
| filter_desc_, |
| my_desc_, |
| conv_desc_, |
| src_desc_, |
| CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST, |
| workspace_byte_limit_, |
| &bp_data_alg_)); |
| |
| size_t fp_byte, bp_data_byte, bp_filter_byte; |
| CHECK_CUDNN(cudnnGetConvolutionForwardWorkspaceSize(handle_, |
| src_desc_, |
| filter_desc_, |
| conv_desc_, |
| my_desc_, |
| fp_alg_, |
| &fp_byte)); |
| CHECK_CUDNN(cudnnGetConvolutionBackwardDataWorkspaceSize(handle_, |
| filter_desc_, |
| my_desc_, |
| conv_desc_, |
| src_desc_, |
| bp_data_alg_, |
| &bp_data_byte)); |
| CHECK_CUDNN(cudnnGetConvolutionBackwardFilterWorkspaceSize(handle_, |
| src_desc_, |
| my_desc_, |
| conv_desc_, |
| filter_desc_, |
| bp_filter_alg_, |
| &bp_filter_byte)); |
| workspace_count_ = std::max(std::max(fp_byte, bp_data_byte), bp_filter_byte) |
| / sizeof(float) + 1; |
| } |
| |
| void CudnnConvLayer::ComputeFeature(int flag, const vector<Layer*>& srclayers) { |
| if (!has_init_cudnn_) |
| InitCudnn(); |
| float alpha = 1.f, beta = 0.f; |
| Blob<float> workspace(vector<int>{static_cast<int>(workspace_count_)}); |
| CHECK_CUDNN(cudnnConvolutionForward(handle_, |
| &alpha, |
| src_desc_, |
| srclayers[0]->data(this).gpu_data(), |
| filter_desc_, |
| weight_->data().gpu_data(), |
| conv_desc_, |
| fp_alg_, |
| workspace.mutable_gpu_data(), |
| workspace_count_ * sizeof(float), |
| &beta, |
| my_desc_, |
| data_.mutable_gpu_data())); |
| if (bias_) { |
| beta = 1.f; |
| |
| #if CUDNN_MAJOR == 4 |
| CHECK_CUDNN(cudnnAddTensor(handle_, |
| &alpha, |
| bias_desc_, |
| bias_->data().gpu_data(), |
| &beta, |
| my_desc_, |
| data_.mutable_gpu_data())); |
| #else |
| CHECK_CUDNN(cudnnAddTensor(handle_, |
| CUDNN_ADD_SAME_C, |
| &alpha, |
| bias_desc_, |
| bias_->data().gpu_data(), |
| &beta, |
| my_desc_, |
| data_.mutable_gpu_data())); |
| #endif |
| } |
| } |
| |
| void |
| CudnnConvLayer::ComputeGradient(int flag, const vector<Layer*>& srclayers) { |
| float alpha = 1.f, beta = 0.f; |
| Blob<float> workspace(vector<int>{static_cast<int>(workspace_count_)}); |
| // LOG(ERROR) << "backward bias"; |
| if (bias_) { |
| CHECK_CUDNN(cudnnConvolutionBackwardBias(handle_, |
| &alpha, |
| my_desc_, |
| grad_.gpu_data(), |
| &beta, |
| bias_desc_, |
| bias_->mutable_grad()->mutable_gpu_data())); |
| } |
| // LOG(ERROR) << "backward w"; |
| CHECK_CUDNN(cudnnConvolutionBackwardFilter_v3(handle_, |
| &alpha, |
| src_desc_, |
| srclayers[0]->data(this).gpu_data(), |
| my_desc_, |
| grad_.gpu_data(), |
| conv_desc_, |
| bp_filter_alg_, |
| workspace.mutable_gpu_data(), |
| workspace_count_ * sizeof(float), |
| &beta, |
| filter_desc_, |
| weight_->mutable_grad()->mutable_gpu_data())); |
| // LOG(ERROR) << "backward src"; |
| if (srclayers[0]->mutable_grad(this) != nullptr) { |
| CHECK_CUDNN(cudnnConvolutionBackwardData_v3(handle_, |
| &alpha, |
| filter_desc_, |
| weight_->data().gpu_data(), |
| my_desc_, |
| grad_.gpu_data(), |
| conv_desc_, |
| bp_data_alg_, |
| workspace.mutable_gpu_data(), |
| workspace_count_ * sizeof(float), |
| &beta, |
| src_desc_, |
| srclayers[0]->mutable_grad(this)->mutable_gpu_data())); |
| } |
| } |
| } // namespace singa |
