谷歌开源 GPipe 库，速度提升 25 倍，CIFAR-10 精度达到 99%

作者：Yanping Huang
来源：机器之心
原文：https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html

BigGan、Bert、GPT 2.0 等模型的出现似乎表明，DNN 模型越大，其准确率越高。但芯片内存的增长速度似乎无法满足这种“暴力美学”的需要。为此，谷歌提出并开源了一种新的方法——GPipe，利用管道并行化扩展 DNN 训练以克服这种局限。它可以让研究者轻松部署更多的加速器，以训练更大的模型，并在不调整超参数的前提下实现性能扩展。

深度神经网络（DNN）推动了许多机器学习任务的发展，包括语音识别、视觉识别、语言处理。BigGan、Bert、GPT 2.0取得的近期进展表明，DNN模型越大，其在任务中的表现越好。视觉识别领域过去取得的进展也表明，模型大小和分类准确率之间存在很强的关联。例如，2014年ImageNet视觉识别挑战赛的冠军GoogleNet以400万的参数取得了74.8%的top-1准确率，但仅仅过了三年，冠军的宝座就被Squeeze-and-ExcitationNetworks抢去，后者以1.458亿（前者的36倍还多）的参数量取得了82.7%的top-1准确率。然而，在这段时间里，GPU的内存只提高了3倍左右，当前最优的图像模型却已经达到了谷歌云 TPUv2的可用内存。因此，我们急需一个能够实现大规模深度学习并克服当前加速器内存局限的可扩展高效架构。

近期有代表性的图像分类模型的ImageNet准确率和模型大小之间的强关联。

谷歌在论文《GPipe: Efficient Training of GiantNeural Networks using Pipeline Parallelism》中，展示了利用管道并行化（pipeline parallelism）扩展DNN训练以克服这种局限。GPipe是一个分布式机器学习库，它使用同步随机梯度下降和管道并行化进行训练，可以应用到包含多个序列层的任意DNN中。重要的是，GPipe可以让研究者轻松部署更多的加速器，以训练更大的模型，并在不调整超参数的前提下实现性能扩展。为了展示GPipe的有效性，谷歌在谷歌云TPUv2上训练了AmoebaNet-B模型，该模型具备5.57亿个参数，输入图像大小为480 x 480。这一模型在多个流行的数据集上表现良好，比如将ImageNet上的single-crop准确率提高到84.3%，将CIFAR-10上的准确率提高到99%，CIFAR-100准确率提高到91.3%。该核心GPipe库已经在Lingvo框架下开源。

开源地址：https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py

从小批量到微小批量

加速中型DNN模型有两种标准方法。第一种是数据并行化方法，它利用更多机器将输入数据分散到这些机器上。另一种方法是将模型移到加速器上，如GPU 或 TPU，这些加速器具备特殊的硬件可以加速模型训练。但是，加速器内存有限、通信带宽也有限。因此，要在加速器上训练大型DNN模型需要模型并行化，将模型分割成多个部分，然后将不同的部分分配到不同的加速器。但是由于DNN的序列本质，这一策略可能导致计算期间只有一个加速器处于活跃状态，而这会导致无法充分利用加速器的计算能力。另外，标准数据并行化方法允许同一个模型在多个加速器上对不同的输入数据执行并行训练，但是这无法增加每个加速器可以支持的最大模型大小。

为了在多个加速器上也能进行高效的模型训练，GPipe将模型分割并分配给不同的加速器，将小批量训练样本自动分割成更小的批量（微小批量）。通过在微小批量样本上管道化整个执行过程，加速器可以实现并行运行。此外，各微小批量上的梯度可以一直累加，这样分区数量就不会影响到模型质量。

上：由于DNN网络的序列本质，初始的模型并行化策略导致计算能力无法充分利用，每次只有一个加速器处于活跃状态。下：GPipe 将输入小批量分割成更小的批量，使得不同的加速器可以同时处理各自分配到的微小批量样本。

最大化内存和效率

GPipe 最大化了模型参数的内存分配。谷歌在云 TPUv2上进行试验，每个芯片上有8个加速器内核和64GB的内存（每个加速器8GB内存）。没有TPUv2的情况下，由于内存限制，单个加速器最多训练8200万模型参数。由于反向传播和批量分割中的重复计算，GPipe 将中间激活内存从6.26GB降至3.46GB，使得单个加速器上可以训练3.18亿个参数。谷歌还发现，如预期所料，在管道并行化的情况下，最大模型大小与分区数量成正比。有了GPipe，AmoebaNet 能够在云TPUv2的8个加速器上加入18亿参数，是没有GPipe时的25倍。

为了测试效率，谷歌衡量了GPipe对AmoebaNet-D模型吞吐量的影响。因为训练至少需要两个加速器来适应模型大小，谷歌衡量了在两个分区但没有管道并行化的naive情况下的加速，发现训练过程中几乎是线性加速。与具有两个分区的naive方法相比，将模型分布在四倍的加速器上实现了3.5倍的加速。虽然本文所有的实验都使用了云TPUv2，但谷歌发现当前可用的云TPUv3会有更好的性能，每个芯片有16个加速器内核和256GB内存（每个加速器16GB）。当在所有16个加速器上分布模型时，GPipe使得具有80亿参数的Transformer 语言模型在1024-token语句上的加速达到11倍。

使用GPipe加速AmoebaNet-D。该模型不适合加速器。基线naive-2是模型被分割成两个分区时native分区方法的性能。Pipeline-k是使用K个加速器将模型分割成K个分区时GPipe的性能。

GPipe还可以通过使用更多的加速器而不改变参数来扩大训练规模。因此，它可以与数据并行相结合，以互补的方式使用更多的加速器来扩大神经网络训练规模。

测试准确率

谷歌使用GPipe 来验证一个假设，即扩大现有神经网络能够提高模型质量。谷歌在 ImageNet ILSVRC-2012 数据集上训练了一个AmoebaNet-B模型，其模型参数为5.57亿，输入图像大小为480 x 480。该网络被分为四个分区，且在模型和数据上应用了并行训练过程。这个巨大的模型在没有任何外部数据的情况下达到了当前最先进的84.3% top-1 / 97% top-5 single-crop验证准确率。大型神经网络不仅适用于如ImageNet这样的数据集，并且通过迁移学习还可适用于其他数据集。已有研究证明，更好的ImageNet模型迁移效果更好。谷歌在CIFAR10 和CIFAR100数据集上进行迁移学习实验。其巨大的模型将 CIFAR-10的准确率提高到99%，将CIFAR-100的准确率提高到91.3%。

表2：单个AmoebaNet-B(6, 512) 模型在ImageNet ILSVRC 2012验证数据集上的single-crop分类准确率与其他已发布最优模型的对比结果。

表3：AmoebaNet-B (6,512) 的迁移学习结果，输入图像大小为480×480，测试时使用的是single crop。这些结果均为5次精调运行后的平均值。

很多机器学习实际应用的持续发展和成功依赖于尽可能高的准确率，如自动驾驶和医疗成像。而这通常需要构建更大型、更复杂的模型，谷歌很高兴将GPipe带给更广大的社区，并希望它能够有效训练大型深度神经网络。

论文：GPipe: Efficient Training of GiantNeural Networks using Pipeline Parallelism

论文链接：https://arxiv.org/pdf/1811.06965.pdf

摘要：GPipe是一个可扩展管道并行化库，可以训练超大型的深度神经网络。它将网络的层分割成多个部分并分配给不同的加速器，将执行过程管道化以达到最高的硬件利用率。例如，在8个加速器上应用分割技术，则GPipe支持25倍大小的神经网络的训练，这充分展示了其可扩展性。它还保证了无论分区数量多少，梯度仍是一致的。在对模型参数几乎不做更改的情况下，它获得了几乎是线性的加速：加速器数量是原来的4倍时，同一模型的训练速度是原来的3.5倍。我们训练了一个具备5.57亿参数的AmoebaNet模型，在 ImageNet 2012 数据集上达到了新的SOTA结果：84.3%的 top-1 准确率和97.0% 的 top-5 准确率。最后，我们将这个学到的模型应用到多个流行图像分类数据集上，同样获得了有竞争力的结果，比如在CIFAR-10数据集上达到了99%的准确率，在CIFAR-100数据集上达到了91.3%的准确率。

原文链接：https://ai.googleblog.com/2019/03/introducing-gpipe-open-source-library.html

推荐阅读
Google：数据并行对神经网络训练用时的影响
微软开源的深度学习模型转换工具 MMdnn

微信公众号: 极市平台（ID: extrememart ）
每天推送最新CV干货