Returned shift and scale back to BN layer

BVLC · Apr 18, 2016 · 2c4cfbd · 2c4cfbd
1 parent c9eda39
commit 2c4cfbd
Show file tree

Hide file tree

Showing 8 changed files with 956 additions and 318 deletions.
diff --git a/include/caffe/layers/batch_norm_layer.hpp b/include/caffe/layers/batch_norm_layer.hpp
@@ -63,17 +63,32 @@ 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_;
+  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_;
+  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 multicast_gpu(int N, int C, int S, const Dtype *x, Dtype *y);
+  virtual void compute_sum_per_channel_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
+
+  // auxiliary arrays
+  Blob<Dtype> ones_N_, ones_HW_;
+
+  Blob<Dtype> temp_;
+  Blob<Dtype> temp_C_;
+  Blob<Dtype> temp_NC_;
 };
 
 }  // namespace caffe

diff --git a/include/caffe/layers/cudnn_batch_norm_layer.hpp b/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_;
 };