Best practice for multiple validation datasets

[ginfantem]

Hello everyone, I have multiple validation datasets that I pass to the trainer as an OrderedDict. At each validation epoch I might want to process one of them exclusively, in an interleaving fashion. However, even I can conditionally process each dataset via the dataloader_idx argument, validation_step is run also for all the batches for the rest of dataloaders (even though inside validation_step the flow discards them due to the if statement). Is there a best practice around this, so this does not slow down the training process?

[stef41]

You can use Lightning's CombinedLoader with different modes to control how validation dataloaders are iterated, and override on_validation_epoch_start to select which ones run.

Approach 1: Use a callback to swap val dataloaders per epoch

class RotatingValCallback(pl.Callback):
    def on_validation_epoch_start(self, trainer, pl_module):
        epoch = trainer.current_epoch
        # Rotate through your val datasets
        active_idx = epoch % len(trainer.val_dataloaders)
        # Store which one is active so validation_step can check
        pl_module._active_val_idx = active_idx

Then in your validation_step, skip inactive datasets efficiently:

def validation_step(self, batch, batch_idx, dataloader_idx=0):
    if dataloader_idx != self._active_val_idx:
        return None  # Skip — Lightning handles None returns gracefully
    # ... actual validation logic ...

This still iterates all dataloaders but the compute inside validation_step is skipped.

Approach 2: Dynamically set val_dataloaders (avoids iterating skipped ones)

Override val_dataloader() to return only the active one:

class MyModel(pl.LightningModule):
    def __init__(self, val_datasets):
        super().__init__()
        self.val_datasets = val_datasets  # OrderedDict

    def val_dataloader(self):
        # Return only the dataset for this epoch
        idx = self.current_epoch % len(self.val_datasets)
        key = list(self.val_datasets.keys())[idx]
        return DataLoader(self.val_datasets[key], ...)

    def on_validation_epoch_start(self):
        idx = self.current_epoch % len(self.val_datasets)
        key = list(self.val_datasets.keys())[idx]
        self.log("val_dataset", float(idx))  # track which one

This is more efficient because Lightning only iterates the single active dataloader.

Approach 3: Use val_check_interval with manual control

If you want dataset A every epoch and dataset B every N epochs:

# In your training loop / callback:
def on_train_epoch_end(self, trainer, pl_module):
    if trainer.current_epoch % 5 == 0:
        # Run expensive validation manually
        trainer.validate(pl_module, dataloaders=expensive_val_dl)

Recommendation: Approach 2 is the cleanest — it avoids iterating unused dataloaders entirely and doesn't require dataloader_idx branching.