Need a strategy to gather and compute validation loss on the whole validation dataset

[JohnHerry]

Description & Motivation

After lightning 2.0, the validation_epoch_end(self, outputs): function is deprecated. instead, we are asked to use a so called on_validation_epoch_end(self) callback, but this call back takes no parameters. so we should save the validation outpus, into a self.validataion_outpus params inside the lightning_module,
but in the DDP traning, the self.validation_outpus can only save part of the validataion mini-batch results on current rank. I want to get ALL validation results some where, but there is no such convinience.

Pitch

example code:

class My(pl.LightningModule):
    def __init__(self, ...):
       ...
       self.all_validation_step_outptus = []

   def validation_step(self, batch, *args, **kwargs):
       with torch.no_grad():
           gt_data, infer_inputs = batch
          gen_feature = self.backbone(infer_inputs)
          gen_data = self.decoder(gen_feature)
          batch_output = {
              "gt_data": gt_data,
              "gen_feature": gen_feature,
              "gen_data": gen_data
          }
          self.all_validation_step_outpus.append(batch_output)

    def on_validation_epoch_end(self):
         outputs = self.all_validation_step_outputs()  # Here I want ALL validataion dataset results, but can ONLY get local rank results!
        gt_data = [x['gt_data'] for x in outputs]
        gen_data = [x['gen_data'] for x in outputs]
        if self.global_rank == 0:
            val_loss = compute_val_loss(gt_data, gen_data).mean()
            self.log("val_loss": val_loss)
           for  idx, item in outputs:
               self.logger.add_image(f"val/feature_{idx}", plot_spectorgram_to_img(item["gen_feature"]))
           

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[tvelovraf]

Hi @JohnHerry!
I'm not sure it's exactly what you need, but you could try smth like:

class My(pl.LightningModule):
    def __init__(self, ...):
       ...
       # self.all_validation_step_outptus = []

   def validation_step(self, batch, *args, **kwargs):
       with torch.no_grad():
           gt_data, infer_inputs = batch
          gen_feature = self.backbone(infer_inputs)
          gen_data = self.decoder(gen_feature)
          val_loss = compute_val_loss(gt_data, gen_data).mean()
          self.log("val_loss": val_loss, logger=True, on_step=False, on_epoch=True, sync_dist=True)
          ...

sync_dist=True is used to average val_loss across all devices.
And on_epoch=True is used to average val_loss results from all batches at the end of epoch.

But using self.logger.add_image() in that way may be more diffcult

[JohnHerry]

Hi @JohnHerry! I'm not sure it's exactly what you need, but you could try smth like:

class My(pl.LightningModule):
def init(self, ...):
...
# self.all_validation_step_outptus = []

def validation_step(self, batch, *args, **kwargs):
with torch.no_grad():
gt_data, infer_inputs = batch
gen_feature = self.backbone(infer_inputs)
gen_data = self.decoder(gen_feature)
val_loss = compute_val_loss(gt_data, gen_data).mean()
self.log("val_loss": val_loss, logger=True, on_step=False, on_epoch=True, sync_dist=True)
...
sync_dist=True is used to average val_loss across all devices. And on_epoch=True is used to average val_loss results from all batches at the end of epoch.

But using self.logger.add_image() in that way may be more diffcult

Thank you very much, that is the case. the sync_dist config can compute average loss across devices, but it is difficute to gather all validation outputs.

[SkafteNicki]

Thanks for raising this issue. Happy to report that there is a build in way in Lightning to tgather all the validation outputs, which is the self.all_gather method of LightningModule
https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#all-gather

A basic example:

import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import DataLoader, TensorDataset


class SimpleModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
        self.validation_step_outputs = []
    
    def forward(self, x):
        return self.net(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = nn.functional.mse_loss(self(x).squeeze(), y)
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x).squeeze()
        loss = nn.functional.mse_loss(y_hat, y)
        self.validation_step_outputs.append({'predictions': y_hat, 'targets': y})
        return loss
    
    def on_validation_epoch_end(self):
        if not self.validation_step_outputs:
            return
        
        # Concatenate local validation outputs
        local_predictions = torch.cat([x['predictions'] for x in self.validation_step_outputs])
        local_targets = torch.cat([x['targets'] for x in self.validation_step_outputs])
        
        all_predictions = self.all_gather(local_predictions)
        all_targets = self.all_gather(local_targets)
        
        # Compute global metrics on complete dataset
        global_loss = nn.functional.mse_loss(all_predictions, all_targets)
        global_mae = torch.abs(all_predictions - all_targets).mean()
        
        if self.global_rank == 0:
            print(f"Total samples: {len(all_predictions)}")
            print(f"Global Loss: {global_loss:.4f}, MAE: {global_mae:.4f}")
        
        self.log('val_loss', global_loss, sync_dist=False)
        self.log('val_mae', global_mae, sync_dist=False)
        self.validation_step_outputs.clear()
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)


def main():
    # Create dummy data
    X = torch.randn(1000, 10)
    y = X.sum(dim=1) + 0.1 * torch.randn(1000)
    train_data = TensorDataset(X[:800], y[:800])
    val_data = TensorDataset(X[800:], y[800:])
    
    train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_data, batch_size=16, shuffle=False)
    
    model = SimpleModel()
    trainer = pl.Trainer(max_epochs=2, devices=2)
    
    trainer.fit(model, train_loader, val_loader)
    

if __name__ == "__main__":
    main()

Closing issue, but feel free to ping and reopen if necessesary.