用了 Lightning 训练速度反而更慢，你遇到过这种情况吗？
PyTorch Lightning 是一种重构 PyTorch 代码的工具，它可以抽出代码中复杂重复的部分，使得 AI 研究可扩展并且可以快速迭代。然而近日一位名为 Florian Ernst 的博主却发现 PyTorch Lightning 存在一个 bug——让原本应该加速的训练变得更慢了。
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,
)
我在 GitHub 上创建了一个 issue，希望 Lightning 团队意识这个问题，接下来我要寻找问题根源。GitHub 地址：https://github.com/PyTorchLightning/pytorch-lightning/issues/10389Lightning 的 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 非常简单：我将 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© THE END
转载请联系机器之心公众号获得授权
投稿或寻求报道：content@jiqizhixin.com
