利用数据回传加速神经网络训练

在过去十年中，神经网络训练速度的大幅提升使得深度学习技术能够应用于许多重要问题。在摩尔定律的暮色中，随着通用处理器的改进趋于平稳，机器学习社区越来越多地转向专用硬件来产生额外的加速。例如，GPU和TPU针对高度可并行的矩阵运算进行了优化，这是神经网络训练算法的核心组件。从高层次上讲，这些加速器可以通过两种方式加速训练。首先，它们可以并行处理更多的训练示例，其次，它们可以更快地处理每个训练示例。我们知道并行处理更多训练示例的加速是有限的，但构建更快的加速器是否会继续加快训练速度？

不幸的是，训练管道中的并非所有操作都在加速器上运行，因此不能仅仅依靠更快的加速器来继续提高训练速度。例如，训练管道中的早期阶段（如磁盘 I/O 和数据预处理）涉及的操作不会从 GPU 和 TPU 中受益。随着加速器改进的速度超过 CPU 和磁盘的改进，这些早期阶段将越来越成为瓶颈，浪费加速器容量并限制训练速度。

代表许多大型计算机视觉程序的示例训练管道。应用小批量随机梯度下降(SGD) 更新之前的阶段通常不会受益于专用硬件加速器。

考虑这样一种情况，加速器上游代码的运行时间是加速器上运行代码的两倍——这种情况对于当今的一些工作负载来说已经很现实了。即使代码被流水线化以并行执行上游和下游阶段，上游阶段也会占据训练时间的主导地位，而加速器将有 50% 的时间处于空闲状态。在这种情况下，构建更快的加速器根本不会提高训练速度。也许可以通过投入工程精力和额外的计算资源来加快输入管道的速度，但这种努力非常耗时，而且会分散提高预测性能的主要目标。对于非常小的数据集，可以离线预计算增强数据集并将整个预处理数据集加载到内存中，但这不适用于大多数 ML 训练场景。

在“使用数据回送加快神经网络训练”中，我们提出了一种简单的技术，可以重用（或“回送”）早期管道阶段的中间输出以回收空闲的加速器容量。我们无需等待更多数据可用，只需利用已有的数据来保持加速器忙碌即可。

左图：没有数据回送时，下游计算能力有 50% 的时间处于闲置状态。右图：回送因子为 2 的数据回送可回收下游计算能力。

重复数据以加快训练速度

想象一下这样一种情况，读取和预处理一批训练数据所需的时间是对该批次执行单个优化步骤的两倍。在这种情况下，在对预处理批次执行第一个优化步骤后，我们可以重复使用该批次并在下一批准备好之前执行第二步。在最佳情况下，重复数据与新鲜数据一样有用，我们会看到训练速度翻倍。实际上，数据回传提供的加速略小，因为重复数据不如新鲜数据有用 - 但与让加速器闲置相比，它仍然可以提供显着的加速。

通常有几种方法可以在给定的神经网络训练管道中实现数据回传。我们提出的技术涉及将数据复制到训练管道中的某个混洗缓冲区中，但我们可以在给定管道中产生瓶颈的任何阶段之后的任何位置插入此缓冲区。当我们在批处理之前插入缓冲区时，我们将我们的技术称为示例回传，而当我们在批处理之后插入它时，我们将我们的技术称为批处理回传。示例回显在示例级别打乱数据，而批次回显则打乱重复批次的序列。我们还可以在数据增强之前插入缓冲区，这样重复数据的每个副本都略有不同（因此更接近新鲜示例）。在将打乱缓冲区放置在不同阶段之间的不同版本的数据回显中，提供最大加速的版本取决于特定的训练管道。

跨工作负载的数据回显那么重用数据有多大用处？我们在涵盖 3 个不同任务（图像分类、语言建模和对象检测

） 的五个神经网络训练管道上尝试了数据回显，并测量了达到特定性能目标所需的新鲜示例数量。我们选择的目标以匹配超参数调整期间基线可靠实现的最佳结果。我们发现数据回显使我们能够用更少的新鲜示例达到目标性能，这表明重用数据对于减少各种任务的磁盘 I/O 很有用。在某些情况下，重复数据几乎与新鲜数据一样有用：在下图中，增强之前的示例回显几乎通过重复因子减少了所需的新鲜示例的数量。

数据回传（即每个数据项重复两次）会减少或不改变达到样本外目标表现所需的新样本数量。虚线表示如果重复样本与新样本一样有用，我们预期的值。

缩短训练时间

每当加速器上游的计算占据训练时间时，数据回传都可以加快训练速度。我们测量了训练管道中实现的训练加速，该管道因来自云存储的流式传输训练数据而受到输入延迟的瓶颈，这对于当今许多大规模生产工作负载或任何通过网络从远程存储系统流式传输训练数据的人来说都是现实的。我们在ImageNet数据集上训练了一个ResNet-50模型，发现数据回传提供了显著的训练加速，在这种情况下，使用数据回传时速度提高了 3 倍以上。

数据回传可以减少 ResNet-50 在 ImageNet 上的训练时间。在此实验中，从云存储中读取一批训练数据所花的时间是使用每批数据执行训练步骤的代码的 6 倍。图例中的回传因子是指每个数据项重复的次数。虚线表示如果重复的示例与新示例一样有用并且没有回传开销，则预期值。

数据回传保留了预测性能

虽然人们可能会担心重复使用数据会损害模型的最终性能，但我们发现数据回传不会降低我们测试的任何工作负载的最终模型的质量。

比较在有和没有数据回传的情况下训练期间取得最佳样本外表现的单个试验，表明重复使用数据不会损害最终模型质量。此处验证交叉熵相当于对数困惑度。

随着 GPU 和 TPU 等专用加速器的改进继续超越通用处理器，我们预计数据回传和类似策略将成为神经网络训练工具包中越来越重要的部分。

致谢

数据回传项目由 Dami Choi、Alexandre Passos、Christopher J. Shallue 和 George E. Dahl 开展，当时 Dami Choi 是 Google AI 驻场研究员。我们还要感谢 Roy Frostig、Luke Metz、Yiding Jiang 和 Ting Chen 的有益讨论。