文章图片

我开始探究 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 或更高版本！