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

2023-03-12 04:03

选自Sebastian Raschka博客

机器之心编译

机器之心编辑部

关于 PyTorch 炼丹，本文作者表示：「如果你有 8 个 GPU，整个训练过程只需要 2 分钟，实现 11.5 倍的性能加速。」

如何提升 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 类来拟合模型。