Merge pull request BVLC#102 from drnikolaev/pooling-cudnn-fixes

cuDNN and pooling fixes
thatguymike · Jan 19, 2016 · b3a4a1a · b3a4a1a
2 parents 8eebbb4 + 569aa49
commit b3a4a1a
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/caffe/layer_factory.cpp b/src/caffe/layer_factory.cpp
@@ -78,7 +78,14 @@ shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
 #ifdef USE_CUDNN
   } else if (engine == PoolingParameter_Engine_CUDNN) {
     PoolingParameter p_param = param.pooling_param();
+
+// FIXME We assume that precision issue will be fixed in cuDNN v5.0
+#if CUDNN_VERSION >= 5000
     return shared_ptr<Layer<Dtype> >(new CuDNNPoolingLayer<Dtype>(param));
+#else
+    return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
+#endif
+
 #endif
   } else {
     LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";

diff --git a/src/caffe/layers/cudnn_batch_norm_layer.cu b/src/caffe/layers/cudnn_batch_norm_layer.cu
@@ -85,6 +85,10 @@ void CuDNNBatchNormLayer<Dtype>::Backward_gpu(
       mode_,
       cudnn::dataType<Dtype>::one,
       cudnn::dataType<Dtype>::zero,
+#if CUDNN_VERSION >= 4005
+      cudnn::dataType<Dtype>::one,
+      cudnn::dataType<Dtype>::zero,
+#endif
       bottom_desc_,
       bottom_data,
       bottom_desc_,