BN层的添加实战

对于一个小白，从了解Batch Normalization（后面简称BN）到正确使用BN,可谓路漫漫兮。在此做一个记录。

网上搜索关于BN最多的就是原理推导，相关论文出处。

例如：

http://blog.csdn.net/Fate_fjh/article/details/53375881

https://www.jianshu.com/p/0312e04e4e83

但是这个并不能帮助我们实际的使用，对于需要迅速用起来的伙伴帮助不大。我们工程师相信的是先用起来，再去研究原理！呵呵！

有一些文章介绍的BN层的实现，也有代码示例，但能顺利跑起来的寥寥。因为使用BN不像卷积层那样，写个层的实现就可以了。由于BN层会包含两个可训练参数以及两个不可训练参数，所以涉及到在train代码中如何保存的关键问题，以及在inference代码中如何加载的问题。有相关博客介绍到这一步了，很有帮助。

例如：

https://www.cnblogs.com/hrlnw/p/7227447.html

本以为别人都说这么明白了，抄一抄不是很容易的事情吗。可以上的代码是不能让你正确完成BN功能的。也不知是抄错了，还是别人漏掉了一些关键环节。总之你的moving_mean/moving_variance好像就是不太对。基本上中文网页很难在找到这个问题的解了。

现在你需要搜索的关键字可能要变成BN/参数保存/平均滑动等等了。还好tensorflow的github中有了线索：

https://github.com/tensorflow/tensorflow/issues/14809

https://github.com/tensorflow/tensorflow/issues/15250

可见有很多人确实无法正确使用BN功能，然而最有用的一个issues是：

https://github.com/tensorflow/tensorflow/issues/1122#issuecomment-280325584

在这里，我拼凑成了一个完整能用的BN功能代码，解决了我好久的痛苦，让我兴奋一下。

知识来源于网络，奉献给网络。不敢独享这一成果，再此分享给大家。

-----------------------------------------------------------------华丽的分割线----------------------------------------------------------------------------

整个BN功能的实现需要分三个部分：1.BN层实现；2.训练时更新和完成后保存；3.预测时加载。

1.BN层实现：

如果你接触了一段时间后，这里你至少应该知道BN的三种实现方式了，但是我只成功了其中的一种，希望其他朋友能够补充完善。

def bn_layer(x, scope, is_training, epsilon=0.001, decay=0.99, reuse=None):

    """

    Performs a batch normalization layer

    Args:

        x: input tensor

        scope: scope name

        is_training: python boolean value

        epsilon: the variance epsilon - a small float number to avoid dividing by 0

        decay: the moving average decay

    Returns:

        The ops of a batch normalization layer

    """

    with tf.variable_scope(scope, reuse=reuse):

        shape = x.get_shape().as_list()

        # gamma: a trainable scale factor

        gamma = tf.get_variable(scope+"_gamma", shape[-1], initializer=tf.constant_initializer(1.0), trainable=True)

        # beta: a trainable shift value

        beta = tf.get_variable(scope+"_beta", shape[-1], initializer=tf.constant_initializer(0.0), trainable=True)

        moving_avg = tf.get_variable(scope+"_moving_mean", shape[-1], initializer=tf.constant_initializer(0.0), trainable=False)

        moving_var = tf.get_variable(scope+"_moving_variance", shape[-1], initializer=tf.constant_initializer(1.0), trainable=False)

        if is_training:

            # tf.nn.moments == Calculate the mean and the variance of the tensor x

            avg, var = tf.nn.moments(x, np.arange(len(shape)-1), keep_dims=True)

            avg=tf.reshape(avg, [avg.shape.as_list()[-1]])

            var=tf.reshape(var, [var.shape.as_list()[-1]])

            #update_moving_avg = moving_averages.assign_moving_average(moving_avg, avg, decay)

            update_moving_avg=tf.assign(moving_avg, moving_avg*decay+avg*(1-decay))

            #update_moving_var = moving_averages.assign_moving_average(moving_var, var, decay)

            update_moving_var=tf.assign(moving_var, moving_var*decay+var*(1-decay))

            control_inputs = [update_moving_avg, update_moving_var]

        else:

            avg = moving_avg

            var = moving_var

            control_inputs = []

        with tf.control_dependencies(control_inputs):

            output = tf.nn.batch_normalization(x, avg, var, offset=beta, scale=gamma, variance_epsilon=epsilon)

    return output

def bn_layer_top(x, scope, is_training, epsilon=0.001, decay=0.99):

    """

    Returns a batch normalization layer that automatically switch between train and test phases based on the

    tensor is_training

    Args:

        x: input tensor

        scope: scope name

        is_training: boolean tensor or variable

        epsilon: epsilon parameter - see batch_norm_layer

        decay: epsilon parameter - see batch_norm_layer

    Returns:

        The correct batch normalization layer based on the value of is_training

    """

    #assert isinstance(is_training, (ops.Tensor, variables.Variable)) and is_training.dtype == tf.bool

    return tf.cond(

        is_training,

        lambda: bn_layer(x=x, scope=scope, epsilon=epsilon, decay=decay, is_training=True, reuse=None),

        lambda: bn_layer(x=x, scope=scope, epsilon=epsilon, decay=decay, is_training=False, reuse=True),

    )

这里的参数epsilon=0.001, decay=0.99可以自行调整。

 

2.训练时更新和完成后保存：

在训练的代码中增加如下代码:

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

with tf.control_dependencies(update_ops):

    train = tf.train.AdamOptimizer(learning_rate=lr).minimize(cost)

这个是用于更新参数的。

var_list = tf.trainable_variables()

g_list = tf.global_variables()

bn_moving_vars = [gfor gin g_listif 'moving_mean' in g.name]

bn_moving_vars += [gfor gin g_listif 'moving_variance' in g.name]

var_list += bn_moving_vars

train_saver = tf.train.Saver(var_list=var_list)

这个是用于保存bn不可训练的参数。

3.预测时加载：

# get moving avg

var_list = tf.trainable_variables()

g_list = tf.global_variables()

bn_moving_vars = [gfor gin g_listif 'moving_mean' in g.name]

bn_moving_vars += [gfor gin g_listif 'moving_variance' in g.name]

var_list += bn_moving_vars

saver = tf.train.Saver(var_list=var_list)

ckpt_path =""

saver.restore(sess, ckpt_path)

这样就可以找到checkpoint中的参数了。