深度学习流水线提速10倍的六招

面对海量图片数据，如何快速完成实验成了常见难题。

无论是在学术研究还是工业应用中，这个问题都时常困扰着研究者与工程师。

近日有一位国外研究者给出了一种思路：在现有深度学习框架的基础上，显著提升流水线速度，达到约10倍的加速效果。

用他的话来说，就像在爬楼时突然给你装上了电梯。

那么，这样的提升到底是如何实现的呢？

优化机器学习流水线极为重要

无论你处于学术界还是工业界，时间和资源往往成为限制实验进展的关键因素。

随着数据集规模与模型复杂度的上升，实验成本往往变得更高，提速就显得尤为关键。

曾经，像训练一个大型图像模型需要数天的时间，而现在在更大的数据集上也能在较短时间内完成训练。

从某种角度看，这得益于各类技术工具的出现，而其中一种便是用于提升训练效率的框架/工具。

于是，研究者总结出六种“闪电加速”实验周期的方法。

并行数据加载

数据加载与数据增强往往被视为训练流水线的瓶颈之一。

一个典型的训练流水线通常包含：从磁盘读取数据、在运行时执行随机增强、对样本进行批处理。

在这过程中，可以通过多进程并行加速数据加载，但还可以进一步优化：

1、在数据加载器中将 workers 参数设置为可用 CPU 核心数量。

2、在与 GPU 协同工作时，将数据加载器中的 pin_memory 设置为 True，以把数据分配到页锁定内存中，从而提升到 GPU 的传输速度。

使用分布式数据并行进行多GPU训练

相较于单个 CPU，GPU 已大幅缩短了训练与推理时间，但是否存在更优解？答案可能在于多个 GPU。

在深度学习框架中，有多种范式可用于多 GPU 训练。

较为常见的两种是“DataParallel”和“DistributedDataParallel”。

研究者采用了后者，因为它通常具有更好的扩展性。

不过，在改动训练流水线、实现分布式数据加载、以及权重、梯度和指标的同步等方面，难度并不小。

不过，借助相应的框架，能够在多 GPU 上训练模型往往变得更加容易，几乎无需重写大量代码。

混合精度

默认情况下，输入张量与模型权重使用单精度（float32）。

一些运算可以用半精度（float16）完成，因而可显著提升速度、降低显存带宽占用，同时通常不会牺牲模型性能。

通过在训练框架中开启混合精度支持，系统会在可能处自动使用半精度，其他情况仍保持单精度。

仅需最小改动即可使训练速度提升约1.5到2倍。

早停法

在训练深度学习模型时，通常希望获得良好的泛化性能。然而，许多模型容易在训练集上持续降低损失而对验证集表现变差。

因此，早停法在训练过程中介入：当验证集损失在若干轮评估后不再改善时就停止训练。

这样不仅可以降低过拟合风险，还能在若干个 epoch 内找到较优模型。

分布式训练（分布式训练中的零冗余技术）

分布式训练技术通过在多设备/节点上并行计算，显著提升对大模型的处理能力。

其中，零冗余训练（Zero Redundancy Optimizer，ZeRO）等方法，能够减少内存冗余、提升扩展性。

在某些训练场景下，它们对大模型的可扩展性尤为关键。

模型评估与推理中的优化

在评估与推理阶段，梯度并不需要参与前向计算，因此可以将评估代码放在 Torch.no_grad 环境中。

这能避免梯度存储，降低显存占用，从而在相同显存条件下部署更大批量的输入，提升评估与推理效率。

效果如何？

研究者整理了一张表格，详细列出每种方法带来的加速效果。

这些方法对你是否也有帮助？

不妨亲自尝试，看看在你的场景中能带来多少提升。