Add iter information to layer

DancingJane · Aug 23, 2019 · de32caa · de32caa
1 parent 96bb265
commit de32caa
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/include/caffe/layer.hpp b/include/caffe/layer.hpp
@@ -290,8 +290,13 @@ class Layer {
     }
     param_propagate_down_[param_id] = value;
   }
-
-
+  /**
+  */
+  virtual void set_iter(int iter,int max_iter) {
+    iter_ = iter;
+    max_iter_ = max_iter;
+  }
+
  protected:
   /** The protobuf that stores the layer parameters */
   LayerParameter layer_param_;
@@ -305,7 +310,9 @@ class Layer {
   /** The vector that indicates whether each top blob has a non-zero weight in
    *  the objective function. */
   vector<Dtype> loss_;
-
+  /**
+  */
+  int iter_ , max_iter_;
   /** @brief Using the CPU device, compute the layer output. */
   virtual void Forward_cpu(const vector<Blob<Dtype>*>& bottom,
       const vector<Blob<Dtype>*>& top) = 0;

diff --git a/include/caffe/net.hpp b/include/caffe/net.hpp
@@ -251,7 +251,10 @@ class Net {
   void add_after_backward(Callback* value) {
     after_backward_.push_back(value);
   }
-
+  void set_iter(int iter,int max_iter) {
+    iter_ = iter;
+    max_iter_ = max_iter;
+  }
  protected:
   // Helpers for Init.
   /// @brief Append a new top blob to the net.
@@ -335,6 +338,8 @@ class Net {
   vector<Callback*> after_forward_;
   vector<Callback*> before_backward_;
   vector<Callback*> after_backward_;
+  int iter_;
+  int max_iter_;
 
 DISABLE_COPY_AND_ASSIGN(Net);
 };

diff --git a/src/caffe/net.cpp b/src/caffe/net.cpp
@@ -523,6 +523,7 @@ Dtype Net<Dtype>::ForwardFromTo(int start, int end) {
     for (int c = 0; c < before_forward_.size(); ++c) {
       before_forward_[c]->run(i);
     }
+    layers_[i]->set_iter(iter_,max_iter_);
     Dtype layer_loss = layers_[i]->Forward(bottom_vecs_[i], top_vecs_[i]);
     loss += layer_loss;
     if (debug_info_) { ForwardDebugInfo(i); }

diff --git a/src/caffe/solver.cpp b/src/caffe/solver.cpp
@@ -222,6 +222,7 @@ void Solver<Dtype>::Step(int iters) {
   while (iter_ < stop_iter) {
     // zero-init the params
     net_->ClearParamDiffs();
+    net_->set_iter(iter_,stop_iter);
     if (param_.test_interval() && iter_ % param_.test_interval() == 0
         && (iter_ > 0 || param_.test_initialization())) {
       if (Caffe::root_solver()) {