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

新闻 前端
近日一位名为 Florian Ernst 的博主却发现 PyTorch Lightning 存在一个 bug——让原本应该加速的训练变得更慢了。

 [[436675]]

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

[[436676]]

本文作者 Florian Ernst

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

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

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

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

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

[[436677]]

问题出在哪里?

我首先运行 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, 

因此,有两种可能性:

  • 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 秒,相比较而言,速度确实提升了很多。

[[436678]]

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

责任编辑:张燕妮 来源: 机器之心Pro
相关推荐

2019-09-02 08:58:27

Python编译器编程语言

2022-02-03 14:42:37

Windows 11微软Windows

2011-09-02 13:54:24

Windows优化宽带

2023-10-05 09:40:06

Next.jsTurbopackVite

2023-03-13 16:00:05

训练模型

2020-11-09 14:17:13

Rust编码语言

2021-08-03 08:13:48

重构API代码

2012-12-12 16:16:19

微软IE10Windows 8

2022-08-07 21:17:27

微软Windows 11

2009-06-02 09:09:52

2020-10-23 07:39:39

深度学习PyTorch代码

2019-12-23 08:57:50

Python代码单线程

2021-06-09 08:00:00

Python编程语言开发

2024-03-11 15:47:11

RustPython代码

2022-04-27 09:24:22

前端代码速度

2010-05-25 10:37:00

三网融合网络加速网宿科技

2010-10-27 14:37:30

Android

2016-12-26 14:46:09

宽带无线网络Ovum

2017-04-24 16:44:41

宽带千兆速度

2011-05-24 16:03:30

C++编译速度
点赞
收藏

51CTO技术栈公众号