Pytorch Lightning ddp模式复制主脚本

Pytorch Lightning duplicates main script in ddp mode

当我使用 ddp 模式(2 个 GPU)在集群上启动我的主脚本时,Pytorch Lightning 会复制主脚本中执行的任何内容,例如打印或其他逻辑。我需要一些扩展的训练逻辑,我想自己处理。例如。在 Trainer.fit() 之后做某事(一次!)。但是随着主脚本的重复,这并不像我想要的那样工作。我还尝试将它包装在 if __name__ == "__main__" 中,但它不会改变行为。如何解决这个问题?或者,如何在没有重复项的情况下围绕我的 Trainer 对象使用一些逻辑?

从那以后,我开始在 PyTorch 中使用原生“ddp”进行多处理。据我了解,PytorchLightning (PTL) 只是 运行 你的主脚本在多个 GPU 上多次运行。如果您只想在脚本的一次调用中拟合您的模型,这很好。然而,在我看来,一个巨大的缺点是在训练过程中失去了灵活性。与您的实验交互的唯一方法是通过这些(记录不完整的)回调。老实说,在 PyTorch 中使用原生多处理要灵活方便得多。最终它实现起来要快得多,也容易得多,而且你不必通过 PTL 文档搜索年龄来实现简单的事情。 我认为 PTL 正在朝着一个好的方向前进,移除大部分样板,但是,在我看来,Trainer 概念需要一些认真的修改。在我看来太封闭了,违背了PTL自己的“重组PyTorch代码,保留原生PyTorch代码”的理念。 如果你想使用PTL来轻松进行多GPU训练,我个人强烈建议不要使用它,对我来说那是浪费时间,最好学习原生的PyTorch multiprocessing。

您可以在Trainer.fit之后添加以下代码来退出重复的子进程:

import sys
if model.global_rank != 0:
    sys.exit(0)

其中 model 继承自 LightningModule,它有一个 属性 global_rank 指定机器的等级。我们可以大致理解为gpu id或者process id。此代码之后的所有内容只会在主进程中执行,即 global_rank = 0.

的进程

更多信息,请参考文档https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#global_rank

在 GitHub 存储库中询问了这个问题:https://github.com/PyTorchLightning/pytorch-lightning/issues/8563

有不同的训练加速器,虽然 DDP (DistributedDataParallel) 每个 GPU 运行一次脚本,ddp_spawn 而 dp 则不会。

但是,某些插件(如 DeepSpeedPlugin)是基于 DDP 构建的,因此更改加速器不会多次停止 运行 中的主脚本。