Merge pull request #112 from drnikolaev/nvcaffe-0.14-memory-opt

Memory consuming optimization for BN cuDNN algos
NVIDIA · Mar 15, 2016 · a8d1b9d · a8d1b9d
2 parents 192d349 + e8dab0e
commit a8d1b9d
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 3 deletions.
diff --git a/include/caffe/layer.hpp b/include/caffe/layer.hpp
@@ -38,7 +38,8 @@ class Layer {
    * layer.
    */
   explicit Layer(const LayerParameter& param)
-    : layer_param_(param), is_shared_(false) {
+    : layer_param_(param), is_shared_(false),
+      forward_passed_(false), backward_passed_(false) {
       // Set phase and copy blobs (if there are any).
       phase_ = param.phase();
       if (layer_param_.blobs_size() > 0) {
@@ -316,6 +317,21 @@ class Layer {
     param_propagate_down_[param_id] = value;
   }
 
+  bool IsForwardPassed() const {
+    return forward_passed_;
+  }
+
+  void ForwardPassed(bool passed) {
+    forward_passed_ = passed;
+  }
+
+  bool IsBackwardPassed() const {
+    return backward_passed_;
+  }
+
+  void BackwardPassed(bool passed) {
+    backward_passed_ = passed;
+  }
 
  protected:
   /** The protobuf that stores the layer parameters */
@@ -431,6 +447,9 @@ class Layer {
   /** Whether this layer is actually shared by other nets*/
   bool is_shared_;
 
+  bool forward_passed_;
+  bool backward_passed_;
+
   /** The mutex for sequential forward if this layer is shared */
   shared_ptr<boost::mutex> forward_mutex_;
 

diff --git a/src/caffe/layers/cudnn_conv_layer.cpp b/src/caffe/layers/cudnn_conv_layer.cpp
@@ -118,8 +118,12 @@ void CuDNNConvolutionLayer<Dtype>::Reshape(
     cudnn::setConvolutionDesc<Dtype>(&conv_descs_[i], bottom_descs_[i],
         filter_desc_, pad_h, pad_w, stride_h, stride_w);
 
-     // choose forward and backward algorithms + workspace(s)
-      CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm(Caffe::cudnn_handle(),
+    if (!this->IsForwardPassed() || !this->IsBackwardPassed()) {
+      continue;
+    }
+
+    // choose forward and backward algorithms + workspace(s)
+    CUDNN_CHECK(cudnnGetConvolutionForwardAlgorithm(Caffe::cudnn_handle(),
       bottom_descs_[i],
       filter_desc_,
       conv_descs_[i],

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
@@ -580,6 +580,9 @@ Dtype Net<Dtype>::ForwardFromTo(int start, int end) {
     loss += layer_loss;
     if (debug_info_) { ForwardDebugInfo(i); }
   }
+  for (int i = start; i <= end; ++i) {
+    layers_[i]->ForwardPassed(true);
+  }
   return loss;
 }
 
@@ -645,6 +648,9 @@ void Net<Dtype>::BackwardFromTo(int start, int end) {
       if (debug_info_) { BackwardDebugInfo(i); }
     }
   }
+  for (int i = start; i >= end; --i) {
+    layers_[i]->BackwardPassed(true);
+  }
 }
 
 template <typename Dtype>