By default do not choose CUDNN engine in CPU mode even if it is compi…

…led wiht USE_CUDNN Fixes BVLC#3317
dashesy · Jul 3, 2019 · 10666ec · 10666ec
1 parent 04ab089
commit 10666ec
Showing 1 changed file with 15 additions and 10 deletions.
diff --git a/src/caffe/layer_factory.cpp b/src/caffe/layer_factory.cpp
@@ -53,7 +53,7 @@ shared_ptr<Layer<Dtype> > GetConvolutionLayer(
   if (engine == ConvolutionParameter_Engine_DEFAULT) {
     engine = ConvolutionParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    if (!use_dilation) {
+    if (!use_dilation && Caffe::mode() == Caffe::GPU) {
       engine = ConvolutionParameter_Engine_CUDNN;
     }
 #endif
@@ -92,7 +92,7 @@ shared_ptr<Layer<Dtype> > GetDeconvolutionLayer(const LayerParameter& param) {
   if (engine == ConvolutionParameter_Engine_DEFAULT) {
     engine = ConvolutionParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    if (!use_dilation) {
+    if (!use_dilation && Caffe::mode() == Caffe::GPU) {
       engine = ConvolutionParameter_Engine_CUDNN;
     }
 #endif
@@ -122,7 +122,8 @@ shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
   if (engine == PoolingParameter_Engine_DEFAULT) {
     engine = PoolingParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    engine = PoolingParameter_Engine_CUDNN;
+    if (Caffe::mode() == Caffe::GPU)
+        engine = PoolingParameter_Engine_CUDNN;
 #endif
   }
   if (engine == PoolingParameter_Engine_CAFFE) {
@@ -159,10 +160,10 @@ shared_ptr<Layer<Dtype> > GetLRNLayer(const LayerParameter& param) {
   LRNParameter_Engine engine = param.lrn_param().engine();
 
   if (engine == LRNParameter_Engine_DEFAULT) {
+      engine = LRNParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    engine = LRNParameter_Engine_CUDNN;
-#else
-    engine = LRNParameter_Engine_CAFFE;
+      if (Caffe::mode() == Caffe::GPU)
+          engine = LRNParameter_Engine_CUDNN;
 #endif
   }
 
@@ -198,7 +199,8 @@ shared_ptr<Layer<Dtype> > GetReLULayer(const LayerParameter& param) {
   if (engine == ReLUParameter_Engine_DEFAULT) {
     engine = ReLUParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    engine = ReLUParameter_Engine_CUDNN;
+    if (Caffe::mode() == Caffe::GPU)
+        engine = ReLUParameter_Engine_CUDNN;
 #endif
   }
   if (engine == ReLUParameter_Engine_CAFFE) {
@@ -222,7 +224,8 @@ shared_ptr<Layer<Dtype> > GetSigmoidLayer(const LayerParameter& param) {
   if (engine == SigmoidParameter_Engine_DEFAULT) {
     engine = SigmoidParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    engine = SigmoidParameter_Engine_CUDNN;
+    if (Caffe::mode() == Caffe::GPU)
+        engine = SigmoidParameter_Engine_CUDNN;
 #endif
   }
   if (engine == SigmoidParameter_Engine_CAFFE) {
@@ -246,7 +249,8 @@ shared_ptr<Layer<Dtype> > GetSoftmaxLayer(const LayerParameter& param) {
   if (engine == SoftmaxParameter_Engine_DEFAULT) {
     engine = SoftmaxParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    engine = SoftmaxParameter_Engine_CUDNN;
+    if (Caffe::mode() == Caffe::GPU)
+        engine = SoftmaxParameter_Engine_CUDNN;
 #endif
   }
   if (engine == SoftmaxParameter_Engine_CAFFE) {
@@ -270,7 +274,8 @@ shared_ptr<Layer<Dtype> > GetTanHLayer(const LayerParameter& param) {
   if (engine == TanHParameter_Engine_DEFAULT) {
     engine = TanHParameter_Engine_CAFFE;
 #ifdef USE_CUDNN
-    engine = TanHParameter_Engine_CUDNN;
+    if (Caffe::mode() == Caffe::GPU)
+        engine = TanHParameter_Engine_CUDNN;
 #endif
   }
   if (engine == TanHParameter_Engine_CAFFE) {