有bug！用Pytorch Lightning重构代码速度更慢，修复后速度倍增

机器之心

发布于 2023-03-29 18:00:58

8170

发布于 2023-03-29 18:00:58

文章被收录于专栏：机器之心

选自Medium

作者：Florian Ernst

机器之心编译

编辑：小舟、陈萍

用了 Lightning 训练速度反而更慢，你遇到过这种情况吗？

PyTorch Lightning 是一种重构 PyTorch 代码的工具，它可以抽出代码中复杂重复的部分，使得 AI 研究可扩展并且可以快速迭代。然而近日一位名为 Florian Ernst 的博主却发现 PyTorch Lightning 存在一个 bug——让原本应该加速的训练变得更慢了。

本文作者 Florian Ernst

Ernst 撰写博客详细描述了他发现这个 bug 的过程，以下是博客原文。

两周前，我将一些深度学习代码重构为 Pytorch Lightning，预计大约有 1.5 倍的加速。然而，训练、评估和测试任务的速度却降为原来的 1/4。重构之后的神经网络需要运行几天才能得出结果，因此我想找出原因，并尽可能地减少训练时间。

事情是这样的，我使用的是一些开源深度学习代码，这些代码是用来展示某些机器学习任务最新架构的。然而这些代码本身既不整洁也没进行优化。我注意到几个可以加速的地方，并将代码重构为 Pytorch 代码，让训练大约快了 3 倍。

但我认为还有改进的余地。Pytorch Lightning 是一个非常好的工具：它删除了大量样板代码，并配备了一些优化方法，因此我决定使用 Lightning 重构这些代码。

我原本希望代码大约能提速 1.5 倍，但完成重构时，我惊讶地发现迭代时间从 4 秒变成了 15 秒，这使训练时间多了近 3 倍。

问题出在哪里？

我首先运行 Lightning 的分析器来找出问题所在。

基础分析器给了我一个起点：大部分时间都花在运行一个 epoch 上；高级分析器没有给我更多信息。

我想知道我是否在神经网络上错误地配置了一些超参数。我打乱了其中一些超参数，训练速度没有任何变化。

然后我调整了数据加载器，发现改变作业数 n_jobs 会对总训练时间产生影响。然而影响不是加快了计算速度，而是减慢了。

随着 job 数变化，100 个 epoch 花费的时间。

使用 n_jobs=0 完全禁用多处理使我的迭代几乎比使用 6 个内核快了 2 倍。默认情况下，Pytorch 在两个 epoch 之间会 kill 掉运行中的进程（worker）并重新加载，因而需要重新加载数据集。

在我这个例子中，加载数据集非常慢。我将 DataLoader 里的 persistent_workers 参数设置为 True，以防止运行中的进程被杀死，进而防止重新加载数据。

# My data Loader parameters
DataLoader(
  train_dataset, batch_size=64, shuffle=True, num_workers=n_workers,
  persistent_workers=True, pin_memory=True,
)

因此，有两种可能性：

Pytorch Lightning kill 掉 worker，没有考虑 persistent_workers 参数；
问题出在别的地方。

我在 GitHub 上创建了一个 issue，希望 Lightning 团队意识这个问题，接下来我要寻找问题根源。

GitHub 地址：https://github.com/PyTorchLightning/pytorch-lightning/issues/10389

寻找问题根源

Lightning 的 profiler 与上下文管理器一起运行并计算给定块花费的时间。它可以轻松搜索特定的 profiler 操作，以运行「run_training_epoch」为例。

我开始探究 Lightning 源码，查看导致循环（loops）变慢的指令，我发现了一些问题：Loop.run 调用 Loop.on_run_start、Loop.on_run_start 重新加载 dataloader，如下图所示：

Loop.run 调用 Loop.on_run_start…

Loop.on_run_start 重新调用 dataloader

问题看起来确实来自在每个 epoch 中重新加载 DataLoader。查看 DataLoader 的源码，发现是这样的：

当使用 persistent_workers > 0 迭代 DataLoader 时，如果_iterator` 为 None，则使用_get_iterator() 重新加载整个数据集。可以确定的是 Pytorch Lightning 错误地重置了 _iterator，从而导致了这个问题。

为了证实这一发现，我用一个自定义的只能重载的__iter__方法替换了 DataLoader：

正如预期的那样，在迭代之后，_iterator 属性被正确设置，但在下一个 epoch 开始之前被重置为 None。

n_jobs=1，persistent_workers=True

现在，我只需要知道属性何时被设置为 None ，这样就可找到问题的根源。我尝试使用调试器，但由于多进程或 CUDA 而导致程序崩溃。我开始采用 Python 的 getter & setter 用法：

当 DataLoader._iterator 设置为 None 时，将会打印 stack trace

这样做非常有效，会输出如下内容：

File "trainer\trainer.py", line 1314, in _run_train
  self.fit_loop.run()
...
File "loops\fit_loop.py", line 234, in advance
  self.epoch_loop.run(data_fetcher)
File "loops\base.py", line 139, in run
  self.on_run_start(*args, **kwargs)
File "loops\epoch\training_epoch_loop.py", line 142, in on_run_start
  self._dataloader_iter = _update_dataloader_iter(...)
File "loops\utilities.py", line 121, in _update_dataloader_iter
  dataloader_iter = enumerate(data_fetcher, batch_idx)
File "utilities\fetching.py", line 198, in __iter__
  self.reset()
File "utilities\fetching.py", line 212, in reset
  self.dataloader.reset()
...
File "trainer\supporters.py", line 498, in _shutdown_workers_and_reset_iterator
  dataloader._iterator = None

通过跟踪发现每次开始运行时都会调用 DataLoader.reset。通过深入研究代码后，我发现每次迭代都会重置 DataFetcher，从而导致 DataLoader 也被重置。代码中没有条件来避免重置：每个 epoch 都必须重置 DataLoader。

这就是我发现迭代缓慢的根本原因。

修复 bug

既然发现了 bug，就要想办法修复。修复 bug 非常简单：我将 self.reset 行从 DataFetcher 的__iter__ 方法中移除：

通过修改后再次训练，现在一次迭代只需要 1.5 秒，而此前需要 15 秒，使用 vanilla Pytorch 也需要 3 秒，相比较而言，速度确实提升了很多。

我将发现的这个 bug 报告给了 Lightning 团队，他们对问题进行了修复并在第二天推送了修补程序。我随后更新了库，更新后发现他们的修复确实有效。相信更多人将从这次修复中受益，并且他们的 Lightning 模型的训练和测试时间会得到改善。如果你最近还没有更新依赖项，请尝试安装 pytorch-lightning==1.5.1 或更高版本！

原文链接：https://medium.com/@florian-ernst/finding-why-pytorch-lightning-made-my-training-4x-slower-ae64a4720bd1

详解NVIDIA TAO系列分享第2期：

基于Python的口罩检测模块代码解析——快速搭建基于TensorRT和NVIDIA TAO Toolkit的深度学习训练环境

第2期线上分享将介绍如何利用NVIDIA TAO Toolkit，在Python的环境下快速训练并部署一个人脸口罩监测模型，同时会详细介绍如何利用该工具对模型进行剪枝、评估并优化。

TAO Toolkit 内包含了150个预训练模型，用户不用从头开始训练，极大地减轻了准备样本的工作量，让开发者专注于模型的精度提升。本次分享摘要如下：