Proposal: Initialize DataModule from Dataset(s)
🚀 Feature
Add functionality to initialize a pl.DataModule from pre-existing datasets, which could help eliminate some of the boilerplate caused from initializing torch.utils.data.DataLoaders.
Motivation
A common pattern we have seen some users using is creating a DataModule wrapper that instantiates the necessary dataloaders given a train dataset, and optional validation and test datasets (this is my own wrapper):
class DataModule(pl.LightningDataModule):
  def __init__(
        self,
        train_ds: Optional[Dataset] = None,
        valid_ds: Optional[Dataset] = None,
        test_ds: Optional[Dataset] = None,
        batch_size: int = 1,
        num_workers: Optional[int] = None,
    ):

        super().__init__()
        self._train_ds = train_ds
        self._valid_ds = valid_ds
        self._test_ds = test_ds

        if self._train_ds is not None:
            self.train_dataloader = self._train_dataloader

        if self._valid_ds is not None:
            self.val_dataloader = self._val_dataloader

        if self._test_ds is not None:
            self.test_dataloader = self._test_dataloader

        self.batch_size = batch_size

        if num_workers is None:
            num_workers = os.cpu_count()

        if num_workers is None:
            warnings.warn("could not infer cpu count automatically, setting it to zero")
            num_workers = 0
        self.num_workers = num_workers

    def _train_dataloader(self):
        return DataLoader(
            self._train_ds,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True,
        )

    def _val_dataloader(self):
        return DataLoader(
            self._valid_ds,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )

    def _test_dataloader(self):
        return DataLoader(
            self._test_ds,
            batch_size=self.batch_size,
            num_workers=self.num_workers,
            pin_memory=True,
        )
Pitch
Add optional arguments for train, val, test datasets, batch size and num_workers to the base DataModule class:
train_ds = MyDataset(...)
val_ds = MyDataset(...)
data = pl.DataModule(train_ds=train_ds, val_ds=val_ds, batch_size=64, num_workers=8)
Or, have a factory method doing the same thing:
train_ds = MyDataset(...)
val_ds = MyDataset(...)
data = pl.DataModule.from_datasets(train_ds=train_ds, val_ds=val_ds, batch_size=64, num_workers=8)
Would like to here from the community if they have any ideas of preferences here!