Split `reload_dataloaders_every_n_epochs` into separate parameters for train, val and test dataloaders

[windring]

Description & Motivation

Currently, the reload_dataloaders_every_n_epochs parameter applies to all dataloaders (train, validation and test). However, in some use cases, we need different reload behaviors for different types of dataloaders.
Specifically, in my case:

• The train dataloader is an infinite IterableDataset that doesn't need reloading
• The validation and test dataloaders are finite and need to be reloaded after each epoch

This mismatch causes issues when trying to use the current single parameter to control reloading.

Pitch

I propose splitting reload_dataloaders_every_n_epochs into three separate parameters:

1. reload_train_dataloaders_every_n_epochs
2. reload_val_dataloaders_every_n_epochs
3. reload_test_dataloaders_every_n_epochs

This would allow users to have fine-grained control over the reloading behavior for each type of dataloader.

Alternatives

No response

Additional context

Relevant code areas that may need modification:

    def setup_data(self) -> None:
        trainer = self.trainer
        trainer_fn = self._trainer_fn
        >>>if self._combined_loader is not None and trainer_fn == TrainerFn.FITTING and not self._should_reload_val_dl:
            return
    ....

    @property
    def _should_reload_val_dl(self) -> bool:
        """Check if validation dataloader should be reloaded."""
        >>>n_epochs = self.trainer.reload_dataloaders_every_n_epochs
        return bool(n_epochs and self.trainer.current_epoch - self._last_val_dl_reload_epoch >= n_epochs)

    def setup_data(self) -> None:
        >>>if self._combined_loader is not None and not self._should_reload_train_dl:
            return
    ...

    @property
    def _should_reload_train_dl(self) -> bool:
        """Check if train dataloader should be reloaded."""
        >>>n_epochs = self.trainer.reload_dataloaders_every_n_epochs
        return n_epochs and self.trainer.current_epoch - self._last_train_dl_reload_epoch >= n_epochs

cc @Borda

[tshu-w]

You can set reload_dataloaders_every_n_epochs to 1 and reset the dataloader you want as below?

class LitData(L.LightningDataModule):
    ...
    
    def train_dataloader(self) -> TRAIN_DATALOADERS:
        if self.train_dataloader is None:
            self.train_dataloader = DataLoader(...)

        return self.train_dataloader

    def val_dataloader(self) -> EVAL_DATALOADERS:
        return DataLoader(...)

[windring]

good idea, worth trying

[windring]

@tshu-w It works very well, thanks!

You can set reload_dataloaders_every_n_epochs to 1 and reset the dataloader you want as below?

class LitData(L.LightningDataModule):
    ...
    
    def train_dataloader(self) -> TRAIN_DATALOADERS:
        if self.train_dataloader is None:
            self.train_dataloader = DataLoader(...)

        return self.train_dataloader

    def val_dataloader(self) -> EVAL_DATALOADERS:
        return DataLoader(...)