diff options
| author | Felix Abecassis <fabecassis@nvidia.com> | 2016-05-16 14:03:38 -0700 |
|---|---|---|
| committer | Felix Abecassis <fabecassis@nvidia.com> | 2016-05-16 14:03:38 -0700 |
| commit | b43c8e43a95608a00033f8f8867d32a201e5eed4 (patch) | |
| tree | 2117571f885d297e14fa9c629878900bb1d588dd /include | |
| parent | bb0c1a5ba18b685dbd62efbac06b2609e20d1d88 (diff) | |
| download | caffeonacl-b43c8e43a95608a00033f8f8867d32a201e5eed4.tar.gz caffeonacl-b43c8e43a95608a00033f8f8867d32a201e5eed4.tar.bz2 caffeonacl-b43c8e43a95608a00033f8f8867d32a201e5eed4.zip | |
Add cuDNN v5 support, drop cuDNN v3 support
cuDNN v4 is still supported.
Diffstat (limited to 'include')
| -rw-r--r-- | include/caffe/layers/cudnn_relu_layer.hpp | 1 | ||||
| -rw-r--r-- | include/caffe/layers/cudnn_sigmoid_layer.hpp | 1 | ||||
| -rw-r--r-- | include/caffe/layers/cudnn_tanh_layer.hpp | 1 | ||||
| -rw-r--r-- | include/caffe/util/cudnn.hpp | 24 |
4 files changed, 24 insertions, 3 deletions
diff --git a/include/caffe/layers/cudnn_relu_layer.hpp b/include/caffe/layers/cudnn_relu_layer.hpp index e01f568a..a1cb29e7 100644 --- a/include/caffe/layers/cudnn_relu_layer.hpp +++ b/include/caffe/layers/cudnn_relu_layer.hpp @@ -37,6 +37,7 @@ class CuDNNReLULayer : public ReLULayer<Dtype> { cudnnHandle_t handle_; cudnnTensorDescriptor_t bottom_desc_; cudnnTensorDescriptor_t top_desc_; + cudnnActivationDescriptor_t activ_desc_; }; #endif diff --git a/include/caffe/layers/cudnn_sigmoid_layer.hpp b/include/caffe/layers/cudnn_sigmoid_layer.hpp index 9c597958..7b3486f8 100644 --- a/include/caffe/layers/cudnn_sigmoid_layer.hpp +++ b/include/caffe/layers/cudnn_sigmoid_layer.hpp @@ -37,6 +37,7 @@ class CuDNNSigmoidLayer : public SigmoidLayer<Dtype> { cudnnHandle_t handle_; cudnnTensorDescriptor_t bottom_desc_; cudnnTensorDescriptor_t top_desc_; + cudnnActivationDescriptor_t activ_desc_; }; #endif diff --git a/include/caffe/layers/cudnn_tanh_layer.hpp b/include/caffe/layers/cudnn_tanh_layer.hpp index c0f0053f..59e758d7 100644 --- a/include/caffe/layers/cudnn_tanh_layer.hpp +++ b/include/caffe/layers/cudnn_tanh_layer.hpp @@ -37,6 +37,7 @@ class CuDNNTanHLayer : public TanHLayer<Dtype> { cudnnHandle_t handle_; cudnnTensorDescriptor_t bottom_desc_; cudnnTensorDescriptor_t top_desc_; + cudnnActivationDescriptor_t activ_desc_; }; #endif diff --git a/include/caffe/util/cudnn.hpp b/include/caffe/util/cudnn.hpp index 8a7e17c6..a7d8dbba 100644 --- a/include/caffe/util/cudnn.hpp +++ b/include/caffe/util/cudnn.hpp @@ -91,8 +91,13 @@ template <typename Dtype> inline void createFilterDesc(cudnnFilterDescriptor_t* desc, int n, int c, int h, int w) { CUDNN_CHECK(cudnnCreateFilterDescriptor(desc)); +#if CUDNN_VERSION_MIN(5, 0, 0) CUDNN_CHECK(cudnnSetFilter4dDescriptor(*desc, dataType<Dtype>::type, - n, c, h, w)); + CUDNN_TENSOR_NCHW, n, c, h, w)); +#else + CUDNN_CHECK(cudnnSetFilter4dDescriptor_v4(*desc, dataType<Dtype>::type, + CUDNN_TENSOR_NCHW, n, c, h, w)); +#endif } template <typename Dtype> @@ -123,8 +128,21 @@ inline void createPoolingDesc(cudnnPoolingDescriptor_t* pool_desc, 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)); +#if CUDNN_VERSION_MIN(5, 0, 0) + CUDNN_CHECK(cudnnSetPooling2dDescriptor(*pool_desc, *mode, + CUDNN_PROPAGATE_NAN, h, w, pad_h, pad_w, stride_h, stride_w)); +#else + CUDNN_CHECK(cudnnSetPooling2dDescriptor_v4(*pool_desc, *mode, + CUDNN_PROPAGATE_NAN, h, w, pad_h, pad_w, stride_h, stride_w)); +#endif +} + +template <typename Dtype> +inline void createActivationDescriptor(cudnnActivationDescriptor_t* activ_desc, + cudnnActivationMode_t mode) { + CUDNN_CHECK(cudnnCreateActivationDescriptor(activ_desc)); + CUDNN_CHECK(cudnnSetActivationDescriptor(*activ_desc, mode, + CUDNN_PROPAGATE_NAN, Dtype(0))); } } // namespace cudnn |