batch normalization

example_cifar10.py

@@ -47,12 +47,14 @@ class Model(ModelDesc):
 
         l = Conv2D('conv1', image, out_channel=64, kernel_shape=5, padding='SAME',
                   W_init=tf.truncated_normal_initializer(stddev=1e-4))
+        #l = BatchNorm('bn0', l, is_training)
         l = MaxPooling('pool1', l, 3, stride=2, padding='SAME')
         l = tf.nn.lrn(l, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='norm1')
 
         l = Conv2D('conv2', l, out_channel=64, kernel_shape=5, padding='SAME',
                   W_init=tf.truncated_normal_initializer(stddev=1e-4),
                   b_init=tf.constant_initializer(0.1))
+        #l = BatchNorm('bn1', l, is_training)
         l = tf.nn.lrn(l, 4, bias=1.0, alpha=0.001 / 9.0, beta=0.75, name='norm2')
         l = MaxPooling('pool2', l, 3, stride=2, padding='SAME')
 

tensorpack/models/batch_norm.py

+#!/usr/bin/env python2
+# -*- coding: UTF-8 -*-
+# File: batch_norm.py
+# Author: Yuxin Wu <ppwwyyxx@gmail.com>
+
+import tensorflow as tf
+
+from ._common import layer_register
+
+__all__ = ['BatchNorm']
+
+
+# http://stackoverflow.com/questions/33949786/how-could-i-use-batch-normalization-in-tensorflow
+@layer_register()
+def BatchNorm(x, is_training):
+    """
+    x: has to be BHWC for now
+    is_training: bool
+    """
+    is_training = bool(is_training)
+    shape = x.get_shape().as_list()
+    assert len(shape) == 4
+
+    n_out = shape[-1]  # channel
+    beta = tf.get_variable('beta', [n_out])
+    gamma = tf.get_variable('gamma', [n_out],
+                        initializer=tf.constant_initializer(1.0))
+    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)
+
+    if is_training:
+        def mean_var_with_update():
+            with tf.control_dependencies([ema_apply_op]):
+                return tf.identity(batch_mean), tf.identity(batch_var)
+        mean, var = mean_var_with_update()
+    else:
+        mean, var = ema_mean, ema_var
+
+    normed = tf.nn.batch_norm_with_global_normalization(x, mean, var, beta, gamma, 1e-4, True)
+    return normed
+

tensorpack/models/image_sample.py

@@ -69,10 +69,10 @@ def ImageSample(inputs):
     diffy, diffx = tf.split(3, 2, diff)
     neg_diffy, neg_diffx = tf.split(3, 2, neg_diff)
 
-    prod = tf.reduce_prod(diff, 3, keep_dims=True)
-    diff = tf.Print(diff, [tf.is_finite(tf.reduce_sum(diff)), tf.shape(prod),
-                          tf.reduce_max(diff), diff],
-                    summarize=50)
+    #prod = tf.reduce_prod(diff, 3, keep_dims=True)
+    #diff = tf.Print(diff, [tf.is_finite(tf.reduce_sum(diff)), tf.shape(prod),
+                          #tf.reduce_max(diff), diff],
+                    #summarize=50)
 
     return tf.add_n([sample(template, lcoor) * neg_diffx * neg_diffy,
            sample(template, ucoor) * diffx * diffy,