Class CombinedDataModule

flowdock/data/combined_datamodule.py:28–625

`LightningDataModule` for combined training and test datasets. A `LightningDataModule` implements 7 key methods: ```python def prepare_data(self): # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP). # Download data, pre-process, split, save to disk, etc...

class CombinedDataModule(LightningDataModule):
    """`LightningDataModule` for combined training and test datasets.

    A `LightningDataModule` implements 7 key methods:

    ```python
        def prepare_data(self):
        # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP).
        # Download data, pre-process, split, save to disk, etc...

        def setup(self, stage):
        # Things to do on every process in DDP.
        # Load data, set variables, etc...

        def train_dataloader(self):
        # return train dataloader

        def val_dataloader(self):
        # return validation dataloader

        def test_dataloader(self):
        # return test dataloader

        def predict_dataloader(self):
        # return predict dataloader

        def teardown(self, stage):
        # Called on every process in DDP.
        # Clean up after fit or test.
    ```

    This allows you to share a full dataset without explaining how to download,
    split, transform and process the data.

    Read the docs:
        https://lightning.ai/docs/pytorch/latest/data/datamodule.html
    """

    def __init__(
        self,
        data_dir: str = "data/",
        train_datasets: List[str] = ["pdbbind", "moad", "pdbsidechain"],
        test_datasets: List[str] = ["pdbbind", "moad"],
        batch_size: int = 16,
        num_workers: int = 0,
        pin_memory: bool = False,
        stage: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Initialize a `CombinedDataModule`.

        :param data_dir: The data directory. Defaults to `"data/"`.
        :param train_datasets: The datasets to use for training. Defaults to `["pdbbind", "moad", "pdbsidechain"]`.
        :param test_datasets: The datasets to use for testing. Defaults to `["pdbbind", "moad"]`.
        :param batch_size: The batch size. Defaults to `16`.
        :param num_workers: The number of workers. Defaults to `0`.
        :param pin_memory: Whether to pin memory. Defaults to `False`.
        :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to ``None``.