TorchMetrics, Pytorch Lightning and DataParallel

[aretor]

In the TorchMetrics in Pytorch Lightning section there is the following warning:

If using metrics in data parallel mode (dp), the metric update/logging should be done in the _step_end method (where is either training, validation or test). This is due to metric states else being destroyed after each forward pass, leading to wrong accumulation.

Is this true also in the case of distributed data parallel mode (ddp/ddp2)?

My code reports correct metric values only if I follow the instructions above.

[SkafteNicki]

No, in ddp it should not be necessary. You still need to sync the metric at some point between the different devices, which is automatically done when metric.compute() is called. Therefore, something like this should still work:

def training_step(self, batch, batch_idx):
    ...
    self.metric.update(preds, target)
    ...

def training_epoch_end(self, outputs)
    val = self.metric.compute()  # this will sync the metric between devices
    self.log("metric", val)
    self.metric.reset()