前记

最近感觉Batch Normalization在实践上还是很有搞头的，首先在单机TensorFlow里完整实现BN演算法就有很多坑（具体可以看我的文章Michael：[论文阅读]Batch Normalization: Accelerating Deep Netwo，里面也有我对BN的理解），然后在单机多卡上的实现也有些坑，主要是moving mean和moving variance在单机多卡上是怎么实现更新的（分散式的原理类似於单机多卡）。

首先做一些实验来看看BN的运行原理吧，代码已上传至testBNInMultiGPU.py，我改了一些bug后可以完美运行并在实际项目中运行成功的代码在validSyncBN.py。

单机单卡上的BN实验

代码为testBNInSingleGPU()。

结果：

outputArray = [0. 0. 0.]

After normalization:
moving mean = [0.1 0.1 0.1]
moving variance = [0.9 0.9 0.9]

首先看输出全为0表示BN起到了归一化的作用。

之后可以看moving mean和moving variance的值，参考Michael：TensorFlow源码解读之Batch Normalization，得知TF中moving mean和moving variance的更新方式都是momentum，即 x=m imes x+(1-m) imes E_B(x)=x+(1-m) imes(E_B(x)-x) 。

因为moving mean初始化为0，moving variance初始化为1，momentum设置为0.9，代入可得：

moving mean： 0+(1-0.9) imes(1-0)=0.1

moving variance： 1+(1-0.9) imes(0-1)=0.9

和实验是完美符合的。

单机多卡上的运行实验

这里选取了2块卡进行实验，代码为testBNInMultiGPU()。

结果：

outputArray = [0. 0. 0.]

After normalization:
moving mean = [0.2 0.2 0.2]
moving variance = [0.79999995 0.79999995 0.79999995]

由于是两块卡，结合Michael：TensorFlow源码解读之Batch Normalization的分析，再根据实验结果，我推测更新公式是这样的：

统计每张卡上moving mean和moving variance的变化量，即 Delta x=x-x=(1-m) imes(E_B(x)-x) ，将各卡的变化量求和得到总变化量，进而得到更新的结果：

x=x+sum^N_1{Delta x}=x+sum^N_1{(1-m) imes(E_B(x)-x)} ，这里 N 表示GPU的数量。

因为moving mean初始化为0，moving variance初始化为1，momentum设置为0.9，代入可得：

moving mean： 0+(1-0.9) imes(1-0)+(1-0.9) imes(1-0)=0.2

moving variance： 1+(1-0.9) imes(0-1)+(1-0.9) imes(0-1)=0.8

和实验是完美符合的。

另一种更新的策略是，每张卡前向执行一次，更新一次moving mean和moving variance，那么：

moving mean：

第一次： 0+(1-0.9) imes(1-0)=0.1

第二次： 0.1+(1-0.9) imes(1-0.1)=0.19

这与实验结果是不符合的。

这里改变一下输入的值（一个卡上是1，一个卡上是2），以验证我的想法是不是对的：

for item in range(numberGPU):
feedDict[tf.get_default_graph().get_tensor_by_name("tower_%s/data:0" % item)]
= np.ones((3))*(item+1)

结果：

outputArray = [0. 0. 0.]

After normalization:
moving mean = [0.3 0.3 0.3]
moving variance = [0.79999995 0.79999995 0.79999995]

因为moving mean初始化为0，moving variance初始化为1，momentum设置为0.9，代入可得：

moving mean： 0+(1-0.9) imes(1-0)+(1-0.9) imes(2-0)=0.3

moving variance： 1+(1-0.9) imes(0-1)1+(1-0.9) imes(0-1)=0.8

和实验是完美符合的。

影响

对于公式 x=x+sum^N_1{Delta x}=x+sum^N_1{(1-m) imes(E_B(x)-x)} ，如果各个卡上 x 的分布是一致的，那么很清晰就得到：

Delta x=N imes (1-m) imes(E_B(x)-x)

这个实际是会加速moving mean和moving variance的，所以可设置新的moment为： m=frac{N+m-1}{N} 。

第二个影响就是对训练了，从上述的实验可以猜测BN的训练过程：

在每张卡上前向推导，得到BN的moving mean和moving variance的改变数，同时，反向推导，得到 gamma 和 eta 在每张卡上的梯度。

然后，使用tensorflow 多GPU编程 完全指南中描述的方法，得到gamma 和 eta的平均梯度。

所以，TensorFlow的官方方法是没有考虑多GPU编程的，下一个步骤好好搞搞这个吧，目测比较难。

目前找到的一些比较好的资源：

batch normalization的multi-GPU版本该怎么实现？

里面有几个比较好的解决方案：

1.jianlong-yuan/syncbn-tensorflow重写了TensorFlow的官方方法，可以做个实验验证一下。

2.旷视科技：CVPR 2018 | 旷视科技物体检测冠军论文——大型Mini-Batch检测器MegDet

tensorpack/tensorpack里面有该论文的代码实现。

关于第一种解决方案的实验论证与理论解释

jianlong-yuan/syncbn-tensorflow是源代码，我稍微改了改代码，并添加了我的一些注释，代码为sync_batch_norm()。

测试代码为testBNInMultiGPU3()。

输出为：

outputTuple = [array([-0.998006, -0.998006, -0.998006], dtype=float32), array([0.9980061, 0.9980061, 0.9980061], dtype=float32),
[array([0.15, 0.15, 0.15], dtype=float32), array([0.92499995, 0.92499995, 0.92499995], dtype=float32)]]

After normalization:
moving mean = [0.15 0.15 0.15]
moving variance = [0.92499995 0.92499995 0.92499995]

outputTuple一共4项，前两项为2个GPU的输出量，一个为-1，一个为1，后两项分别是moving mean和moving variance的输出。

这里先就数据的观点来验证结果是否正确。

GPU_0的数据是3个1，GPU_1的数据是3个2，参考batch normalization accelerating deep network training by reducing internal covariate shift演算法1，有：

mu_B=frac{3 imes1+3 imes2}{6}=1.5

sigma^2_B=frac{3 imes(1-1.5)^2+3 imes(2-1.5)^2}{6}=0.25

对于GPU_0的数据，有：

hat x=frac{x-mu_B}{sqrt{sigma^2_B+epsilon}}=frac{1-1.5}{sqrt{0.25}}=-1

y=gamma hat x+eta=1 imes-1+0=-1

对于GPU_1的数据，有：

hat x=frac{x-mu_B}{sqrt{sigma^2_B+epsilon}}=frac{2-1.5}{sqrt{0.25}}=1

y=gamma hat x+eta=1 imes1+0=1

这说明输出是符合的。

下面看moving mean和moving variance的值。

moving mean： 0+(1-0.9) imes(1.5-0)=0.15

moving variance： 1+(1-0.9) imes(0.25-1)=0.925

也是符合的，说明程序还是很鲁棒的。

这里来分析程序的运行原理，大部分我都用注释说明了，这里讲里面的两个知识点吧，一个是mean和variance是怎么计算得到的，一个是moving mean和moving variance是怎么更新的。

首先说下mean和variance是怎么计算得到的，是以下的这段代码：

# avarage moving_mean and moving_var in multi GPUs
shared_name = tf.get_variable_scope().name
batch_mean = tf.reduce_mean(inputs, axis=axis)
batch_mean = gen_nccl_ops.nccl_all_reduce(
input=batch_mean,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean) * (1.0 / num_dev)
batch_mean_square = tf.reduce_mean(tf.square(inputs), axis=axis)
batch_mean_square = gen_nccl_ops.nccl_all_reduce(
input=batch_mean_square,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean_square) * (1.0 / num_dev)
mean = batch_mean
var = batch_mean_square - tf.square(batch_mean)

假设有N块卡，每个卡上有 m_i,i=1, 2, ...,N 个训练样本，且令 m_1=m_2=...=m_N ，那么有 m=m_1*N=m_2*N=...=m_N*N ，其中， m 是整个的样本数。

这句代码：

batch_mean = tf.reduce_mean(inputs, axis=axis)

其用数学表达式可以表示成： batch\_mean_i=frac{sum_{j=1}^{m_i}x_j}{m_i},for,i=1,2, ...,N

这句代码：

batch_mean = gen_nccl_ops.nccl_all_reduce(
input=batch_mean,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean) * (1.0 / num_dev)

其会在各卡对相应的变数进行操作，这里进行的操作是求和，最后除N即可，用数学表达式可以表达成：

batch\_mean=frac{sum_i^Nbatch\_mean_i}{N}=frac{sum_i^Nfrac{sum_{j=1}^{m_i}x_j}{m_i}}{N} ，因为有 m_1=m_2=...=m_N ，那么有：

batch\_mean=frac{sum_i^N sum_{j=1}^{m_i}x_j}{Nm_1}=frac{sum_{j=1}^mx_j}{m}

这个表达式就与论文里的表达式一致了，以下用 mu_B 代替 batch\_mean 。

这句代码：

batch_mean_square = tf.reduce_mean(tf.square(inputs), axis=axis)

其用数学表达式可以表示成： batch\_mean\_square_i=frac{sum_{j=1}^{m_i}x_j^2}{m_i},for , i=1,2,...,N

这句代码：

batch_mean_square = gen_nccl_ops.nccl_all_reduce(
input=batch_mean_square,
reduction=sum,
num_devices=num_dev,
shared_name=shared_name + _NCCL_mean_square) * (1.0 / num_dev)

其用数学表达式可以表示成： batch\_mean\_square=frac{sum_i^Nbatch\_mean\_square_i}{N}=frac{sum_i^Nfrac{sum_{j=1}^{m_i}x_j^2}{m_i}}{N}\ =frac{sum_i^Nsum_{j=1}^{m_i}x_j^2}{N m_1}=frac{sum_{j=1}^mx_j^2}{m}

这句代码：

var = batch_mean_square - tf.square(batch_mean)

其用数学表达式可以表示成： var=batch\_mean\_square-mu_B^2=frac{sum_{j=1}^mx_j^2}{m}-mu_B^2=frac{sum_{j=1}^mx_j^2-mmu_B^2}{m}

这里利用的一个trick是 mu_B=frac{sum_j^mx_j}{m} ，那么 var=frac{sum_{i=1}^m(x_i-mu_B)^2}{m}=frac{sum_{i=1}^mx_i^2-2mu_Bsum_{i=1}^mx_i+mu_B^2}{m}=frac{sum_{i=1}^mx_i^2-mu_B^2}{m}

另外一个值得注意的代码就是这个控制条件：

if int(outputs.device[-1]) == 0:

因为每个GPU上的moving mean和moving variance都一致了，所以没必要每个都更新，更新GPU:0的即可。

综上，看起来这个代码还是不错的，待我在实际项目中再验证验证吧。

在实际项目中的实验论证

说实话我是很想在实验中加入进去的，可惜的是，在我本地跑没有问题，但是在伺服器上跑最终依然有问题，所以最后可能得放弃这个了。

最开始碰到的问题在于Python 2中的对str的解释，报错为：

Traceback (most recent call last):
File "/home/dao/anaconda2/envs/python36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2349, in get_attr
c_api.TF_OperationGetAttrValueProto(self._c_op, name, buf)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Operation tower_0/lstm2fc/batch_normalization_0/NcclAllReduce_1 has no attr named _XlaCompile.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/dao/anaconda2/envs/python36/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py", line 391, in _MaybeCompile
xla_compile = op.get_attr("_XlaCompile")
File "/home/dao/anaconda2/envs/python36/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 2353, in get_attr
raise ValueError(str(e))
ValueError: Operation tower_0/lstm2fc/batch_normalization_0/NcclAllReduce_1 has no attr named _XlaCompile.
...

在搜索解决方案的时候我碰巧发现本文所写的第二种解决方案（即旷视论文）与第一种解决方案异曲同工（没兴趣去了解代码的根源了）。

第二种解决方案的代码在tensorpack/tensorpack/blob/master/tensorpack/models/batch_norm.py，首先给了我一些关于TF支持nccl的解释：

assert six.PY2 or TF_version >= (1, 10),
"Cross-GPU BatchNorm is only supported in TF>=1.10 ."
"Upgrade TF or apply this patch manually: https://github.com/tensorflow/tensorflow/pull/20360"

链接的issue是Fix gradient of nccl_ops by ppwwyyxx · Pull Request #20360 · tensorflow/tensorflow，个人猜测1.10以下的版本支持Python 2.7的居多，我用Python 3.6，所以报错了。

升级到了1.10.0的版本，发现代码这段加上也是很不错的：

if TF_version <= (1, 12):
try:
from tensorflow.contrib.nccl.python.ops.nccl_ops import _validate_and_load_nccl_so
except Exception:
pass
else:
_validate_and_load_nccl_so()
from tensorflow.contrib.nccl.ops import gen_nccl_ops
else:
from tensorflow.python.ops import gen_nccl_ops

弄好之后又报错：

tensorflow.python.framework.errors_impl.NotFoundError: libnccl.so.2

参考文章tensorflow1.12 多GPU协同训练报错tensorflow.python.framework.errors_impl.NotFoundError: libnccl.so.2发现需要安装nccl库，Nvidia官方的安装指导在Deep Learning SDK Documentation。

好歹不报错了，但是程序似乎不运行了，也查不出原因，所以决定暂时放弃这种做法了。

后续更新：排除bug

终于搞通了，发现是代码的bug，Tensorflow的条件控制没有if else这么简单，必须写到计算图里面去。

这会导致一个什么结果呢？在推理阶段，不需要更新moving mean和moving variance，所以代码会直接推理得到结果：

outputs, _, _ = tf.nn.fused_batch_norm(inputs, gamma, beta, mean=moving_mean, variance=moving_var, epsilon=epsilon, is_training=False)

但是呢，这是一个else语句，计算图是没法把Python里的判断条件放到计算图里去的，如果列印计算图，会发现此时根本就没有直接推理的结果：

print ([n.name for n in tf.get_default_graph().as_graph_def().node if "fused_batch_norm" in n.name])

使用tf.cond()写好了判断函数后，还发现一个问题，在我真实的项目环境中是可以使用的，但是我在写的测试代码中却会报这样的错：

"operation has been marked as not fetchable"

参考Queue operation in conditional execution context fails with "operation has been marked as not fetchable" · Issue #4094 · tensorflow/tensorflow，控制条件不能有太多控制语句，我发现True分支里写了moving mean和moving variance的更新语句，这个会在分支里引入过多的输入，我猜测TensorFlow对此做了限定，在实验中我不在sess.run()中求tf.get_collection(tf.GraphKeys.UPDATE_OPS)的值也是可以不会报错的，综上，我认为错误在于更新语句，也做了一些改进，详细代码见validSyncBN.py。

【已完结】


推荐阅读：
相关文章