tensorflow中如何正确使用BatchNorm（涵盖所有BN的坑）

                                                      tensorflow中如何正确使用BatchNorm（涵盖所有BN的坑）

      在caffe中使用BatchNorm层很简单，只要注意一点，在训练时将use_global_states设为false，测试前向阶段将use_global_states设为true即可。在tensorflow中使用batchnorm层有几个地方需要注意，不然会踩坑导致训练不收敛或者测试时准确率降低很多，推荐使用tf.layers.batch_normalization函数实现BN归一化。

 

1、训练时training设置为True，测试时training设置为False

 

BN层中有两个参数是需要学习和进行反向传播的，分别是尺度参数gamma和偏移因参数beta，目的是利用这两个参数进行线性变换，以恢复输入数据本身的表征能力。 所以在训练阶段和测试阶段需要对trainable设置不同的值，以保证待学习的参数在训练过程中可以进行反向传播。

训练：

x = tf.layers.batch_normalization(x, axis=3, training=True)

测试：

x = tf.layers.batch_normalization(x, axis=3, training=False)

2、均值和方差的更新

moving_mean和moving_variance是不可训练的变量，计算的是滑动平均和滑动方差，所以

需要注意的是，该方法内部维护的moving_mean和moving_variance（实际上底层是某个keras.layers.BatchNormalization类对象维护的成员变量）需要在训练时被更新，而更新操作符在tf.GraphKeys.UPDATE_OPS中，因此需要加到train_op的依赖里train_op的依赖中，这样在训练过程中，这些变量才会被更新：

 g_vars = tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES, 'generator_net')
        with tf.name_scope('train_optimizer'):
            update_ops_g = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies([tf.group(*update_ops_g)]):
                train_step_g = tf.train.AdamOptimizer(learning_rate=self.lr_g, beta1=self.beta1,
                                                   beta2=self.beta2).minimize(losses, var_list=g_vars)
saver = tf.train.Saver(max_to_keep=self.max_to_keep)

由于我对于tf.train.Saver的参数var_list没有赋值，所以默认是所有变量都会加入到训练更新中，可以通过var_list = saver._var_list()来获取验证是否所有trainable和非trainable的变量都会在var_list中，通过var_list = saver._var_list，打印信息为：

我们发现可训练与需要更新的变量都在var_list中，这样就可以正确使用BN层了。不过有时这样保存所有模型参数会导致保存的ckpt模型很大，所以我们可以选择只保存可训练的网络参数与BN层的更新参数，代码为：

g_vars = tf.get_collection(
            tf.GraphKeys.TRAINABLE_VARIABLES, 'generator_net')
        with tf.name_scope('train_optimizer'):
            update_ops_g = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies([tf.group(*update_ops_g)]):
                train_step_g = tf.train.AdamOptimizer(learning_rate=self.lr_g, beta1=self.beta1,
                                                      beta2=self.beta2).minimize(losses, var_list=g_vars)

        list_all = tf.global_variables()
        bn_moving_vars = [g for g in list_all if 'moving_mean' in g.name]
        bn_moving_vars += [g for g in list_all if 'moving_variance' in g.name]
        var_all = g_vars+bn_moving_vars

        saver = tf.train.Saver(max_to_keep=self.max_to_keep,var_list = var_all)

由于tf.GraphKeys.TRAINABLE_VARIABLES只会收集可训练的变量，是无法收集到BN的moving_mean和moving_variance的，所以需要先利用tf.global_variables()获取到所有变量，然后通过关键字符‘moving_mean’和‘moving_variance’找出待更新的变量合并到可训练的变量列表中即可。这样保存的ckpt模型尺寸比较小，也可以正确使用BN层进行训练。

3、batch_size的设置

训练含有BN层的深度网络时，切记batch_size不要设置为1，要大于1，不然无法正确计算更新滑动均值与滑动方差，测试时提取出滑动均值moving_mean会发现特别小，而滑动方差moving_variance直接为nan，详细查bug之旅可见我的另一篇博客：https://blog.csdn.net/LxDamon/article/details/108762202。

有问题欢迎评论交流，一起进步！