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

Validate on entire validation set when using ddp backend with PyTorch Lightning

我正在使用 PyTorch Lightning 和 运行ning 在具有多个 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吗?有没有其他方法可以实现这种聚合行为?

我认为您正在寻找 training_step_end/validation_step_end

...So, when Lightning calls any of the training_step, validation_step, test_step you will only be operating on one of those pieces. (...) For most metrics, this doesn’t really matter. However, if you want to add something to your computational graph (like softmax) using all batch parts you can use the training_step_end step.

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() 将为您进行平均。

如果不能为每个 GPU 单独计算指标然后取平均值,它可能会变得更具挑战性。可以在每一步更新一些状态变量,然后在一个纪元结束时同步状态变量并计算度量。推荐的方法是创建一个 class 派生自 TorchMetrics 项目的 Metric class。使用 add_state() 在构造函数中添加状态变量并覆盖 update()compute() 方法。 API 将负责同步 GPU 进程之间的状态变量。

TorchMetrics 中已有准确度指标,source code 是如何使用 API.

的一个很好的示例