pytorch-lightning - 在多 GPU Pytorch Lightning 上训练时在 1 个 GPU 上运行验证

Question

有什么方法可以在使用 DDP 的多个 GPUvalidation_step时在单个 GPU 上执行方法。training_step

我想做的原因是因为我想要实现的几个指标需要完全访问数据，并且在单个 GPU 上运行将确保这一点。我已经尝试过validation_step_end方法，但不知何故我只得到了部分数据。那篇文章在这里：Stack Overflow Post

Answer 1

恐怕这是不可能的。但是有一个TorchMetrics包，它在开发时考虑到了多 GPU 支持，因此当您的自定义指标源自 TM 时，您甚至可以在多 GPU 设置上运行。