使用 PyTorch Lightning 将深度学习管道速度提高 10 倍

前言

本文介绍了如何使用 PyTorch Lightning 构建高效且快速的深度学习管道，主要包括有为什么优化深度学习管道很重要、使用 PyTorch Lightning 加快实验周期的六种方法、以及实验总结。

当 Alex Krizhevsky、Ilya Sutskever 和 Geoffrey Hinton 在 2012 年设计 AlexNet 时，训练 6000 万参数模型需要五到六天的时间。八年后的 2020 年，微软 DeepSpeed 团队在不到 44 分钟的时间内成功训练了一个 3.5 亿参数的 Large-Bert 模型！

九年后，我们现在看到，AlexNet 只是机器学习革命的冰山一角。今天，我们知道许多尚未开发的潜在训练技术和深度学习模型架构都在我们的掌握之中！

不幸的是，由于数据的规模和新的深度学习模型架构的规模，其中许多进步对于普通研究人员来说就像多汁苹果对于没有梯子的水果采摘者一样难以获得。有这么多卓有成效的模型架构挂在深度学习潜力之树上，我们应该问自己，“我们如何才能达到它们？”

答案很简单：要达到这些富有成效的架构，我们需要梯子！Alex Krizhevsky 构建了他自己的梯子来逐块到达 AlexNet，但是今天，像 PyTorch Lightning 这样的解决方案为您提供了自己的现成梯子——甚至是自动扶梯！

本文介绍了如何使用 PyTorch Lightning 构建高效且快速的深度学习管道，还解释了这些优化如何通过显着加快研发实验周期来快速尝试各种研究想法！

为什么优化深度学习管道很重要
使用 PyTorch Lightning 加快实验周期的六种方法
结果总结

为什么优化深度学习管道很重要

无论是在学术界还是在工业界从事研究，研发探索和尝试新想法的时间和资源总是有限的。随着数据集的规模和深度学习模型的复杂性不断增加，对最新机器学习模型和技术的实验日益复杂和耗时。如何应对这些挑战（并使研发周期更有效率）对项目的整体成功起着至关重要的作用。

如今，存在各种解决方案来克服这些障碍，例如 Grid.ai、WandB 和 PyTorch Lightning。本文将重点介绍 PyTorch Lightning，并解释如何使用它使深度学习管道更快，并在需要最少代码更改的情况下在幕后提高内存效率。使用这些解决方案，可以使实验更具可扩展性和迭代速度更快，同时最大限度地减少潜在的错误。进行这些更改将减少实验所需的时间，节省的时间可以用来尝试更多的想法。

使用 PyTorch Lightning 加快实验周期的六种方法

优化深度学习管道的六种方法：

并行数据加载
多GPU训练
混合精度训练(Mixed precision training)
分片训练(Sharded training)
提前停止(Early stopping)
模型评估和推理期间的优化

对于每一种方法，我们都会简要解释它的工作原理，如何实现它，最后，分享我们是否发现它对我们的项目有帮助！

并行数据加载

数据加载和增强步骤成为训练管道中的瓶颈是很常见的。

典型的数据管道包含以下步骤：

从磁盘加载数据
即时创建随机增强
将每个样本整理成批

数据加载和增强过程非常容易并行，可以通过使用多个 CPU 进程并行加载数据来优化。这样一来，昂贵的 GPU 资源就不会在训练和推理过程中受到 CPU 的阻碍。

为了尽快加载数据以训练深度学习模型，可以执行以下操作：

将 DataLoader 中的 `num_workers` 参数设置为 CPU 的数量。
使用 GPU 时，将 DataLoader 中的 `pin_memory` 参数设置为 True。这会将数据分配到页面锁定内存中，从而加快向 GPU 传输数据的速度。

补充说明：

如果处理流数据（即`IterableDataset`），还需要配置每个worker以独立处理传入的数据。
种子初始化错误困扰着许多开源深度学习项目。为避免该错误，请在 `worker_init_fn` 中定义工作进程的进程种子。从 PyTorch Lightning 1.3 开始，这会使用 `seed_everything(123, workers=True)` 自动处理。
从 PyTorch 1.8 开始，可以使用可选的 `prefetch_factor` 参数更好地控制加载性能行为。将此设置为更高的整数以提前加载更多批次，但会占用更大的内存。

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

GPU 为 CPU 的训练和推理时间提供了巨大的加速。什么比 GPU 更好？多个GPU！

PyTorch 中有一些范例可用于训练具有多个 GPU 的模型。两个更常见的范例是“DataParallel”和“DistributedDataParallel”，其中“DistributedDataParallel”是一种更具可扩展性的方法。

在 PyTorch（和其他平台）中修改训练管道并非易事。人们必须考虑诸如以分布式方式加载数据以及权重、梯度和指标的同步等问题。

使用 PyTorch Lightning，能够非常轻松地在多个 GPU 上训练 PyTorch 模型，几乎无需更改代码！

混合精度

默认情况下，输入张量以及模型权重以单精度 (float32) 定义。但是，某些数学运算可以半精度 (float16) 执行。这会显着提高速度并降低模型内存带宽，而不会牺牲模型性能。

通过在 PyTorch Lightning 中设置混合精度标志，框架会在可能的情况下自动使用半精度，同时在其他地方保留单精度。通过最少的代码修改，能够将模型训练时间提高 1.5 到 2 倍。

提前停止

模型需要训练大量的 epoch，但实际上模型在训练过程的早期就很可能过度拟合了训练数据。因此，需要在训练管道中实施提前停止。提前停止被配置为在预定义的评估次数后验证损失停止减少时结束训练。通过这样做，不仅可以防止过度拟合，而且还可以节省时间，在数十个而不是数百个 epoch 内找到最佳模型。