updating batch norm

pkmital · pkmital · commit fc8f7e939867 · 2016-03-05T22:23:35.000-08:00
diff --git a/python/06_modern_convnet.py b/python/06_modern_convnet.py
@@ -16,15 +16,27 @@
 mnist = MNIST()
 x = tf.placeholder(tf.float32, [None, 784])
 y = tf.placeholder(tf.float32, [None, 10])
+
+# %% We add a new type of placeholder to denote when we are training.
+# This will be used to change the way we compute the network during
+# training/testing.
+is_training = tf.placeholder(tf.bool, name='is_training')
+
+# %% We'll convert our MNIST vector data to a 4-D tensor:
+# N x W x H x C
 x_tensor = tf.reshape(x, [-1, 28, 28, 1])
 
-# %% Define the network:
-bn1 = batch_norm(-1, name='bn1')
-bn2 = batch_norm(-1, name='bn2')
-bn3 = batch_norm(-1, name='bn3')
-h_1 = lrelu(bn1(conv2d(x_tensor, 32, name='conv1')), name='lrelu1')
-h_2 = lrelu(bn2(conv2d(h_1, 64, name='conv2')), name='lrelu2')
-h_3 = lrelu(bn3(conv2d(h_2, 64, name='conv3')), name='lrelu3')
+# %% We'll use a new method called  batch normalization.
+# This process attempts to "reduce internal covariate shift"
+# which is a fancy way of saying that it will normalize updates for each
+# batch using a smoothed version of the batch mean and variance
+# The original paper proposes using this before any nonlinearities
+h_1 = lrelu(batch_norm(conv2d(x_tensor, 32, name='conv1'),
+                       is_training, scope='bn1'), name='lrelu1')
+h_2 = lrelu(batch_norm(conv2d(h_1, 64, name='conv2'),
+                       is_training, scope='bn2'), name='lrelu2')
+h_3 = lrelu(batch_norm(conv2d(h_2, 64, name='conv3'),
+                       is_training, scope='bn3'), name='lrelu3')
 h_3_flat = tf.reshape(h_3, [-1, 64 * 4 * 4])
 h_4 = linear(h_3_flat, 10)
 y_pred = tf.nn.softmax(h_4)
@@ -48,9 +60,10 @@
     for batch_i in range(mnist.train.num_examples // batch_size):
         batch_xs, batch_ys = mnist.train.next_batch(batch_size)
         sess.run(train_step, feed_dict={
-            x: batch_xs, y: batch_ys})
+            x: batch_xs, y: batch_ys, is_training: True})
     print(sess.run(accuracy,
                    feed_dict={
                        x: mnist.validation.images,
-                       y: mnist.validation.labels
+                       y: mnist.validation.labels,
+                       is_training: False
                    }))
diff --git a/python/libs/batch_norm.py b/python/libs/batch_norm.py
@@ -1,88 +1,63 @@
 """Batch Normalization for TensorFlow.
-Parag K. Mital, Jan 2016."""
+Parag K. Mital, Jan 2016.
+"""
 
 import tensorflow as tf
+from tensorflow.python import control_flow_ops
 
 
-class batch_norm(object):
-    """Basic usage from: http://stackoverflow.com/a/33950177
-
-    Parag K. Mital, Jan 2016
-
-    Attributes
-    ----------
-    batch_size : int
-        Size of the batch.  Set to -1 to fit to current net.
-    beta : Tensor
-        A 1D beta Tensor with size matching the last dimension of t.
-        An offset to be added to the normalized tensor.
-    ema : tf.train.ExponentialMovingAverage
-        For computing the moving average.
-    epsilon : float
-        A small float number to avoid dividing by 0.
-    gamma : Tensor
-        If "scale_after_normalization" is true, this tensor will be multiplied
-        with the normalized tensor.
-    momentum : float
-        The decay to use for the moving average.
-    name : str
-        The variable scope for all variables under batch normalization.
+def batch_norm(x, phase_train, scope='bn', affine=True):
     """
+    Batch normalization on convolutional maps.
 
-    def __init__(self, batch_size, epsilon=1e-5,
-                 momentum=0.1, name="batch_norm"):
-        """Summary
-
-        Parameters
-        ----------
-        batch_size : int
-            Size of the batch, or -1 for size to fit.
-        epsilon : float, optional
-            A small float number to avoid dividing by 0.
-        momentum : float, optional
-            Decay to use for the moving average.
-        name : str, optional
-            Variable scope will be under this prefix.
-        """
-        with tf.variable_scope(name) as scope:
-            self.epsilon = epsilon
-            self.momentum = momentum
-            self.batch_size = batch_size
-            self.ema = tf.train.ExponentialMovingAverage(decay=self.momentum)
-            self.name = name
-
-    def __call__(self, x, train=True):
-        """Applies/updates the BN to the input Tensor.
+    from: https://stackoverflow.com/questions/33949786/how-could-i-
+    use-batch-normalization-in-tensorflow
 
-        Parameters
-        ----------
-        x : Tensor
-            The input tensor to normalize.
-        train : bool, optional
-            Whether or not to train parameters.
+    Only modified to infer shape from input tensor x.
 
-        Returns
-        -------
-        x_normed : Tensor
-            The normalized Tensor.
-        """
+    Parameters
+    ----------
+    x
+        Tensor, 4D BHWD input maps
+    phase_train
+        boolean tf.Variable, true indicates training phase
+    scope
+        string, variable scope
+    affine
+        whether to affine-transform outputs
+
+    Return
+    ------
+    normed
+        batch-normalized maps
+    """
+    with tf.variable_scope(scope):
         shape = x.get_shape().as_list()
 
-        # Using a variable scope means any new variables
-        # will be prefixed with "variable_scope/", e.g.:
-        # "variable_scope/new_variable".  Also, using
-        # TensorBoard, this will make everything very
-        # nicely grouped.
-        with tf.variable_scope(self.name) as scope:
-            self.gamma = tf.get_variable(
-                "gamma", [shape[-1]],
-                initializer=tf.random_normal_initializer(1., 0.02))
-            self.beta = tf.get_variable(
-                "beta", [shape[-1]],
-                initializer=tf.constant_initializer(0.))
-
-            mean, variance = tf.nn.moments(x, [0, 1, 2])
-
-            return tf.nn.batch_norm_with_global_normalization(
-                x, mean, variance, self.beta, self.gamma, self.epsilon,
-                scale_after_normalization=True)
+        beta = tf.Variable(tf.constant(0.0, shape=[shape[-1]]),
+                           name='beta', trainable=True)
+        gamma = tf.Variable(tf.constant(1.0, shape=[shape[-1]]),
+                            name='gamma', trainable=affine)
+
+        batch_mean, batch_var = tf.nn.moments(x, [0, 1, 2], name='moments')
+        ema = tf.train.ExponentialMovingAverage(decay=0.9)
+        ema_apply_op = ema.apply([batch_mean, batch_var])
+        ema_mean, ema_var = ema.average(batch_mean), ema.average(batch_var)
+
+        def mean_var_with_update():
+            """Summary
+
+            Returns
+            -------
+            name : TYPE
+                Description
+            """
+            with tf.control_dependencies([ema_apply_op]):
+                return tf.identity(batch_mean), tf.identity(batch_var)
+        mean, var = control_flow_ops.cond(phase_train,
+                                          mean_var_with_update,
+                                          lambda: (ema_mean, ema_var))
+
+        normed = tf.nn.batch_norm_with_global_normalization(
+            x, mean, var, beta, gamma, 1e-3, affine)
+    return normed
