Support all iterator modes for fit/validate/test/predict

[carmocca]

Description & Motivation

trainer.fit only works with CombinedLoader(..., mode="max_size_cycle"|"min_size")

trainer.{validate,test,predict} only works with CombinedLoader(..., mode="sequential")

This constraint is checked in the top-level loops:
https://github.com/Lightning-AI/lightning/blob/0009cde1db1a9ab4e2f1e0a9f69a4affb59d5134/src/lightning/pytorch/loops/fit_loop.py#L351-L354
https://github.com/Lightning-AI/lightning/blob/0009cde1db1a9ab4e2f1e0a9f69a4affb59d5134/src/lightning/pytorch/loops/evaluation_loop.py#L182-L183

Pitch

Have all trainer functions support all modes

TODO:

• FitLoop
• EvaluationLoop (Support all CombinedLoader modes during evaluation #17163)
• PredictionLoop

Alternatives

Not do it

Additional context

This builds on top of #16726

cc @Borda @justusschock @awaelchli

[mees]

I am migrating my code to PL 2 and it seems that for the val dataloader getting a batch to be of the form {"key_a": batch_dataloader_a, "key_b": batch_dataloader_b} is not implemented in PL 2 yet. Here my old code as a reference:

def val_dataloader(self):
    val_dataloaders = {
        key: DataLoader(
            dataset,
            batch_size=dataset.batch_size,
            shuffle=False,
            num_workers=dataset.num_workers,
            pin_memory=False,
        )
        for key, dataset in self.val_datasets.items()
    }
    combined_val_loaders = CombinedLoader(val_dataloaders, "max_size_cycle")
    return combined_val_loaders

[carmocca]

@mees I added support for that in #17163, if you want to give it a try. The PR only implements it for validation and testing.

[bkmi]

@mees I added support for that in #17163, if you want to give it a try. The PR only implements it for validation and testing.

really helpful! I hope this gets into "stable" soon.... or even the next release!

[FarzanT]

I really wish there was sequential support in the training loop. Right now, it's not clear how one should handle batches of potentially different sizes in the training_step. We'd have to collate inside the training_step and ensure the given batch size is divided by the number of data loaders to keep gradient accumulation consistent etc. It gets pretty ugly.
@carmocca Thank you for your work on this issue. Not to rush you, but any update on the sequential support in the training loop?
Thanks again!

[carmocca]

Unfortunately, I dont have bandwidth to work on this now. If somebody wants to try, I can help getting the PR merged.
You can follow the structure in the EvaluationLoop. The training hooks will need an optional dataloader_idx argument

[surya-narayanan]

@mees I added support for that in #17163, if you want to give it a try. The PR only implements it for validation and testing.

really helpful! I hope this gets into "stable" soon.... or even the next release!

Me too! Is there any release timeline / nightly version with this supported? I can't use lightning without this and really would like to leverage its other features!

[spfrommer]

Ditto! FYI for others pulling nightly will get the feature: #17163

[chenhaomingbob]

Thanks! I also need this great feature.

[johnathanchiu]

+1, please release this feature asap!

[lukas-folle-snkeos]

Is this feature currently worked on?

[carmocca]

As far as I know, nobody is currently working on it, Lukas

[astirn]

I would also really like this feature. I use CombinedDataloader to encapsulate modality-specific Dataloaders to recycle modalities with fewer data than our largest modality. For this reason, I use CombinedDatloader with "max_size_cycle"/"min_cycle" for train/validation and would like to be able for predict as well. Thanks for considering it!