"""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
[docs]
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.
"""
[docs]
def __init__(self, data_cfg: DataConfig) -> None:
"""Creates an instance of the class."""
super().__init__()
self.data_cfg = data_cfg
[docs]
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)
[docs]
def test_dataloader(self) -> list[DataLoader[DictData]]:
"""Return dataloaders for testing."""
return instantiate_classes(self.data_cfg.test_dataloader)
[docs]
def val_dataloader(self) -> list[DataLoader[DictData]]:
"""Return dataloaders for validation."""
return self.test_dataloader()