本文转载自机器之心

仅用于学术分享，若侵权请联系删除

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

CV各大方向专栏与各个部署框架最全教程整理

计算机视觉入门1v3辅导班

如何提升 PyTorch「炼丹」速度？

最近，知名机器学习与 AI 研究者 Sebastian Raschka 向我们展示了他的绝招。据他表示，他的方法在不影响模型准确率的情况下，仅仅通过改变几行代码，将 BERT 优化时间从 22.63 分钟缩减到 3.15 分钟，训练速度足足提升了 7 倍。

作者更是表示，如果你有 8 个 GPU 可用，整个训练过程只需要 2 分钟，实现 11.5 倍的性能加速。

下面我们来看看他到底是如何实现的。

让 PyTorch 模型训练更快

首先是模型，作者采用 DistilBERT 模型进行研究，它是 BERT 的精简版，与 BERT 相比规模缩小了 40%，但性能几乎没有损失。其次是数据集，训练数据集为大型电影评论数据集 IMDB Large Movie Review，该数据集总共包含 50000 条电影评论。作者将使用下图中的 c 方法来预测数据集中的影评情绪。

基本任务交代清楚后，下面就是 PyTorch 的训练过程。为了让大家更好地理解这项任务，作者还贴心地介绍了一下热身练习，即如何在 IMDB 电影评论数据集上训练 DistilBERT 模型。如果你想自己运行代码，可以使用相关的 Python 库设置一个虚拟环境，如下所示：

相关软件的版本如下：

现在省略掉枯燥的数据加载介绍，只需要了解本文将数据集划分为 35000 个训练示例、5000 个验证示例和 10000 个测试示例。需要的代码如下：

代码部分截图

完整代码地址：https://github.com/rasbt/faster-pytorch-blog/blob/main/1_pytorch-distilbert.py

然后在 A100 GPU 上运行代码，得到如下结果：

正如上述代码所示，模型从第 2 轮到第 3 轮开始有一点过拟合，验证准确率从 92.89% 下降到了 92.09%。在模型运行了 22.63 分钟后进行微调，最终的测试准确率为 91.43%。

使用 Trainer 类

接下来是改进上述代码，改进部分主要是把 PyTorch 模型包装在 LightningModule 中，这样就可以使用来自 Lightning 的 Trainer 类。部分代码截图如下：

完整代码地址：https://github.com/rasbt/faster-pytorch-blog/blob/main/2_pytorch-with-trainer.py

上述代码建立了一个 LightningModule，它定义了如何执行训练、验证和测试。相比于前面给出的代码，主要变化是在第 5 部分（即 ### 5 Finetuning），即微调模型。与以前不同的是，微调部分在 LightningModel 类中包装了 PyTorch 模型，并使用 Trainer 类来拟合模型。

之前的代码显示验证准确率从第 2 轮到第 3 轮有所下降，但改进后的代码使用了 ModelCheckpoint 以加载最佳模型。在同一台机器上，这个模型在 23.09 分钟内达到了 92% 的测试准确率。

需要注意，如果禁用 checkpointing 并允许 PyTorch 以非确定性模式运行，本次运行最终将获得与普通 PyTorch 相同的运行时间（时间为 22.63 分而不是 23.09 分）。

自动混合精度训练

进一步，如果 GPU 支持混合精度训练，可以开启 GPU 以提高计算效率。作者使用自动混合精度训练，在 32 位和 16 位浮点之间切换而不会牺牲准确率。

在这一优化下，使用 Trainer 类，即能通过一行代码实现自动混合精度训练：

上述操作可以将训练时间从 23.09 分钟缩短到 8.75 分钟，这几乎快了 3 倍。测试集的准确率为 92.2%，甚至比之前的 92.0% 还略有提高。

使用 Torch.Compile 静态图

最近 PyTorch 2.0 公告显示，PyTorch 团队引入了新的 toch.compile 函数。该函数可以通过生成优化的静态图来加速 PyTorch 代码执行，而不是使用动态图运行 PyTorch 代码。

由于 PyTorch 2.0 尚未正式发布，因而必须先要安装 torchtriton，并更新到 PyTorch 最新版本才能使用此功能。

然后通过添加这一行对代码进行修改：

在 4 块 GPU 上进行分布式数据并行

上文介绍了在单 GPU 上加速代码的混合精度训练，接下来介绍多 GPU 训练策略。下图总结了几种不同的多 GPU 训练技术。

想要实现分布式数据并行，可以通过 DistributedDataParallel 来实现，只需修改一行代码就能使用 Trainer。

经过这一步优化，在 4 个 A100 GPU 上，这段代码运行了 3.52 分钟就达到了 93.1% 的测试准确率。

DeepSpeed

最后，作者探索了在 Trainer 中使用深度学习优化库 DeepSpeed 以及多 GPU 策略的结果。首先必须安装 DeepSpeed 库：

接着只需更改一行代码即可启用该库：

这一波下来，用时 3.15 分钟就达到了 92.6% 的测试准确率。不过 PyTorch 也有 DeepSpeed 的替代方案：fully-sharded DataParallel，通过 strategy="fsdp" 调用，最后花费 3.62 分钟完成。

以上就是作者提高 PyTorch 模型训练速度的方法，感兴趣的小伙伴可以跟着原博客尝试一下，相信你会得到想要的结果。

原文链接：https://sebastianraschka.com/blog/2023/pytorch-faster.html

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

计算机视觉入门1v3辅导班

【技术文档】《从零搭建pytorch模型教程》122页PDF下载

QQ交流群：470899183。群内有大佬负责解答大家的日常学习、科研、代码问题。

其它文章

​AI最全资料汇总 | 基础入门、技术前沿、工业应用、部署框架、实战教程学习

改变几行代码，PyTorch炼丹速度狂飙、模型优化时间大减

AAAI 2023 | 轻量级语义分割新范式： Head-Free 的线性 Transformer 结构

TSCD：弱监督语义分割新方法，中科院自动化所和北邮等联合提出

如何用单个GPU在不到24小时的时间内从零开始训练ViT模型？

CVPR 2023 | 基于Token对比的弱监督语义分割新方案！

比MobileOne还秀，Apple将重参数与ViT相结合提出FastViT

CVPR 2023 | One-to-Few：没有NMS检测也可以很强很快

ICLR 2023 | Specformer: Spectral GNNs Meet Transformers

重新审视Dropout

RestoreDet：低分辨率图像中目标检测

AAAI 2023 | 打破NAS瓶颈，AIO-P跨任务网络性能预测新框架

CLIP：语言-图像表示之间的桥梁​

目标检测Trick | SEA方法轻松抹平One-Stage与Two-Stage目标检测之间的差距

少样本学习综述：技术、算法和模型

CVPR 2023 | 标注500类，检测7000类！清华大学等提出通用目标检测算法UniDetector

CVPR 2023｜基于多层多尺度重建任务的MIM改进算法

CVPR 2023 | 超越MAE！谷歌提出MAGE：图像分类和生成达到SOTA！

称霸Kaggle的十大深度学习技巧

CVPR 2023 | 用于半监督目标检测的知识蒸馏方法

目标跟踪方向开源数据集资源汇总

CVPR2023 | 书生模型霸榜COCO目标检测，研究团队解读公开

Vision Transformer的重参化也来啦 | RepAdpater让ViT起飞

高效压缩99%参数量！轻量型图像增强方案CLUT-Net开源

一文了解 CVPR 2023 的Workshop 都要做什么

CVPR'23 最新 70 篇论文分方向整理｜包含目标检测、图像处理、人脸、医学影像、半监督学习等方向

目标检测无痛涨点新方法 | DRKD蒸馏让ResNet18拥有ResNet50的精度

PyTorch 2.0正式版来了！

CVPR2023最新Backbone | FasterNet远超ShuffleNet、MobileNet、MobileViT等模型

CVPR2023 | 集成预训练金字塔结构的Transformer模型

AAAI 2023 | 一种通用的粗-细视觉Transformer加速方案

大核分解与注意力机制的巧妙结合，图像超分多尺度注意网络MAN已开源！

计算机视觉各个方向交流群与知识星球

CV小知识讨论与分析(7) 寻找论文创新点的新方式

CV小知识分析与讨论(6）论文创新的一点误区

计算机视觉入门1v3辅导班

计算机视觉交流群

聊聊计算机视觉入门