Properly support Iterables as an alternative to DataLoaders in Trainer.{fit, validate, ...}

[awaelchli]

🚀 Feature

1. Support passing dataloaders of type Iterable to Trainer.{fit, validate, test, predict}
2. Support returning dataloaders of type Iterable in LightningModule.*_dataloder() and LightningDataModule.*_dataloder()

Motivation

Lightning today supports iterating over a DataLoader (when creating an iterator from it) or multiple dataloaders. Some claims have been made in the past that the Trainer also supports Iterable, but this is not true today because a trivial example can demonstrate that it fails. In any case, our Trainer methods as well as the Lightning{Data}Module methods do not reflect this in their signature types.
Some applications require operating directly on an instance of Iterable and wrapping them inside a torch DataLoader/Dataset would not be desirable or feasible. An example from the medical domain: https://github.com/MIC-DKFZ/batchgenerators

Furthermore, the introduction of Iterables would most likely facilitate the integration of torch/data DataPipes which are soon to be included in PyTorch.

Pitch

Support the Iterable type.

Usage example:

trainer = Trainer()
model = Model()
data = [batch0, batch1, batch2]  # a list of batches
trainer.fit(model, train_dataloaders=data)

Several places inside data_loading.py would have to be updated with branching logic that excludes setup steps normally done for DataLoaders (see limitations section below).

Limitations and Challenges

Some features in Lightning will not work for Iterables. Most notably:

• Automatic replacement of samplers for distributed data-parallel: Samplers do not apply here for Iterable datastructures. However, the user will still have to take care of correctly returning the data in case of distributed training.
• Fault-tolerant training: Not supported out of the box. The user may have to switch to manual-mode fault-tolerant training Add support for manual fault tolerant #10568
• Ambiguity of the input type for Trainer.fit(): The .fit() method and it's sisters currently support these types:
  https://github.com/PyTorchLightning/pytorch-lightning/blob/48cf1adfd3ad9c7e659083a4afc334dafb331f28/pytorch_lightning/utilities/types.py#L36-L43
  However, if we now add Iterable, we run into an ambiguous case because Dict and Sequence are subtypes of Iterable.

Alternatives

• Do not support Iterable. The user will have to wrap their Iterables into IterableDataset and compose them into a DataLoader.
• Add support only in Lite. Drawback: Users converting to Lightning will have questions.

Additional context

#10279 attempted to add support for Iterable types in Lite, but it was reverted later on due to incompleteness.

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @justusschock @awaelchli @ninginthecloud @tchaton @rohitgr7

[justusschock]

+1 for this.

Regarding the limitations:
Back when there was still support for that, stuff like patching the sampler has always been done only for DataLoaders or subclasses. Using other iterables required to manually take care of such things and I think that this is the only way we can actually support this given the arbitrary implementations of such iterables.

Same goes for the fault tolerance IMO.

Regarding the ambiguity of the types for the loader:
One Option would be to only allow support for a single loader (which could then be a combined loader).
While this would be more explicit (and also allow us to remove some arguments from the trainers init), it would be a step back in terms of UX. For the dicts you could explicitly check for mappings and for the sequences maybe investigate the type of the first element?

[tchaton]

Hey @awaelchli,

Did you find a demand for such a feature from the community?

[justusschock]

@tchaton we once had support for it and also I mentioned that I know a few people not using PyTorch dataloaders (sometimes including myself). Not sure when the regression happened since we don't have tests for it, but I think this is a pretty strong restriction to make.

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[awaelchli]

@Stale you blind??? This had a MILESTONE!!!

[alessiamarcolini]

Hey maintainers :) I would be really looking forward to this feature 🚀