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

 [[436675]]

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 或更高版本！