BN without momentum

deephealthproject · Nov 26, 2019 · 17b86d2 · 17b86d2
1 parent 050424b
commit 17b86d2
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/layers/normalization/layer_batchnorm.cpp b/src/layers/normalization/layer_batchnorm.cpp
@@ -167,10 +167,18 @@ void LBatchNorm::forward() {
     }
   }
   else {
-    Tensor::copy(mean->output,layers[0]->output);
+    if (momentum!=0.0)
+      Tensor::copy(mean->output,layers[0]->output);
+    else layers[0]->forward();
+
     layers[1]->forward();
 
-    Tensor::copy(variance->output,layers[3]->output);
+    if (momentum!=0.0)
+      Tensor::copy(variance->output,layers[3]->output);
+    else {
+      layers[2]->forward();
+      layers[3]->forward();
+    }
     layers[4]->forward();
     layers[5]->forward();
     layers[6]->forward();