选自Medium
作者：Florian Ernst
机器之心编译
编辑：小舟、陈萍

用了 Lightning 训练速度反而更慢 ， 你遇到过这种情况吗？

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」为例。