File size: 1,341 Bytes
f71ac1d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """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()
|