使用 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 内找到最佳模型。

分片训练

分片训练基于微软的 ZeRO 研究和 DeepSpeed 库，这使得训练大型模型具有可扩展性和简单性。这是通过使用各种内存和资源间通信优化来实现的。实际上，分片训练可以训练大型模型，否则这些模型将不适合单个 GPU 或在训练和推理期间使用更大的批次大小。

PyTorch Lightning 在其 1.2 版本中引入了对分片训练的支持。在我们的用例中，我们没有观察到训练时间或内存占用的任何显着改进。但是，我们的见解可能无法推广到其他问题和设置，可能值得一试，尤其是当处理不使用单个 GPU 的大型模型时。

模型评估和推理期间的优化

在模型评估和推理期间，模型的前向传递不需要梯度。因此，可以将评估代码包装在一个 `torch.no_grad` 上下文管理器中。这可以防止在前向传递期间存储梯度，从而减少内存占用。因此，可以将更大的批次输入模型中，从而实现更快的评估和推理。

默认情况下，PyTorch Lightning 在幕后管理这些优化。

结果总结

在我们的实验中，我们发现所有优化都独立地减少了训练深度学习模型的时间，除了分片训练，我们没有观察到任何速度或内存改进。

下表是改进深度学习管道所做的每项优化，以及观察到的性能提升。

通过这些优化，我们把深度学习管道提高了10倍的速度，从两周节省到只要10小时。

作者：Georgian

编译：CV技术指南

原文链接：https://devblog.pytorchlightning.ai/how-we-used-pytorch-lightning-to-make-our-deep-learning-pipeline-10x-faster-731bd7ad318a

欢迎关注公众号 CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读。

在公众号中回复关键字 “技术总结”可获取公众号原创技术总结文章的汇总pdf。

其它文章

归一化方法总结 | 欠拟合与过拟合技术总结

NMS总结 | 损失函数技术总结

注意力机制技术总结 | 特征金字塔技术总结

池化技术总结 | 数据增强方法总结

论文创新的常见思路总结| GPU多卡并行训练总结

CNN可视化技术总结（四）可视化工具与项目

计算机视觉中的图像标注工具总结

各种 Optimizer 梯度下降优化算法回顾和总结

汇总 | 国内外经典开源数据集

Softmax 函数和它的误解

提高机器学习模型性能的常用策略

资源分享 | SAHI：超大图片中对小目标检测的切片辅助超推理库

计算机视觉中的图像标注工具总结

Batch Size对神经网络训练的影响

神经网络超参数的调参方法总结

使用 Ray 将 PyTorch 模型加载速度提高 340 倍

计算机视觉中的图像标注工具总结

2021年小目标检测最新研究综述

经典论文系列--胶囊网络：新的深度学习网络

计算机视觉专业术语总结(一)构建计算机视觉的知识体系

计算机视觉中的小样本学习综述