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

时间:2022-11-21 17:43:02

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

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

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

本文作者 Florian Ernst

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

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

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

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

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

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

问题出在哪里?

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

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

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

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

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

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

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

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

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


  1. # My data Loader parameters
  2. DataLoader(
  3. train_dataset, batch_size=64, shuffle=True, num_workers=n_workers,
  4. persistent_workers=True, pin_memory=True,
  5. )

因此,有两种可能性:

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

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

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

寻找问题根源

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

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

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

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

Loop.run 调用 Loop.on_run_start…

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

Loop.on_run_start 重新调用 dataloader

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

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

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

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

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

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

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

n_jobs=1,persistent_workers=True

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

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

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

这样做非常有效,会输出如下内容:

									
  1. File "trainer\trainer.py", line 1314, in _run_train
  2. self.fit_loop.run()
  3. ...
  4. File "loops\fit_loop.py", line 234, in advance
  5. self.epoch_loop.run(data_fetcher)
  6. File "loops\base.py", line 139, in run
  7. self.on_run_start(*args, **kwargs)
  8. File "loops\epoch\training_epoch_loop.py", line 142, in on_run_start
  9. self._dataloader_iter = _update_dataloader_iter(...)
  10. File "loops\utilities.py", line 121, in _update_dataloader_iter
  11. dataloader_iter = enumerate(data_fetcher, batch_idx)
  12. File "utilities\fetching.py", line 198, in __iter__
  13. self.reset()
  14. File "utilities\fetching.py", line 212, in reset
  15. self.dataloader.reset()
  16. ...
  17. File "trainer\supporters.py", line 498, in _shutdown_workers_and_reset_iterator
  18. dataloader._iterator = None

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

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

修复 bug

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

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

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

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

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

原文链接:https://www.toutiao.com/a7034391328750371361/