Add init iteration param for BatchNormLayer

kkhoot · kkhoot · commit 18ef30f82ada · 2015-11-10T22:52:29.000+09:00
diff --git a/include/caffe/common_layers.hpp b/include/caffe/common_layers.hpp
@@ -130,6 +130,7 @@ class BatchNormLayer : public Layer<Dtype> {
   Blob<Dtype> mean_, variance_, temp_, x_norm_;
   bool use_global_stats_;
   BatchNormParameter_AvgType avg_type_;
+  int init_iterations_;
   Dtype moving_average_fraction_;
   int channels_;
   Dtype eps_;
diff --git a/src/caffe/layers/batch_norm_layer.cpp b/src/caffe/layers/batch_norm_layer.cpp
@@ -13,6 +13,7 @@ void BatchNormLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
   moving_average_fraction_ = param.moving_average_fraction();
   use_global_stats_ = this->phase_ == TEST;
   avg_type_ = param.avg_type();
+  init_iterations_ = param.init_iterations();
   if (param.has_use_global_stats())
     use_global_stats_ = param.use_global_stats();
   if (bottom[0]->num_axes() == 1)
@@ -125,7 +126,8 @@ void BatchNormLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
         variance_.mutable_cpu_data());  // E((X_EX)^2)
 
     // compute and save moving average
-    if (avg_type_ == BatchNormParameter_AvgType_EXPONENTIAL) {
+    if (init_iterations_ < this->blobs_[2]->cpu_data()[0] ||
+        avg_type_ == BatchNormParameter_AvgType_EXPONENTIAL) {
       Dtype scale_factor = this->blobs_[2]->cpu_data()[0] == 0 ?
           1 : 1 - moving_average_fraction_;
       caffe_cpu_axpby(mean_.count(), scale_factor,
diff --git a/src/caffe/layers/batch_norm_layer.cu b/src/caffe/layers/batch_norm_layer.cu
@@ -60,7 +60,8 @@ void BatchNormLayer<Dtype>::Forward_gpu(const vector<Blob<Dtype>*>& bottom,
         variance_.mutable_gpu_data());  // E((X_EX)^2)
 
     // compute and save moving average
-    if (avg_type_ == BatchNormParameter_AvgType_EXPONENTIAL) {
+    if (init_iterations_ < this->blobs_[2]->cpu_data()[0] ||
+        avg_type_ == BatchNormParameter_AvgType_EXPONENTIAL) {
       Dtype scale_factor = this->blobs_[2]->cpu_data()[0] == 0 ?
           1 : 1 - moving_average_fraction_;
       caffe_gpu_axpby(mean_.count(), scale_factor,
diff --git a/src/caffe/proto/caffe.proto b/src/caffe/proto/caffe.proto
@@ -481,9 +481,12 @@ message BatchNormParameter {
   // How much does the moving average decay each iteration?
   // Used when avg_type = EXPONENTIAL
   optional float moving_average_fraction = 3 [default = .999];
+  // How many iterations to use to make initial moving average value?
+  // Used when avg_type = EXPONENTIAL
+  optional int32 init_iterations = 4 [default = 100];
   // Small value to add to the variance estimate so that we don't divide by
   // zero.
-  optional float eps = 4 [default = 1e-5];
+  optional float eps = 5 [default = 1e-5];
 }
 
 message ContrastiveLossParameter {
