单GPU用CIFAR10训练达到94%只要26秒的秘密

作者 | David Page

译者 | 核子可乐

编辑 | 陈思

AI 前线导读：本文将进行全面探讨，包括如何加快单 GPU 训练实现，最终在其它多 GPU 竞争方案中脱颖而出。本文提出了一系列标准以及不那么标准的实施技巧，其能够将训练时长缩短至 34 秒；如果配合测试时间增强，则可进一步缩短至 26 秒。

更多优质内容请关注微信公众号“AI 前线”（ID：ai-front）

此前，我们在超参数调优、权重衰减以及批量规范等指标上的成绩已经被其他人超越人，因此我们 94% 的 CIFAR 训练准确率已经在 DAWNBnch 排行榜上下滑了五位：

上榜的六大系统都使用 9 层 ResNets，可以说各套解决方案之间算是堂兄弟或者多胞胎的关系。第一位是来自 Kakao Brain 的 4-GPU 实现方案，其仅用 37 秒就完成了任务。同一套方案的单 GPU 版本位列第三，耗时为 68 秒，比去年的单 GPU 方案缩短了 7 秒。不过认真观察，我们发现新方案中采用了测试时间增强（TTA）技术。在文章末尾，我们将讨论这种方法的作用（这里先来说一点结论，我们认为任何合理的限制都应该与总体推理成本相符，因此在这里稍微使用 TTA 以及轻量级网络并不是坏事。）需要注意的是，我们之前的提交成果也允许使用相同的 TTA，其能够在不变更任何其它元素的基础上，直接将成绩提升为 60 秒（19 轮训练周期）。

到本文末尾，我们的单 GPU 实现方案将成功超越顶级多 GPU 方案，以 34 秒的总时长重夺 DAWNBench 榜首地位，并在单 GPU 层面实现高达 10 倍的性能改进。如果配合 Kakao Brain 所使用的 TTA，那么时长将进一步缩短至 26 秒。我们通过将一系列小小的改进性结合起来（通常能够为测试贡献 0.1% 到 0.3% 的绝对准确率），最终实现了可观的性能提升，即以更短的训练周期完成执行。而这一切，基于一系列标准以及不那么标准的小技巧。

我们的主要武器就是统计意义。音效训练的测试准确率标准差约为 0.15%，两次运行蝗比较需要乘以√2。这要远高于我们测量到的多数结果。考虑到训练时间很快就会下降至一分钟以内，我们可以通过 10 到 100 次实验以确保改进的真实性。也只有这样，我们才能持续取得进展。

明确的实验结果对于推进项目非常重要，但如果基准调整不佳或者运行次数太少，那么我们的验证结果也将失去价值。本次发布的一大主要目标，就是提供一套经过良好调整的基准用于测试新技术。其能够帮助用户在几分钟之内在 GPU 上完成拥有显著统计学意义的训练。我们在文章末尾将确认，如果继续进行训练，那么训练速度的提高将最终转化为准确率的提升。

GPU 上的预处理（70 秒）

我们先从代码的实际优化问题开始。我们之前提交至 DAWNBench 的日志显示，数据的预处理中存在 3 秒钟的浪费，这无疑延长了整个训练时间。有鉴于此，我们在训练之前对数据集进行了规范化、转置与填充，以避免在各个阶段中进行重复工作。

我们可以将数据传输至 GPU，并在那里进行预处理，然后再传输回 CPU 进行随机数据扩充与批处理——这能够进一步提升训练效率。将整体数据集（uint8 格式）迁移至 GPU 只需要可以忽略不计的 40 毫秒，而在 GPU 上完成预处理步骤的速度则更快，只需要 15 毫秒。但是，大部分时间都被耗费在了将预处理后的数据集传回 CPU 身上，这需要将近半秒钟。虽然比我们之前的 3 秒钟好了太多，但其中应该仍有改进空间，因为数据需要在批处理与增强之后再次交付至 GPU，这意味着每个训练步骤当中都会积累这一延迟。那么，我们可以把数据扩充转移到 GPU 上，从而回避这部分时间浪费吗？

答案是肯定的，但处理的过程却需要小心谨慎。如果我们简单将增强应用于单一训练示例（就像在 CPU 上那样），则会产生大量开销——系统会启动多个 GPU 核心处理各个条目。为了避免这种情况，我们可以对示例组应用同样的扩充方式，同时配合对数据进行预先混洗以保持随机性。

例如，在对 CIFAR10 图像进行 8 x 8 剪切增强时：由于 32 x 32 图像当中共存在 625 个可能的 8 x 8 剪切区域，因此我们可以通过混洗数据集并将其分为 625 组的方式实现随机增强，每个可能的剪切区域一个。如果我们选择的组拥有完全均匀的大小，那么即使与完全随机的示例选择仍不相同（其中的组大小并不规则），但二者的结果也已经相当接近。要做出进一步优化，如果扩充组的数量足够大，我们可以考虑对其做出合理的限制——例如在每个时间点上选取 200 个随机组。

我们的基本实现非常简单，只需要大概 35 行代码（不需要任何 Pytorch DataLoaders）。以下是对相同 4 幅图像的两次随机扩充，用以显示系统的实际效果：

更重要的是，它的速度很快，能够在 400 毫秒以内完成 24 轮数据训练周期迭代，期间应用随机裁剪、水平翻转与剪切数据增强、改组与批处理等。整个周期仍比将数据集传输至 CPU 一次所耗费的时间少！此外，由于我们不再针对 GPU 争用 CPU 预处理队列，因此完全不必担心数据加载问题——即使训练速度更快，一切也仍将井然有序。

注意：我们这套方案的前提在于数据集足够小，完全能够以整体形式在 GPU 内存中进行存储与操作；对于更复杂的场景，可能需要更完善的实现方案或者采用英伟达 DALI 这类工业级解决方案。

如果我们使用新的 GPU 数据处理机制处理之前 DAWNBench 提交中的网络与训练流程，那么训练时间能够下降至 70 秒以内，这将使我们在排行榜上晋升两位！

旁白：混合精度训练

在我们最初提交至 DAWNBench 的测试代码当中，我们简单地将模型转换为 float16，而没有对全部细节进行混合精度训练。当然，我们还是通过加和方式强调了“损失伸缩”，这种作法至少要比批量取平均损失值好得多。虽然采用适当的混合精度训练并不困难人，但这会增加约 1 秒的整体训练时长，而且我们发现其对最终准确率并没有影响，所以这一次我们继续选择忽略。

移动最大池层（64 秒）

最大池化与 ReLU 等整体递增激活函数进行通信，因此先应用池化应该能够提高效率。我们可以利用合适的编译器完成这项工作，但这里我们先手动进行切换。以下是池化之前的典型 conv-pool 块：

池化之后：

顺序的切换让 24 轮训练时间进一步减少了 3 秒，而网络计算出的特征则完全没有变化！也许我们可以尝试更激进的方法，即在批量规范化之前移动最大池。这将进一步提高效率，但会对网络本身造成影响，因此我们需要测试这究竟会如何给训练带来怎样的变化。

结果是，与之前 94.1% 的基准相比，我们的测试准确率只受到了一丁点影响，目前为 94.0%（50 次运行取平均值）。更重要的是，这让训练时间缩短了 5 秒。我们可以通过在训练中添加额外的周期来恢复之前的准确性。这也是本篇文章中，我们使用的唯一一次会导致准确率下降的“改进”！这 5 秒优化来自更高效的网络体系，而且高于引入新的训练周期所带来的 2.5 秒损失。到这里，净收益让我们的成绩提升到了 64 秒，位列榜单第三位。

标签平滑（59 秒）

标签平滑在分类改进问题当中的一种相当成熟的神经网络训练速度与泛化技巧。其实质，是将独热目标概率与交叉熵损失内的各个类标签进行均匀混合。这有助于稳定输出分布结果，并防止网络进行过度自信的预测（可能抵制后续训练效果）。我们首先尝试将标签平滑参数设定为 0.2，这是一种比较粗略的手工优化方式，但结果显示具体参数选择并不会造成太大的影响。

标签平滑损失

这一步让我们的准确率提升至 94.2%（50 次运行取平均值）。我们可以通过减少训练周期的次数换取训练速度的加快。根据经验，准确率每提升 0.1%，就意味着我们可以减少一个训练周期。这大体反映出额外训练周期的增益效果。我们借此降低了预热期（在此期间内学习率呈线性增加）与总周期数之间的比例。

经过 23 轮训练，我们的训练准确率为 94.1%，训练时长也终于降低到了 1 分钟以内！

CELU 激活（52 秒）

我们希望能够利用平滑激活函数对流程进行优化，借此取代初始阶段的 ReLU 及其具有曲率的 delta 函数。这种作法也能够提升模型的推广能力，因为平滑函数是一种表达式较少的函数类——在整体平滑限制之下，我们能够借此还原出一套线性网络。

我们对 ReLU 的其它表现非常满意，因此我们决定选择一种简单的平滑替代方案。我们的选择主要集中在连续可微分指数线性单元（CELU）激活函数身上，因为它既具备平滑特性（与 ELU 不同），同时 PyTorch 实现的运行速度要比其它同类 Softplus 激活函数更快。除了平滑之外，CELU 还能够将 x 与 y 位移引入 ReLU，如下图所示。但由于我们使用的是批量规范，所以基本不会造成什么影响。

CELU 激活

我们粗略将平滑参数α手动调整为 0.075——请注意，这要比默认值 1 低得多。如此一来，我们的测试准确率达到 94.3%（50 次运行取平均值），这无疑是种显著的改善——意味着我们可以进一步减少 3 轮训练周期，20 轮训练周期的总耗时将为 52 秒，准确率则保持为 94.1%。

幽灵批量规范（46 秒）

批量规范最适合的批量大小一般在 32 左右。其中的原因可能与批量统计中的噪声有关，中等批量大小能够在有效规范化效果与少量噪声之间取得比较理想的平衡。

我们的批量大小为 512，而且进一步减少会严重增加训练时长；但我们发现，可以将批量规范分别应用于训练中的各批量子集。这种被称为“幽灵”批量规范的技术通常用于分布式设置，但也能够在单一节点的大批量处理流程中发挥重要作用。PyTorch 虽然不直接提供支持，但我们仍然能够较为轻松地建立自己的实现。

在 20 轮的训练周期之内，其将测试准确率提升到了 94.2%。随着训练时长变得越来越短，我们神经网络系统的学习率也在不断提高。如果我们将最大学习率提高到 50%，则完全能够在 18 个周期之内获得 94.1% 的准确率，这时的训练时长为 46 秒。

冻结批量规范比例（43 秒）

批量规范虽然对每条通道的均值与方差进行了标准化，但接下来我们还需要面对可学习规模与方差问题。我们的批量标准层由（平滑的）ReLU 继承，因此可学习方差允许网络对各通道的稀疏程度做出优化。另一方面，如果通道比例发生重大变化，则有可能会减少通道的有效数量并引发瓶颈。下面，我们一起来看训练期间这些参数的具体动态：

图中体现出不少情况，但至少有一点可以确定，即其比例并没有发生太大的学习与演变，这主要是因为我们对权重衰减做出严格控制。我们尝试将其冻结为 1/4 这一恒定值，大致相当于它们在训练中的中位平均值。最后一层的可学习比例稍高，但我们可以调整网络输出的比例进行补偿。

实际上，如果我们将 CELU 的α参数重新调整为补偿因子 4，并将批量规范方差的学习率与权重衰减分别设定为 42 与 1/42，那么则可将批量标准比例修正为 1。我们更喜欢这样的方式，因为它能够更明确地反映出通道比例对方差学习率动态的影响。

在 18 轮训练周期之后，我们的测试准确率提高到 94.2%。有趣的是，如果不提高批量规范方差的学习率，那么准确率会出现快速下降。这表明拥有学习能力的方差确实发挥着重要的作用——要么是在学习适当的稀疏度，要么是添加了正则化噪声。实际上，通过将方差的学习率提升至因子 4 并使用权重衰减除以这个因子，我们进一步使准确率实现小幅提升。

最后，我们可以使用增加的准确率将训练周期减少至 17 轮。新的测试准确率仍然是 94.1%，但最重要的是我们已经超越了使用 8 个 GPU 的 BaiduNet9P，目前成绩为 43 秒，位列排行榜第二！

输入补丁白化（36 秒）

批量规范可以很好地控制各个通道的分布，但并不能解决通道与像素之间的协方差问题。利用批量规范的“白化”版本，我们有望控制内部层的协方差；但这需要涉及额外的计算量以及相当复杂的实现过程。因此，我们决定把注意力放在更简单的输入层身上。

去除输入相关性的经典方法，在于执行全局 PCA（或者 ZCA）白化。我们提出了一种基于补丁的方法，该方法与总图像大小无关，而且更符合卷积网络的结构特点。我们将 PCA 白化应用在 3 x 3 的输入补丁当中，并将其作为具有固定（不可学习）权重的初始 3 x 3 卷积。接下来，我们再创建一个可学习的 1 x 1 卷积。指向该层的 27 条输入通道即是原始 3 x 3 x 3 输入块的转换后版本，其协方差矩阵基本恒等，这就使得优化变得更为轻松。

首先，我们需要对输入数据中 3 x 3 补丁的协方差矩阵的前导特征向量进行绘制。括号中的数字代表的是对应特征值的平方根，用以显示沿这些方向的相对变化比例；此外，我们还绘制包含两个符号的特征向量，用以说明变化的方向。正如我们所预料，图像中的局部亮度变化所占比例最高。

现在，我们用固定的 3 x 3 白化卷积替换网络中的第一个 3 x 3 卷积，用以均衡以上本征补丁的比例，而后添加一个可学习 1 x 1 卷积以观察训练结果中出现的变化。

在 17 轮训练之后，测试准确率达到 94.4%，我们可以进一步将训练周期减少 2 轮。15 轮训练共耗时 39 秒，测试准确率为 94.1%，已经接近 4-GPU 测试时间增强辅助下的 DAWNBench 冠军！如果我们将最大学习率提高 50% 并减少剪切增强量（即将补丁由 8 x 8 调整为 5 x 5）以补偿更高学习率带来的额外正则化，那么我们将能够在 36 秒之内达到 94.1% 的测试准确率，这时成绩就已经能够勉强在排行榜中登顶！

指数移动平均化（34 秒）

高学习率是快速训练的必要前提，因为其允许随机梯度下降能够在有限的时间内在参数空间中推进必要的距离。在另一方面，学习率还需要在训练结束时进行退火，以确保模型能够优化参数空间中的浮动与噪声区间。参数平均化方法使得训练能够以更高学习率推进，同时通过多次迭代求平均的方式尽可能缩小浮动与噪声范围。

我们采用的是参数的指标移动平均化技术，这是一种标准方法。为了保持效率，我们每 5 个批次进行一次平均化更新，因为我们发现更高的更新频率无法带来有意义的改进。我们需要选择一种新的学习率模式，确保能够在训练结束时带来更高的学习率以及平均化变化趋向。在学习率方面，最简单的选择当然是坚持我们一直使用的分段线性时间表，在最后 2 轮周期内以低固定值为基础——这里我们选择 0.99 的动量，以便在最后一轮训练中以大致相同的时间尺度进行平均化。

测试准确率提高到 94.3%，这意味着我们又有了减少训练周期的空间。13 轮训练之后，测试准确率为 94.1%，而且与文章开头时相比，我们的训练时长已经减少了 34 秒，性能比当时最先进的单一 GPU 技术提高了 10 倍！

测试时间增强（26 秒）

相信大家都希望自己的网络在输入内容进行水平翻转之后，能够以相同的方式进行图像分类。我们一直使用的方法是向网络提供大量数据，这样标签就会保留左右翻转情况以实现增强，而网络本身将最终通过广泛的训练过程学习到特征的一致性。

另一种方法则排除掉一切随机因素，其既能够显示输入图像，又能够显示水平翻转图像，并通过平均两个版本的网络输出结果达成共识，最终保证结果的一致性。这种非常聪明的方法，正是我们已经多次提到的测试时间增强技术。

在训练过程中，我们仍然向网络展示每幅图像的单一版本——可能包含随机翻转以作为数据增量，以便在不同的训练时期内呈现出不同的版本。另一种方法则是在训练过程中，使用与测试时相同的流程，并同时呈现图像及其镜像翻转版本。在后一种情况下，我们可以通过声明将图像拆分成两个独立的分支，其一为原始图像、其二为翻转图像，并在最后进行合并。通过这种方式，原始训练就成了一种权重约束之下的双分支网络随机训练流程，其中每个训练样例都会在完成后被“退出（dropout）”。

通过这种 dropout 训练，我们可以明确意识到任何有可能引入违反 TTA 规则的尝试都会遇到阻碍。从这个角度来看，我们实际上引入了一套更大的网络，并在其中采用有效的随机训练方法。另一方面，如果我们不限制测试过程当中的准备工作量，那么还有其它一些稍差的解决方案可供选择——这类方案中的训练花费时间更短，但数据集的存储时间却更长。

这些增强方法不仅与人工基准相关，也与最终用例有着密切的联系。在某些应用当中，我们需要保证分类准确率，在这种情况下必须使用 TTA。在另一些场景下，我们可能对推理时间有着严格要求，这时候最合理的方法当然是在约束范围之内最大程度提高准确率。这可能也是训练基准测试的最佳方法。

回到当前的案例，Kakao Brain 团队使用到我们之前提到的简单 TTA 方法——在推理时同时处理图像及其镜像翻转版本，这使得计算负担直接加倍。对于其它对称性差异（例如平移对称、高度 / 颜色变化等），我们当然可以采用更为广泛的 TTA 形式，但其同样会带来更高的计算成本。

由于此次项目基于轻量级计算系统——9 层 REsNet，所以包括 TTA 在内的总推理时间可能远低于竞赛早期出现的那些超过 100 层的网络。根据之前的讨论，任何用于限制此类方案的合理规则，都应基于推理时间约束而非对任意特征的实现方式。因此从这个角度来看，TTA 方法是完全可以接受的。

下面，我们看看 TTA 到底带来了哪些改进。我们将讨论范围控制在镜像翻转 TTA，这是为了与当前 DAWNBench 的提交方案保持一致，而且这似乎也是准确率与推理成本之间的最佳平衡点。凭借我们目前的网络方案与 13 轮训练周设置，TTA 的测试准确率为 94.6%，这也是我们目前研究中见到过的最佳效果。

如果我们移除余下的剪切数据增强——这会在本项目这类短期训练案例中造成负面影响——则能够将训练周期再减少 10 轮！这意味着只需要 26 秒，我们就能够实现 94.1% 的 TTA 测试准确率。

融合训练

这是一项简单的实验，用于研究我们观察到的训练速度增益是否能够转化为模型的最终准确率增益，即其是否会随着训练过程而出现收敛。我们的预测有着充分的理由，因为我们目前所使用的许多技术，其实都源自 ImageNet 的准确率融合场景！如果能够在 CIFAR10 数据集上将训练速度加快的技术，也可以提高 ImageNet 上的融合准确率，那就证明我们的成果完全可以应用到后一类问题当中。

与此前的实验不同，这一次我们只是粗略进行检验，更细致的探索将在后续逐步展开。我们将选择一个固定的学习率水平，将其设置得低一些以匹配较长时间的训练，同时将剪切增强量提升至 12 x 12，以确保在不造成过度拟合的前提下完成长时间训练。之前提到的其它超参数都保持不变，基准网络与最终网络的训练轮次数分别为 24 轮与 100 轮。最后，我们会打破所有规则，每项实验只运行 5 次！好了，下面来看具体结果：

虽然没有对最终训练设置中的超参数进行额外调整，也没有进行更长时间的运行，但我们的模型似乎仍然在整个 100 轮训练与接近收敛的阶段内保持着超越基准模型的领先优势。我们小小的 9 层 ResNet 在 80 轮训练周期内的最终 TTA 准确率为 96.1%——没有经过任何优化，准确率远高于 94% 的结果真的令人相当兴奋！如果配合适当的超参数优化，相信结果还会更上一层楼。

在大约 70 轮训练之后，也就是大约 3 分钟的总训练时间，我们的神经网络达到了 96% 的准确率。这也回答了一个始终困扰着我自己以及不少其他朋友的重要问题——DAWNBench 设定的 94% 准确率阈值似乎太低了。请注意，在到达 96% 准确率后我们并没有进行优化，因此后续训练可能反而导致准确率有所下降。

总结陈词

本系列文章到这里就结束了，最后还想跟大家再聊几句。首先感谢所有为项目做出贡献、提供支持以及反馈的朋友们。这里特别要感谢 Sam Davis，感谢 Thomas Read 在去年夏季的文章中提出的权重衰减观点，也要感谢 Myrtle 的每一位同事。

通过这个项目，我发现探索神经网络训练的动态、扩展他人的工作成果，并推动训练时间逐步缩短的过程非常有趣。我希望各位读者能够从中得到启发，并坚信通过算法的进一步完善，训练时间仍有巨大的压缩空间（或者说准确率还有巨大的提升空间，具体取决于您的要求）。

在本系列文章的开头，我曾半开玩笑地提到，如果我们要达到 100% 的计算效率，那么训练应该能够在 40 秒左右完成。但最终结果让我惊讶地发现，实际情况要比预想快得多，而且计算效率也远远没有达到 100%！这让我更加坚定了一切皆可进一步优化的信念。

原文链接：

https://myrtle.ai/how-to-train-your-resnet-8-bag-of-tricks/

本文亦发表在 Colab 笔记当中，详见此处：

https://colab.research.google.com/github/davidcpage/cifar10-fast/blob/master/bag_of_tricks.ipynb

你也「在看」吗？👇