`LightningDataModule` for wrapping custom PlinderDatasets. 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...
| 24 | |
| 25 | |
| 26 | class PlinderDataModule(LightningDataModule): |
| 27 | """`LightningDataModule` for wrapping custom PlinderDatasets. |
| 28 | |
| 29 | A `LightningDataModule` implements 7 key methods: |
| 30 | |
| 31 | ```python |
| 32 | def prepare_data(self): |
| 33 | # Things to do on 1 GPU/TPU (not on every GPU/TPU in DDP). |
| 34 | # Download data, pre-process, split, save to disk, etc... |
| 35 | |
| 36 | def setup(self, stage): |
| 37 | # Things to do on every process in DDP. |
| 38 | # Load data, set variables, etc... |
| 39 | |
| 40 | def train_dataloader(self): |
| 41 | # return train dataloader |
| 42 | |
| 43 | def val_dataloader(self): |
| 44 | # return validation dataloader |
| 45 | |
| 46 | def test_dataloader(self): |
| 47 | # return test dataloader |
| 48 | |
| 49 | def predict_dataloader(self): |
| 50 | # return predict dataloader |
| 51 | |
| 52 | def teardown(self, stage): |
| 53 | # Called on every process in DDP. |
| 54 | # Clean up after fit or test. |
| 55 | ``` |
| 56 | |
| 57 | This allows you to share a full dataset without explaining how to download, |
| 58 | split, transform and process the data. |
| 59 | |
| 60 | Read the docs: |
| 61 | https://lightning.ai/docs/pytorch/latest/data/datamodule.html |
| 62 | """ |
| 63 | |
| 64 | def __init__( |
| 65 | self, |
| 66 | data_dir: str = "data/PLINDER/", |
| 67 | batch_size: int = 16, |
| 68 | num_workers: int = 0, |
| 69 | pin_memory: bool = False, |
| 70 | stage: Optional[str] = None, |
| 71 | plinder_offline: bool = False, |
| 72 | **kwargs: Any, |
| 73 | ) -> None: |
| 74 | """Initialize a `PlinderDataModule`. |
| 75 | |
| 76 | :param data_dir: The data directory. Defaults to `"data/"`. |
| 77 | :param batch_size: The batch size. Defaults to `16`. |
| 78 | :param num_workers: The number of workers. Defaults to `0`. |
| 79 | :param pin_memory: Whether to pin memory. Defaults to `False`. |
| 80 | :param plinder_offline: Whether to use the offline version of PLINDER. Defaults to `False`. |
| 81 | :param stage: The stage to setup. Either `"fit"`, `"validate"`, `"test"`, or `"predict"`. Defaults to ``None``. |
| 82 | """ |
| 83 | super().__init__() |