Merge pull request #3439 from flx42/cudnn_v4

Add ifdef in CuDNNConvolutionLayer for cuDNN v4
This commit is contained in:
Evan Shelhamer 2015-12-12 01:00:00 -03:00
Родитель 541386054c eb2b848df1
Коммит 79539180ed
2 изменённых файлов: 11 добавлений и 0 удалений

Просмотреть файл

@ -7,6 +7,9 @@
#include "caffe/common.hpp"
#include "caffe/proto/caffe.pb.h"
#define CUDNN_VERSION_MIN(major, minor, patch) \
(CUDNN_VERSION >= (major * 1000 + minor * 100 + patch))
#define CUDNN_CHECK(condition) \
do { \
cudnnStatus_t status = condition; \

Просмотреть файл

@ -30,11 +30,19 @@ void CuDNNConvolutionLayer<Dtype>::Forward_gpu(
// Bias.
if (this->bias_term_) {
const Dtype* bias_data = this->blobs_[1]->gpu_data();
#if CUDNN_VERSION_MIN(4, 0, 0)
CUDNN_CHECK(cudnnAddTensor(handle_[g],
cudnn::dataType<Dtype>::one,
bias_desc_, bias_data + bias_offset_ * g,
cudnn::dataType<Dtype>::one,
top_descs_[i], top_data + top_offset_ * g));
#else
CUDNN_CHECK(cudnnAddTensor(handle_[g], CUDNN_ADD_SAME_C,
cudnn::dataType<Dtype>::one,
bias_desc_, bias_data + bias_offset_ * g,
cudnn::dataType<Dtype>::one,
top_descs_[i], top_data + top_offset_ * g));
#endif
}
}