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

链接：https://www.zhihu.com/question/59321480

编辑：深度学习与计算机视觉

声明：仅做学术分享，侵删

batch normalization中 的multi-GPU实现涉及到大量GPU间通信，这时效率就会很慢。当前各个平台（caffe， torch）的batch normalization的实现在multi-GPU情况下是否只考虑了单个GPU上的均值与方差？

作者：张航
https://www.zhihu.com/question/59321480/answer/198077371

终于搞定了Synchronized BatchNorm，来回答一下这个问题。

首先针对问题本身，目前所有的framework，包括Caffe，Torch，TF，PyTroch等等，BatchNorm的实现都是只考虑了single gpu。也就是说BN使用的均值和标准差是单个gpu算的，相当于缩小了mini－batch size。至于为什么这样实现，1）因为没有sync的需求，因为对于大多数vision问题，单gpu上的mini－batch已经够大了，完全不会影响结果。2）影响训练速度，BN layer通常是在网络结构里面广泛使用的，这样每次都同步一下GPUs，十分影响训练速度。

关于如何实现，这里不得不写一下心得，从发现问题到实现synchronize，花了将近一周的时间，一开始把这个问题想简单了，希望对大家有所帮助吧。偶然中发现自己做的task十分依赖synchronzie BN （因为在公司实习， 不方便说project）。然后开始实现：

1、因为通常训练过程是将网络复制到不同的gpu上，然后进行forward和backward，之后只需要collect gradient，再更新主gpu上的网络，然后下一个iteration再复制一遍。在这种方式下，不同gpu的BN layer是没办法通信的。所以第一步是重做了DataParallel架构，可以enable 每层的通信。

2、第二步就是推推公式，相信各位大神花点儿时间动动笔就搞定，我就不赘述了。然后开始编程写kernel，核心思想是每次forward 时候sync一下均值和方差，然后backward的时候sync一些grad of 均值和方差，再继续backward。

3、本身前两步搞定之后，应该一切顺利了，可是训练的时候总是跑几十个iteration就卡死。这个其实是pytorch autograd engine 的问题， 因为每个BN layer的均值和方差都是cross gpu 的grad graph，而我们又是大量使用BN，所以成个back－prop的graph破坏了pytorch grad engine。解决方案是写一个cross gpu的autograd function来handle。

大体思路是这样的，可能发paper的时候再release。不过大部分task应该不涉及这个题。

作者：知乎用户
https://www.zhihu.com/question/59321480/answer/164103885

暂时主流的做法不是 Per-Node / Per-Card 的做么？毕竟这里 Sync 的 Cost 太大了，最后可以 Post BN 的算回。据我所知，CNTK，Torch 和 Tensorflow，Neon 也都类似。我觉得还有部分原因是因为大多现在主流框架都利用 cuDNN 去计算了，相对会更不好处理。

这样也能保证每次 BN 都是 32 BatchSize 的算，跟原 Paper 一样，以我个人有限的实验经验来看，Per-Node 的 BN Batch-Size 更大（总 Batch-Size 不变）结果不小的变化。

当然我主要工作是做验证的。这部分训练的时候，该 Sync 不该 Sync 我觉得还是结果说话的，我觉得值得一试，如果确实结果提升大，那确实要考虑改了，如果不大，毕竟现在也 Work，效率也高得多。

当然对于部分 Task，比如 Action Recognition，这么做确实有特殊的必要，毕竟如果 3D-Conv 的网络，Per-Node 也没多少个 Sample，受限于显存大小。

我觉得非得这么样 Sync 的算，从优化的角度主要也还是节点通信的 Cost，这个又回到了老问题上面，显卡间就是 NCCL 那种策略，或者最简单的 Tree-based 相邻节点的 Sync （先同一个 PCI-E Switch 的，再同一个 CPU Node 的，然后 QPI 的，再 IB 的）

作者：Misquith
https://www.zhihu.com/question/59321480/answer/536088395

Tensorflow的话 tensorpack 可以实现

用tf做图像处理，做FRCNN的时候搜到tensorpack/tensorpack，看到这个example代码支持cross-gpu batchnorm，准备学下这包咋使的，用到自己的代码里。。

再添加一个，本人做医学图像的，tensorpack本身的batchnorm不支持3d图像的5维输入，安利一下这个：tkuanlun350/3DUnet-Tensorflow-Brats18：https://github.com/tkuanlun350/3DUnet-Tensorflow-Brats18，当我正对着tensorpack里的assert ndims in [2, 4]抓狂的时候看到了这个custom的版本，好感动。。

☆ END ☆

如果看到这里，说明你喜欢这篇文章，请转发、点赞。微信搜索「uncle_pn」，欢迎添加小编微信「 mthler」，每日朋友圈更新一篇高质量博文。

↓扫描二维码添加小编↓