| """Data module composing the data loading pipeline.""" |
|
|
| from __future__ import annotations |
|
|
| import lightning.pytorch as pl |
| from torch.utils.data import DataLoader |
|
|
| from vis4d.config import instantiate_classes |
| from vis4d.config.typing import DataConfig |
| from vis4d.data.typing import DictData |
|
|
|
|
| class DataModule(pl.LightningDataModule): |
| """DataModule that wraps around the vis4d implementations. |
| |
| This is a wrapper around the vis4d implementations that allows to use |
| pytorch-lightning for training and testing. |
| """ |
|
|
| def __init__(self, data_cfg: DataConfig) -> None: |
| """Creates an instance of the class.""" |
| super().__init__() |
| self.data_cfg = data_cfg |
|
|
| def train_dataloader(self) -> DataLoader[DictData]: |
| """Return dataloader for training.""" |
| if self.trainer is not None and hasattr(self.trainer, "seed"): |
| seed = self.trainer.seed |
| else: |
| seed = None |
| return instantiate_classes(self.data_cfg.train_dataloader, seed=seed) |
|
|
| def test_dataloader(self) -> list[DataLoader[DictData]]: |
| """Return dataloaders for testing.""" |
| return instantiate_classes(self.data_cfg.test_dataloader) |
|
|
| def val_dataloader(self) -> list[DataLoader[DictData]]: |
| """Return dataloaders for validation.""" |
| return self.test_dataloader() |
|
|
