谷歌开源GPipe，高效训练大型深度神经网络

谷歌的人工智能研究部门今天开源GPipe，这是一个能“有效”训练深度神经网络（在神经元之后建模的分层函数）的Lingvo程序库，Lingvo是用于序列建模的TensorFlow框架。它适用于由多个连续层组成的任何网络。据谷歌人工智能软件工程师黄燕平在博客文章中介绍，GPipe使研究人员能够“轻松”扩展性能。

“深度神经网络（DNN）推动了许多机器学习任务，包括语音识别，视觉识别和语言处理。更大的DNN模型带来更好的任务绩效，而视觉识别任务的进展也佐证了模型大小和分类准确性之间的强相关性，“黄燕平补充说：“通过GPipe …我们展示了使用管道并行性扩展DNN训练，可以克服模型大小的限制。”

正如Huang及其同事在一篇随附的论文（“ GPipe：使用管道并行性的高效神经网络的高效训练 ”）中所解释的那样，GPipe实现了两种出色的AI训练技术。一种是同步随机梯度下降，用于更新给定AI模型参数的优化算法；另一种是流水线并行——一种任务执行系统，其中一步的输出作为下一步的输入。

大多数GPipe的性能提升得益于更好的AI模型内存分配。在第二代谷歌云tensor处理器（TPU）上，每个处理器包含8个处理器核心和64 GB内存（每个核心8 GB），GPipe将中间内存使用量从6.26 GB降低到3.46GB，在单个加速器核心上实现了3.18亿个参数。黄燕平指出：如果没有GPipe，单个核心只能训练至多8200万个模型参数。

这不是GPipe的唯一优势。它将模型划分为不同的加速器，并自动将微型批次（即“迷你批次”）的训练样本分成较小的“微批次”，并在管道中执行。这使得核心能够并行操作，并且还在微批次上累积梯度，从而防止分区影响模型质量。

在一项实验中，Google训练了一种深度学习算法 – AmoebaNet-B–在TPU上有5.57亿个模型参数和样本图像，每个TPU上包含18亿个参数（比没有GPipe的情况下多25倍）。黄燕平说，它在流行的数据集上表现“很好”，将单一对象的ImageNet精度推高到84.3％，CIFAR-10精度达到99％，CIFAR-100精度达到91.3％。

训练速度也有所提高。在采用AmoebaNet-D算法的单独测试中，将模型分配给四倍数量的第二代TPU核心，实现了3.5倍的加速。当谷歌研究人员采用第三代TPU核心（最新版本）上测试具有80亿个参数的Transformer语言模型时 ，每个核心都有16个核心和256GB内存（每个核心16GB），它们的速度提高了11倍。

“许多机器学习应用的持续发展和成功，例如自动驾驶和医学成像，取决于实现最高精度，”黄燕平写道：“由于这通常需要构建更大，更复杂的模型，我们很乐意为更广泛的研究社区提供GPipe，并希望它是有效培训大规模DNN的有用基础设施。”