今年 6 月,Facebook 实现 1 小时训练 ImageNet 的成绩之后,通过增加批量大小以加快并行训练速度的方式引发了人们的关注。最近 UC Berkeley 的研究人员们为我们展示了 24 分钟训练 ImageNet 的成绩,他们将批量大小增加到了 32k。研究人员表示,在同样成绩下,新的方法使用的计算设备成本(120 万美元)大大低于 Facebook(410 万美元)。
对于深度学习应用而言,模型、数据集越大,结果就越精确,但同时也意味着训练时间的增长。例如,用 ResNet-50 模型在 ImageNet-1k 上完成 90 个 epoch 需要 14 天时间,这样的训练共需要 10^18 次单精度计算(single precision operations)。另一方面,世界目前最为强大的超级计算机可以在每秒钟进行 2×10^17 次单精度计算。如果我们可以完全利用超级计算机的能力,那么使用 ResNet-50 模型在 ImageNet-1k 上应该可以在五秒内完成 90 个 epoch 的训练。然而,目前我们面临的 DNN 训练瓶颈在算法层面上。具体来说,目前我们使用的批量过小(如 512),这让处理器并行处理的效率很低。所以,我们需要为可扩展的 DNN 训练设计有效的并行算法。
使用大批量SGD是大规模 DNN 训练的基础。但是,此前的方法只能将 AlexNet 和 ResNet-50 的批量大小扩大到 1024 和 8192 这样的程度。更大的批量大小会导致结果准确度的显著降低。在 UC Berkeley 最近发表的论文中,研究人员使用层级对应的适应率缩放(Layer-wise Adaptive Rate Scaling,LARS)算法,让批量大小可以扩大到更大的级别(如 32k)而不损失结果准确度。
大批量可以让我们更加有效地利用分布式系统的计算能力。UC Berkeley 的研究人员使用 AlexNet 模型在 ImageNet 上训练 100 个 epoch 只花费了 24 分钟,这是一个新的世界纪录。和 Facebook 此前的结果类似(参见:Facebook 新研究:大批量 SGD 准确训练 ImageNet 仅需 1 小时)新的方法也可以在一小时内使用 ResNet-50 模型在 ImageNet 上训练 90 个 epoch,但硬件花费约为 120 万美元——远少于 Facebook 今年 6 月时使用的 410 万美元。
图 1.(a)是数据并行化的一个示例。每个工作器将自己的梯度∇w^j 发送给控制器,随后控制器更新自己的权重。随后,控制器将更新后的权重 w˜发送给所有工作器。(b)是模型平行化的示例。其中包含一个五层神经网络和本地连接,它被分割成四个部分。只有在跨区域边缘位置的节点需要在连接处(粗线)共享它们的状态。即使节点位置与多条边相连,它也只会对其他部分的节点总共发送一次自己的状态。
表 1. 在 ImageNet 数据集上训练神经网络,tcomp 是计算时间,tcomm 是通信时间。研究人员将 epoch 数量设为 100。越大的批量大小需要越小的迭代次数。研究人员首先将批量大小设为 512,随后不断增加计算设备数量,由于训练 ResNet-50 的 Infiniband 网络和 GPU 效率足够高,单次迭代时间可能接近常数,总训练时间会不断减少。
图 2. 在一定范围内,大批量提升了系统(GPU)的处理性能。本图中的数据是 AlexNet 在 ImageNet 数据集上的训练效果,使用的 GPU 是英伟达 M40,每个 GPU 的批量大小为 512——这是我们能得到的最高速度。Batch=1024 会导致内存溢出。
此前,英伟达曾经表示使用 DGX-1 工作站在 ImageNet 中训练 AlexNet 90 个 epoch 需要花费两个小时。然而,他们使用了半精度(half-precision)或 FP16,计算资源消耗是标准单精度计算的一半。在本研究中,UC Berkeley 的研究者们使用了标准单精度计算,在英伟达 DGX-1 工作站中,批量大小为 512 时花费了 6 小时 9 分钟。因为 LARS 算法的存在,我们可以在保证精度的情况下大幅增加批量大小。如果我们将批量大小增加至 4096,在英伟达 DGX-1 上我们就只需要 2 小时 10 分钟即可完成任务了(这也是单机最快的成绩)。所以,通过 LARS 增加批量大小的方法可以显著提高 DNN 的训练速度。在 AlexNet 上,当批量大小为 32k 时,研究人员将算法扩展到 512 个 KNL 芯片(约 32,000 个处理器核心)。每个 KNL 的批量大小为 64,所以总批量大小为 32678。最终,UC Berkeley 实现了 24 分钟训练 100 个 epoch 的成绩,这也是目前业内最佳的成绩。
表 7. 训练 AlexNet 的速度和硬件花费,批量大小为 32k,研究人员已将 AlexNet 中的局部响应范数(local response norm)改为批范数。
在今年 6 月份,Facebook 在 ImageNet 上训练 ResNet-50 时(90 epoch)花费一小时,使用了 32 个 CPU 和 256 个英伟达 P100 GPU(32 个 DGX-1 工作站),目前单个 DGX-1 的价格为 129,000 美元,整个硬件系统的成本为 410 万美元。
在将批量大小扩展到 32k 之后,UC Berkeley 的研究人员们将计算任务交给了相对便宜的计算机芯片——KNL 上,同样在一个小时的时间里跑完了 90 个 epoch。他们使用的 KNL 搭载英特尔 Xeon Phi 7250 处理器,单个售价 2436 美元,整个系统花费 120 万美元。
表 8. 训练 ResNet-50 的速度和硬件成本,研究中未使用数据增强。
论文:ImageNet Training in 24 Minutes
论文链接:https://www2.eecs.berkeley.edu/Pubs/TechRpts/2017/EECS-2017-155.html
摘要:想要在 ImageNet-1k 上完成 90 个 epoch 的 ResNet-50 模型训练,使用英伟达 M40GPU 需要花上 14 天。这种训练共需要 10^18 次单精度计算(single precision operations)。另一方面,世界目前最快的超级计算机可以每秒钟进行 2×10^17 次单精度计算。如果我们可以充分利用超级计算机的性能用于深度神经网络讯联,我们应该可以使用 ResNet-50 模型在 5 秒钟内完成 90 个 epoch 训练。然而,目前的 DNN 训练瓶颈在算法层面上。具体来说,目前我们使用的批量大小(如 512)过于小,无法有效利用多处理器。
对于大规模 DNN 训练,我们专注于使用大批量数据并行化同步(large-batch data-parallelism synchronous)SGD 在固定 epoch 中保证不损失精度。You、Gitman 与 Ginsburg 在 2017 年提出的 LARS 算法允许我们将批量大小扩大到 32k 的水平。在此之上,我们只花费了 24 分钟即完成了在 ImageNet 上,100 个 epoch 的 AlexNet 训练——这是一个世界纪录。与 Facebook 的结果相同,我们在 ImageNet 上训练 ResNet-50 花费了一个小时,而使用的硬件成本仅为 120 万美元,比 Facebook 使用的 410 万美元降低了很多。