将 ddp 后端与 PyTorch Lightning 一起使用时验证整个验证集

Question

我正在使用 PyTorch Lightning 训练图像分类模型，并在具有多个 GPU 的机器上运行，因此我使用推荐的分布式后端以获得最佳性能 ddp (DataDistributedParallel)。这自然会拆分数据集，因此每个 GPU 只会看到数据的一部分。

但是，对于验证，我想计算整个验证集的准确度等指标，而不仅仅是部分。我该怎么做？我找到了some hints in the official documentation，但它们没有按预期工作或让我感到困惑。发生的事情是validation_epoch_end 被称为num_gpus 次，每个验证数据为1/num_gpus。我想汇总所有结果，只运行一次validation_epoch_end。

在this section 中，他们声明当使用 dp/ddp2 时，您可以添加一个像这样调用的附加函数

def validation_step(self, batch, batch_idx):
    loss, x, y, y_hat = self.step(batch)
    return {"val_loss": loss, 'y': y, 'y_hat': y_hat}

def validation_step_end(self, self, *args, **kwargs):
    # do something here, I'm not sure what, 
    # as it gets called in ddp directly after validation_step with the exact same values
    return args[0]

但是，结果并未汇总，validation_epoch_end 仍被称为num_gpu 次。这种行为是不是ddp不可用？还有其他方法可以实现这种聚合行为吗？

Answer 1

training_epoch_end() 和 validation_epoch_end() 聚合来自特定过程的所有训练/验证批次的数据。他们只是创建一个列表，列出您在每个训练或验证步骤中返回的内容。

使用 DDP 后端时，每个 GPU 都会运行一个单独的进程。没有简单的方法可以访问另一个进程正在处理的数据，但是有一种机制可以在进程之间同步特定的张量。最简单的方法是分段计算度量，然后例如通过取平均值来同步张量。当您使用 sync_dist=True 时，self.log() 调用将 automatically synchronize the value between GPUs。值的同步方式由reduce_fx 参数决定，默认为torch.mean。

如果您也对批次的平均指标感到满意，则无需覆盖 training_epoch_end() 或 validation_epoch_end() — self.log() 将为您进行平均。

也可以在每一步更新一些状态变量，然后定义一个自定义指标，该指标将在每个 epoch 之后根据您保存的值进行计算。推荐的方法是创建一个派生自Metric 类的类（现在移至TorchMetrics 项目）。使用add_state() 在构造函数中添加状态变量并覆盖update() 和compute() 方法。 API 将负责同步进程之间的状态变量。

TorchMetrics 中已经有一个准确度指标，source code 是如何使用 API 的一个很好的例子。

Answer 2

我想你正在寻找training_step_end/validation_step_end。

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