跨卡同步 Batch Normalization[转]

作者:张航 Amazon AI Applied Scientist

跨卡同步(Cross-GPU Synchronized)Batch Normalization或称为同步BN(SyncBN)。

写在前面:为什么要跨卡同步Batch Normalization

现有的便准Batch Normalization因为使用数据并行(Data Parallel),是单卡的实现模式,只对单个卡上对样本进行归一化,相当于减小了批量大小(batch-size)。对于比较消耗显存的训练任务时,往往单卡上的相对批量过小,影响模型的收敛效果。之前在我们在图像语义分割的实验中,发现使用大规模的效果反而变差,实际上就是BN在作怪。跨卡同步Batch Normalization 可以使用全局的样本进行归一化,这样相当于‘增大’了批大小,这样训练效果不再受到使用GPU数量的影响。最近在图像分割、物体检测的论文中,使用跨卡BN也会显著地提高实验结果,所以跨卡BN已然成为竞赛刷分、发论文的必备神器。

Batch Normalization如何工作

既然是技术贴,读者很多都是神学大牛,为什么还要在这里赘述BatchNorm这个简单概念呢?其实不然,很多做科研的朋友如果没有解决过相关问题,很容易混淆BN在训练和测试时候的工作方式。记得在17年CVPR的tutoial上,何凯明和RBG两位大神分别在自己的talk上都特意强调了BN的工作原理,可见就算台下都是CV学者,也有必要复习一遍这些知识。

  • 工作原理:

    BN有效地加速了模型训练,加大learning rate,让模型不再过度依赖初始化。它在训练时在网络内部进行归一化(normalization),为训练提供了有效的regularization,抑制过拟合,用原作者的话时防止了协方差偏移。这里上一张图来展示训练模式的BN:

    img

    其中输入样本,其均值为u,方差为sigma,BN的输出,时可学习对参数。个人认为,这种强大的效果其实来自于back-propagation时候,来自于均值和方差对输入样本的梯度。这也是BN在训练模式与其测试模式的重要区别,在测试模式时(evaluation mode)下,使用训练集上累积的均值和方差,在back-propagation的时候他们对输入样本没有梯度(gradient)。

  • 数据并行:

    深度学习平台在多卡(GPU)运算的时候都是采用的数据并行(DataParallel),如下图:

    img

    每次迭代,输入被等分成多份,然后分别在不同的卡上面前向(forward)和后向(backward)运算,并且求出梯度,在迭代完成后合并梯度、更新参数,再进行下一次迭代。因为再前向和后向运算的时候,每个卡上的模型是单都运算的,所以相应的Batch Normalization也是在卡内完成,所以实际BN所归一化的样本数量仅仅局限在卡内,相当于批量大小(batch-size)见笑了。

如何实现SyncBN

跨卡同步BN的关键是前向运算的时候拿到全局的均值和方差,在后向运算时得到相应的全局梯度。最简单的实现方法是先同步求均值,再发回各卡然后同步求方差,但是这样就同步了两次。实际上需要同步一次就可以了,我们使用了一个非常简单的技巧,如下图:

img

这样在前向运算的时候,我们只需要在各卡上算出与,再跨卡求出全局的和即可得到正确的均值和方差, 同理我们在后向运算的时候只需同步一次,求出相应的梯度。