validation_epoch_end 与 DDP Pytorch 闪电

Posted 2023-03-27

【中文标题】validation_epoch_end 与 DDP Pytorch 闪电

【英文标题】：validation_epoch_end with DDP Pytorch Lightning

【发布时间】：2021-03-03 08:48:03

【问题描述】：

你的问题是什么？

我正在尝试实现一个需要访问整个数据的指标。因此，我没有在 *_step() 方法中更新指标，而是尝试在 *_epoch_end() 方法中收集输出。但是，输出仅包含每个设备获取的数据分区的输出。基本上，如果有 n 个设备，那么每个设备将获得总输出的 1/n。

你的环境是什么？

OS: ubuntu
Packaging: conda
Version [1.0.4
Pytorch: 1.6.0

【问题讨论】：

我也面临类似的问题，但不是指标。我正在尝试返回测试数据集的预测。它只返回块。他们宣传一个简单的多 GPU 设置，但它不是那样的，至少对我来说是这样。

【参考方案1】：

请参阅pytorch-lightningmanual。我认为您正在寻找 training_step_end/validation_step_end（假设您使用的是 DP/DDP2）。

...因此，当 Lightning 调用任何 training_step、validation_step、test_step 时，您将只对其中一个部分进行操作。 (...) 对于大多数指标，这并不重要。但是，如果您想使用所有批处理部分向您的计算图添加一些东西（如 softmax），您可以使用 training_step_end 步骤。

【讨论】：

是的。这正是我在training_step_end 方法而不是training_step 中实现我的逻辑的原因。但是，我在那里只得到部分数据。

你可能在使用 ddp 吗？根据手册，通过 ddp，设计：“每个 GPU 都可以看到整个数据集的一个子集。它只会看到那个子集。”

这适用于validation_step 或training_step 方法。对于*_step_end，方法应该在从每台机器汇总输出后调用。

【参考方案2】：

使用 DDP 后端时，每个 GPU 都会运行一个单独的进程。他们无法访问彼此的数据，但有一些特殊操作（reduce、all_reduce、gather、all_gather）可以使进程同步。当您在张量上使用此类操作时，进程将等待彼此到达同一点并以某种方式组合它们的值，例如从每个进程中获取总和。

理论上，可以从所有进程中收集所有数据，然后在一个进程中计算指标，但这很慢并且容易出现问题，因此您希望尽量减少传输的数据。最简单的方法是分段计算指标，然后取平均值。当您使用 sync_dist=True 时，self.log() 调用将 do this automatically。

如果您不想取 GPU 进程的平均值，也可以在每一步更新一些状态变量，然后在 epoch 同步状态变量并根据这些值计算您的指标。推荐的方法是创建一个使用 Metrics API 的类，该类最近从 PyTorch Lightning 转移到了 TorchMetrics 项目。

如果存储一组状态变量还不够，您可以尝试让您的指标收集来自所有进程的所有数据。从Metric 基类派生您自己的指标，覆盖update() 和compute() 方法。使用add_state("data", default=[], dist_reduce_fx="cat") 创建一个列表，您可以在其中收集计算指标所需的数据。 dist_reduce_fx="cat" 将导致来自不同进程的数据与torch.cat() 合并。它在内部使用torch.distributed.all_gather。这里棘手的部分是它假设所有进程都创建相同大小的张量。如果大小不匹配，同步将无限期挂起。