Merge pull request #1062 from dwf/possibly_speedup_batch_norm

Speed up BatchNormalization training graphs
mila-iqia · Apr 20, 2016 · 0c9a029 · 0c9a029
2 parents 17373d4 + 1ef9ddb
commit 0c9a029
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/blocks/bricks/bn.py b/blocks/bricks/bn.py
@@ -190,8 +190,10 @@ def _compute_training_statistics(self, input_):
         if self.mean_only:
             stdev = tensor.ones_like(mean)
         else:
-            stdev = tensor.sqrt(tensor.var(input_, axis=axes, keepdims=True) +
-                                numpy.cast[theano.config.floatX](self.epsilon))
+            var = (tensor.sqr(input_).mean(axis=axes, keepdims=True) -
+                   tensor.sqr(mean))
+            eps = numpy.cast[theano.config.floatX](self.epsilon)
+            stdev = tensor.sqrt(var + eps)
             assert (stdev.broadcastable[1:] ==
                     self.population_stdev.broadcastable)
             add_role(stdev, BATCH_NORM_MINIBATCH_ESTIMATE)