[RFC] Introduce a dataloader factory class to better manage data modules
🚀 Feature

Create a new class to manage dataloaders
Proposal: https://docs.google.com/document/d/1c0dBmASUfQy0kIpliGD7sGmdzC0sgbuOQStkM02UySM/edit
Motivation

DataModules bundle together training, validation, and testing dataloaders
Often times, we want to configure different dataloader settings for each of these phases

Example: configure larger batch sizes for validation (no gradients so we can use more memory for batches/activations)
Example: handle uneven end of data differently across training vs validation (drop tail in training, wraparound for validation)


This is combination of bundling of different phases along with a variety of knobs creates complex datamodule initialization logic
Furthermore, a generic datamodule is difficult to implement as it can be used for training (where the train dataloader must be defined and optionally the val dataloader) or validation (where only the val dataloader must be defined) or testing (where only the test dataloader must be defined) or prediction (where only the predict dataloader must be defined).

Pitch
from typing import Any, Dict, List, Mapping, Optional
from abc import ABC, abstractmethod

import torch
from pytorch_lightning.core.hooks import CheckpointHooks
import pytorch_lightning as pl
from torch.utils.data import DataLoader

class DataLoaderFactory(ABC, CheckpointHooks):

    def __init__(self) -> None:
        # Pointer to the trainer object
        # Placeholder until we define a proper TrainerContext class (e.g. frozen dataclass)
        # to pass things like progress tracking, rank, or world size to the factory
        self.trainer = None

        # Private attrs to keep track of whether or not data hooks have been called yet
        self._prepare_data_called: bool = False
        self._setup_called: bool = False
        self._teardown_called: bool = False

        # This should not be a trainer setting
        # https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#prepare-data-per-node
        self.prepare_data_per_node: bool = True

    def prepare_data(self) -> None:
        pass

    def setup(self) -> None:
        pass
    
    @abstractmethod
    def get_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
    
    def teardown(self) -> None:
        pass

    @property
    def prepare_data_called(self) -> bool:
        return self._prepare_data_called

    @prepare_data_called.setter
    def prepare_data_called(self, val: bool) -> None:
        self._prepare_data_called = val

    @property
    def setup_called(self) -> bool:
        return self._setup_called

    @setup_called.setter
    def setup_called(self, val: bool) -> None:
        self._setup_called = val

    @property
    def teardown_called(self) -> bool:
        return self._teardown_called

    @teardown_called.setter
    def teardown_called(self, val: bool) -> None:
        self._teardown_called = val
We can add optional attributes inside of the datamodule for these classes, one for each of train/val/test/predict as a convenience, along with a convenience method to instantiate a datamodule from these factories similar to this: 
  
    
      pytorch-lightning/pytorch_lightning/core/datamodule.py
    
    
        Lines 343 to 398
      in
      a72a799
    
  
  
    

        
          
               @classmethod 
        

        
          
               def from_datasets( 
        

        
          
                   cls, 
        

        
          
                   train_dataset: Optional[Union[Dataset, Sequence[Dataset], Mapping[str, Dataset]]] = None, 
        

        
          
                   val_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None, 
        

        
          
                   test_dataset: Optional[Union[Dataset, Sequence[Dataset]]] = None, 
        

        
          
                   batch_size: int = 1, 
        

        
          
                   num_workers: int = 0, 
        

        
          
               ): 
        

        
          
                   r""" 
        

        
          
                   Create an instance from torch.utils.data.Dataset. 
        

        
          
            
        

        
          
                   Args: 
        

        
          
                       train_dataset: (optional) Dataset to be used for train_dataloader() 
        

        
          
                       val_dataset: (optional) Dataset or list of Dataset to be used for val_dataloader() 
        

        
          
                       test_dataset: (optional) Dataset or list of Dataset to be used for test_dataloader() 
        

        
          
                       batch_size: Batch size to use for each dataloader. Default is 1. 
        

        
          
                       num_workers: Number of subprocesses to use for data loading. 0 means that the 
        

        
          
                           data will be loaded in the main process. Number of CPUs available. 
        

        
          
            
        

        
          
                   """ 
        

        
          
            
        

        
          
                   def dataloader(ds, shuffle=False): 
        

        
          
                       return DataLoader( 
        

        
          
                           ds, 
        

        
          
                           batch_size=batch_size, 
        

        
          
                           shuffle=shuffle, 
        

        
          
                           num_workers=num_workers, 
        

        
          
                           pin_memory=True, 
        

        
          
                       ) 
        

        
          
            
        

        
          
                   def train_dataloader(): 
        

        
          
                       if isinstance(train_dataset, Mapping): 
        

        
          
                           return {key: dataloader(ds, shuffle=True) for key, ds in train_dataset.items()} 
        

        
          
                       if isinstance(train_dataset, Sequence): 
        

        
          
                           return [dataloader(ds, shuffle=True) for ds in train_dataset] 
        

        
          
                       return dataloader(train_dataset, shuffle=True) 
        

        
          
            
        

        
          
                   def val_dataloader(): 
        

        
          
                       if isinstance(val_dataset, Sequence): 
        

        
          
                           return [dataloader(ds) for ds in val_dataset] 
        

        
          
                       return dataloader(val_dataset) 
        

        
          
            
        

        
          
                   def test_dataloader(): 
        

        
          
                       if isinstance(test_dataset, Sequence): 
        

        
          
                           return [dataloader(ds) for ds in test_dataset] 
        

        
          
                       return dataloader(test_dataset) 
        

        
          
            
        

        
          
                   datamodule = cls() 
        

        
          
                   if train_dataset is not None: 
        

        
          
                       datamodule.train_dataloader = train_dataloader 
        

        
          
                   if val_dataset is not None: 
        

        
          
                       datamodule.val_dataloader = val_dataloader 
        

        
          
                   if test_dataset is not None: 
        

        
          
                       datamodule.test_dataloader = test_dataloader 
        

        
          
                   return datamodule 
        
    
  


class LightningDataModule(...)
    train_dataloader_factory: Optional[DataLoaderFactory] = None
    val_dataloader_factory: Optional[DataLoaderFactory] = None
    test_dataloader_factory: Optional[DataLoaderFactory] = None
    predict_dataloader_factory: Optional[DataLoaderFactory] = None

   @classmethod
    def from_dataloader_factories(
        cls,
        train_dataloader_factory: Optional[DataLoaderFactory] = None,
        val_dataloader_factory: Optional[DataLoaderFactory]  = None,
        test_dataloader_factory: Optional[DataLoaderFactory] = None,
        predict_dataloader_factory: Optional[DataLoaderFactory] = None
    ):
        datamodule = cls()
        if train_dataloader_factory is not None:
            datamodule.train_dataloader = train_dataloader_factory.get_dataloader
        if val_dataloader_factory is not None:
            datamodule.val_dataloader = val_dataloader_factory.get_dataloader
        if test_dataloader_factory is not None:
            datamodule.test_dataloader = test_dataloader_factory.get_dataloader
        if predict_dataloader_factory is not None:
            datamodule.predict_dataloader = predict_dataloader_factory.get_dataloader
        return datamodule
This can also replace the raw dataloaders that are currently accepted on the Trainer.fit()/validate()/test()/predict() APIs - the classes here are an improvement as they can have access to trainer context which the existing dataloaders do not have.
cc @justusschock @awaelchli @carmocca
Alternatives

Additional context