谷歌 人工智能部门 https://ai.google/ 最近开源了 GPipe , 这是一个用于快速训练大规模深度学习模型的 TensorFlow https://www.tensorflow.org/ 类库.
深层神经网络 (DNN) 主要用于解决自然语言处理和视觉目标识别等人工智能任务. 以视觉识别为例, 该领域的最新方法通常以 ImageNet 挑战赛 http://image-net.org/ 的获胜方案为基准. 每一届冠军的成绩都优于前一届; 当然, 模型的复杂度也会相应增加. 2014 年的冠军 GoogLeNet 通过使用 400 万个模型参数达到了 74.8% 的 top-1 准确率, 而 2017 年的冠军 Squeeze-and-Excitation Networks 则使用了 1.458 亿个参数并达到了 82.7% 的 top-1 准确率.
在训练神经网络的时候, 模型大小的增加通常会引起问题. 为了在合理的时间内完成训练, 我们把大部分的计算任务委托给了加速器: 诸如 GPU 和 TPU https://cloud.google.com/tpu/ 之类的专用硬件. 但是这些设备的内存有限, 这也就限制了训练模型的大小. 我们可以通过一些方法来减少模型对内存的依赖, 比如将数据从加速器内存中置换出去, 但这会大大减慢训练速度. 另一种解决方案则是模型分区, 这可以让模型同时在多个加速器中并行执行. 对顺序性 DNN 来说, 最好的策略是按层划分模型, 然后由不同的加速器来训练不同的层. 但是由于 DNN 的顺序性本质, 有些时候可能只有一个加速器在工作, 别的加速器则因为需要等待其它层的训练结果而闲置下来.
GPipe 通过进一步细化训练任务解决了这个问题, 它将批量任务分解为更细小的 "微批量" 任务, 并在每一层中管道化执行这些 "微批量" 任务. 这样, 下一层的加速器就可以优先处理上一层已完成的 "微批量" 任务结果, 而不需要等待整个训练过程的结束.
通过使用 GPipe 以及 8 个 TPUv2(第二代 TPU 芯片), 谷歌研究人员能够用 18 亿个参数来训练视觉目标识别模型: 在使用 GPipe 的情况下, 单个 TPUv2 可训练的参数量增加了 5.6 倍. 通过此次训练的大规模模型, ImageNet 数据验证的准确率达到了 84.7%, 超过了 2017 年夺冠时的 82.7%.
GPipe 的模型分区除了能支持更大的模型以外, 它也允许多个加速器并行训练所指定的模型. 研究报告称, 使用 4 倍以上的加速器可以达到 3.5 倍的加速效果.
Gpipe 目前是 Lingvo 框架 https://github.com/tensorflow/lingvo 的一部分, 该框架主要用来在 TensorFlow 中构建顺序神经网络模型.
来源: http://www.tuicool.com/articles/3q2eamB