通过数据回送提升神经网络训练速度

文 / 学生研究员 Dami Choi 和高级研究员 George Dahl

Google Research

在过去的十年中，随着神经网络的训练速度显著提升，我们可以应用深度学习技术来解决越来越多的重要问题。在摩尔定律 (https://www.blogger.com/u/2/wiki) 日渐失色之际，随着通用处理器的改进陷入瓶颈，为提升速度，机器学习社区越来越推崇专用硬件。例如，GPU 和 TPU 可优化高度并行的矩阵运算，而这些运算是神经网络训练算法的核心组成部分。这些加速器可通过两种方式显著提升训练速度。首先，可以并行地处理更多训练样本；其次，还可以更快速地处理每个训练样本。我们知道并行处理更多的训练样本会使提速受到限制，那提升加速器的速度会不断提升训练速度吗？

• 我们知道
  https://ai.googleblog.com/2019/03/measuring-limits-of-data-parallel.html

遗憾的是，并非所有训练流水线中的操作都在加速器上运行，因此，仅依赖提升加速器速度，无法达到不断提升训练速度的目的。例如，训练流水线中的前期操作（如磁盘 I/O 和数据预处理）就不能受益于 GPU 和 TPU。随着加速器的速度提升超过 CPU 和磁盘，这些前期操作将逐渐成为瓶颈，不但会浪费加速器的计算能力，还会限制训练速度。

多个大规模计算机视觉程序的典型训练流水线示例：小批次随机梯度下降 (SGD) 更新之前的阶段通常无法受益于专用硬件加速器

假定加速器上游的代码花费的时间是加速器上所运行代码的两倍，如今这种情况已经是某些工作负载的实际问题。即使将代码流水线处理成能够并行执行上游和下游的操作，上游操作也会占据训练时间，而加速器将有 50% 的时间处于闲置状态。在这种情况下，提升加速器的速度根本无法提升训练速度。专注于工程工作和增加计算资源或许可以提升输入流水线的速度，但这些工作不仅耗时，还会偏离提高预测能力的主要目标。对于规模非常小的数据集，可以离线预计算扩增的数据集，然后将整个预处理的数据集加载到内存中，但此方法并不适用于大多数机器学习训练场景。

在《通过数据回送提升神经网络训练速度》(Faster Neural Network Training with Data Echoing) 一文中，我们提出了一项简单技术，即重用（或“回送”）流水线前期的中间输出，以利用加速器闲置的计算能力。我们不必等待更多的数据，只需利用已有的数据使加速器保持忙碌状态。

左：未使用数据回送，下游计算能力闲置率达 50%；右：使用回送因子为 2 的数据回送，利用下游计算能力

• 通过数据回送提升神经网络训练速度
  https://arxiv.org/abs/1907.05550

重复数据，加快训练

假定读取并预处理一批训练数据的时间是对其执行一次优化步骤所需时间的两倍。在这种情况下，对这批预处理数据执行首次优化步骤后，我们可以在下一批数据准备就绪之前重用这批数据再次执行一次优化步骤。在最佳情况下，如果重复的数据与新数据一样有用，那么训练速度会提高一倍。实际上，由于重复的数据不如新数据有用，数据回送的提速影响力略小一些，但与加速器闲置相比，此技术仍可显著提升训练速度。

在给定的神经网络训练流水线中，通常有几种方法可实现数据回送。我们提出的技术是将数据复制到训练流水线中某个位置的随机缓冲区，无论在给定流水线的哪个阶段遇到瓶颈，都可以将此缓存区插入任意位置。在批处理前插入缓存区时，我们可以调用 样本回送 技术，而在批处理后插入缓存数据时，我们可以调用 批回送 技术。样本回送是在样本级别随机打乱数据，而批回送是随机打乱重复批次的序列。我们还可以在数据增强前插入缓存数据，以使重复数据的每个副本略有不同（因此更接近新样本）。不同版本的数据回送会将随机打乱的缓存数据置于不同阶段之间，哪个版本的提速能力最高要取决于具体的训练流水线。

各工作负载中的数据回送

重用数据有多有用？我们在 5 个神经网络训练流水线中尝试使用数据回送，其中涵盖 3 项不同的任务（图像分类、语言建模和对象检测），并测量了达到特定性能目标所需的新样本数量。我们选择了几个目标，旨在拟合基准训练在超参数调整时所达到的最佳效果。我们发现，数据回送可让我们使用更少的新样本来达到目标性能，这表明在各种任务中重用数据能够有效减少磁盘 I/O。在某些情况下，重复的数据几乎与新数据一样有用：在下图中，数据增强前的样本回送减少了重复因子所需的新样本数量。

当每个数据项重复两次时，数据回送会减少或不更改达到目标样本性能所需的新样本数量。虚线表示重复样本与新样本一样有用时的预期值

缩短训练时间

在加速器上游的计算影响训练时间时，数据回送可以提升训练速度。我们测量了特定训练流水线的训练提速情况，此流水线因采用串流技术传送云存储空间中的训练数据产生输入延迟瓶颈，这种情况是当今许多大规模生产工作负载，或采用串流技术传送远程存储系统中的网络训练数据的人员都需要面对的实际问题。我们在 ImageNet 数据集上训练了一个 ResNet-50 模型，发现数据回送可显著提升训练速度，在本例中，使用数据回送后，训练提速 3 倍以上。

数据回送可缩短 ImageNet 上 ResNet-50 的训练时间。在此实验中，与使用每一批数据执行训练步骤的代码相比，读取云存储空间中一批训练数据所耗费的时间是它的 6 倍。图例中的回送因子是指重复每个数据项的次数。虚线表示重复样本与新样本一样有用且没有回送开销时的预期值

数据回送不会影响预测能力

有人可能会担心重用数据有损模型的最终性能，但我们发现，数据回送在我们已测试的所有工作负载中均未降低最终模型的质量。

比较训练时使用和不使用数据回送达到最佳样本性能的个别试验后，结果表明重用数据不会影响最终的模型质量。这里的验证交叉熵相当于 log perplexity

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

致谢

参与数据回送项目的人员有 Dami Choi、Alexandre Passos、Christopher J. Shallue 和 George E. Dahl，其中 Dami Choi 是 Google AI 定向实习生。同时感谢 Roy Frostig、Luke Metz、Yiding Jiang 和 Ting Chen 提供实用的讨论意见。