谷歌提出AI训练提速新方法，榨干GPU空闲时间，最高提速3倍多

即将开播：5月14日，Jenkins在K8S下的三种部署流程和实战演示

本文经AI新媒体量子位（公众号ID:QbitAI）授权转载，转载请联系出处。

因为通用计算芯片不能满足神经网络运算需求，越来越多的人转而使用GPU和TPU这类专用硬件加速器，加快神经网络训练的速度。

但是，用了更快的GPU和TPU就一定能加速训练吗？

训练流水线的所有操作并不都是在加速器上运行。上游数据处理（如磁盘I/O和数据预处理）就不能在加速器上运行。

随着GPU等加速器越来越快，超过了CPU、磁盘处理数据的速度，上游就逐渐成了训练瓶颈。

在某些情况下，GPU上游的代码花费的时间甚至是GPU本身运行时间的几倍。上游没做完，下游只能空等，浪费了大量时间。

为此，Google AI团队，提出一种简单的数据回波（Data Echoing）方法，可以解决这个问题。该方法最多能将训练速度加速3倍以上，且对模型的精度没有明显影响。

Jeff Dean也在Twitter上转发点赞。

重复数据让GPU不空等

很多情况下，上游花费的时间比加速器更长，使用更快的GPU或TPU根本不会提高训练速度。如果投入大量的工程工作以及额外的计算资源，确实可以加快流水线的速度。

对于非常小的数据集，可以离线预先计算扩增的数据集，并将整个预处理的数据集加载到内存中。

但这种方法不适用于大多数机器学习训练场景，既耗时又分散了改善推理性能的主要目标。

与其等待更多的数据，不如利用已有的数据来使加速器保持忙碌状态。

在加速器空置50%情况下，预处理batch的第一个优化步骤之后，我们可以重复利用该batch再进行一次训练。

如果重复数据与新数据一样有用，训练效率会提高一倍。

实际上，由于重复数据不如新数据有用，因此数据回波提供的加速要小一些，但和加速器处于空闲状态相比，仍然可以提供明显的加速。

通常有几种方法可以在给定的神经网络训练管道中实现数据回波。

Google提出的技术，是将数据复制到训练管道中某个位置的随机缓冲区中，无论在哪个阶段产生瓶颈之后，都可以将缓存数据插入任意位置。

数据回波在样本级别对数据进行混洗，而batch回波则对重复批次的序列进行混洗。另外还可以在数据扩充之前插入缓冲区，以使重复数据的每个副本略有不同，因此不是简单机械重复，而是更接近一个新样本。

加速3倍多，精度无损失

那么数据回波到底多有用呢？

Google在五个神经网络训练管道上尝试了数据回波，涵盖了图像分类、语言建模和目标检测3个不同的任务，并测量了达到特定性能目标所需的新样本数量。

Google发现，数据回波可以用更少的新样本来达到目标性能，这表明重复使用数据对于降低磁盘I/O很有用。在某些情况下，重复数据几乎与新数据一样有用。

在ImageNet的ResNet-50训练任务中，数据回波可以显著提高训练速度，加速了3倍以上。

从云存储中读取一批训练数据所花的时间是使用每一批数据执行训练步骤的6倍。因此数据最多可以重复5次。

我们将回波因子定义为重复每个数据项的次数，对于以上任务，回波因子最大是5。如果重复样本与新样本一样有用，则应该带来6倍加速，而实际上只加速了3.25倍，但也相当可观。

可能有人会担心重复使用数据会损害模型的最终性能，但实验发现，测试的任何工作负载，数据回波都不会降低最终模型的质量。

随着GPU和TPU性能的继续提升，和通用处理器的差距会越来越大，Google期望数据回波和类似策略将成为神经网络培训工具包中越来越重要的一部分。