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

Question

I'm training an image classification model with PyTorch Lightning and running on a machine with more than one GPU, so I use the recommended distributed backend for best performance ddp (DataDistributedParallel). This naturally splits up the dataset, so each GPU will only ever see one part of the data.

However, for validation, I would like to compute metrics like accuracy on the entire validation set and not just on a part. How would I do that? I found some hints in the official documentation, but they do not work as expected or are confusing to me. What's happening is that validation_epoch_end is called num_gpus times with 1/num_gpus of the validation data each. I would like to aggregate all results and only run the validation_epoch_end once.

In this section they state that when using dp/ddp2 you can add an additional function called like this

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]

However, the results are not being aggregated and validation_epoch_end is still called num_gpu times. Is this kind of behavior not available for ddp? Is there some other way how achieve this aggregation behavior?

Answer 1

training_epoch_end() and validation_epoch_end() receive data that is aggregated from all training / validation batches of the particular process. They simply receive a list of what you returned in each training or validation step.

When using the DDP backend, there's a separate process running for every GPU. There's no simple way to access the data that another process is processing, but there's a mechanism for synchronizing a particular tensor between the processes.

The easiest approach for computing a metric on the entire validation set is to calculate the metric in pieces and then synchronize the resulting tensor, for example by taking the average. self.log() calls will automatically synchronize the value between GPUs when you use sync_dist=True. How the value is synchronized is determined by the reduce_fx argument, which by default is torch.mean.

If you're happy with averaging the metric over batches too, you don't need to override training_epoch_end() or validation_epoch_end() — self.log() will do the averaging for you.

If the metric cannot be calculated separately for each GPU and then averaged, it can get a bit more challenging. It's possible to update some state variables at each step, and then synchronize the state variables at the end of an epoch and calculate the metric. The recommended way is to create a class that derives from the Metric class from the TorchMetrics project. Add the state variables in the constructor using add_state() and override the update() and compute() methods. The API will take care of synchronizing the state variables between the GPU processes.

There's already an accuracy metric in TorchMetrics and the source code is a good example of how to use the API.

Answer 2

I think you are looking for 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.