Returned shift and scale back to BN layer + stability improvement

Author: drnikolaev

include/caffe/layers/batch_norm_layer.hpp

@@ -7,6 +7,9 @@
 #include "caffe/layer.hpp"
 #include "caffe/proto/caffe.pb.h"
 
+#define BN_VARIANCE_CLIP_START 200
+#define BN_VARIANCE_CLIP_CONST 4.0
+
 namespace caffe {
 
 /**
@@ -63,17 +66,30 @@ class BatchNormLayer : public Layer<Dtype> {
   virtual void Backward_gpu(const vector<Blob<Dtype>*>& top,
      const vector<bool>& propagate_down, const vector<Blob<Dtype>*>& bottom);
 
-  Blob<Dtype> mean_, variance_, temp_, x_norm_;
+  virtual void multicast_cpu(int N, int C, int S, const Dtype *x, Dtype *y);
+  virtual void compute_sum_per_channel_cpu(int N, int C, int S,
+      const Dtype *x, Dtype *y);
+  virtual void compute_mean_per_channel_cpu(int N, int C, int S,
+      const Dtype *x, Dtype *y);
+#ifndef CPU_ONLY
+  virtual void compute_sum_per_channel_gpu(int N, int C, int S,
+      const Dtype *x, Dtype *y);
+  virtual void multicast_gpu(int N, int C, int S, const Dtype *x, Dtype *y);
+  virtual void compute_mean_per_channel_gpu(int N, int C, int S,
+      const Dtype *x, Dtype *y);
+#endif
+
+  Blob<Dtype> mean_, variance_, inv_variance_, x_norm_;
   bool use_global_stats_;
   Dtype moving_average_fraction_;
   int channels_;
   Dtype eps_;
-
-  // extra temporarary variables is used to carry out sums/broadcasting
-  // using BLAS
-  Blob<Dtype> batch_sum_multiplier_;
-  Blob<Dtype> num_by_chans_;
-  Blob<Dtype> spatial_sum_multiplier_;
+  int iter_;
+  // auxiliary arrays
+  Blob<Dtype> ones_N_, ones_HW_, ones_C_;
+  Blob<Dtype> temp_;
+  Blob<Dtype> temp_C_;
+  Blob<Dtype> temp_NC_;
 };
 
 }  // namespace caffe

include/caffe/layers/cudnn_batch_norm_layer.hpp

@@ -16,7 +16,7 @@ template <typename Dtype>
 class CuDNNBatchNormLayer : public BatchNormLayer<Dtype> {
  public:
   explicit CuDNNBatchNormLayer(const LayerParameter& param)
-      : BatchNormLayer<Dtype>(param), epsilon_(1e-4), handles_setup_(false) {}
+      : BatchNormLayer<Dtype>(param), handles_setup_(false) {}
   virtual void LayerSetUp(const vector<Blob<Dtype>*>& bottom,
       const vector<Blob<Dtype>*>& top);
   virtual void Reshape(const vector<Blob<Dtype>*>& bottom,
@@ -34,7 +34,6 @@ class CuDNNBatchNormLayer : public BatchNormLayer<Dtype> {
   cudnnTensorDescriptor_t scale_bias_mean_var_desc_;
   cudnnBatchNormMode_t mode_;
 
-  double epsilon_;
   Blob<Dtype> save_mean_, save_inv_var_;
   bool handles_setup_;
 };

include/caffe/util/math_functions.hpp

@@ -33,6 +33,16 @@ template <typename Dtype>
 void caffe_cpu_axpby(const int N, const Dtype alpha, const Dtype* X,
     const Dtype beta, Dtype* Y);
 
+// y[i]= max(a*x[i], b*y[i])
+template <typename Dtype>
+void caffe_cpu_eltwise_max(const int N, const Dtype alpha, const Dtype* X,
+    const Dtype beta, Dtype* Y);
+
+// y[i]= min(a*x[i], b*y[i])
+template <typename Dtype>
+void caffe_cpu_eltwise_min(const int N, const Dtype alpha, const Dtype* X,
+    const Dtype beta, Dtype* Y);
+
 template <typename Dtype>
 void caffe_copy(const int N, const Dtype *X, Dtype *Y);
 
@@ -246,6 +256,17 @@ void caffe_gpu_fabs(const int n, const Dtype* x, Dtype* y);
 template <typename Dtype>
 void caffe_gpu_scale(const int n, const Dtype alpha, const Dtype *x, Dtype* y);
 
+// y[i]= max(a*x[i], b*y[i])
+template <typename Dtype>
+void caffe_gpu_eltwise_max(const int n, const Dtype alpha, const Dtype* x,
+    const Dtype beta, Dtype* y);
+
+// y[i]= min(a*x[i], b*y[i])
+template <typename Dtype>
+void caffe_gpu_eltwise_min(const int n, const Dtype alpha, const Dtype* x,
+    const Dtype beta, Dtype* y);
+
+
 #define DEFINE_AND_INSTANTIATE_GPU_UNARY_FUNC(name, operation) \
 template<typename Dtype> \
 __global__ void name##_kernel(const int n, const Dtype* x, Dtype* y) { \