| 56 | |
| 57 | |
| 58 | class DataModule(LightningDataModule): |
| 59 | dataset_cfgs: list[DatasetCfgWrapper] |
| 60 | data_loader_cfg: DataLoaderCfg |
| 61 | step_tracker: StepTracker | None |
| 62 | dataset_shim: DatasetShim |
| 63 | global_rank: int |
| 64 | |
| 65 | def __init__( |
| 66 | self, |
| 67 | dataset_cfgs: list[DatasetCfgWrapper], |
| 68 | data_loader_cfg: DataLoaderCfg, |
| 69 | step_tracker: StepTracker | None = None, |
| 70 | dataset_shim: DatasetShim = lambda dataset, _: dataset, |
| 71 | global_rank: int = 0, |
| 72 | ) -> None: |
| 73 | super().__init__() |
| 74 | self.dataset_cfgs = dataset_cfgs |
| 75 | self.data_loader_cfg = data_loader_cfg |
| 76 | self.step_tracker = step_tracker |
| 77 | self.dataset_shim = dataset_shim |
| 78 | self.global_rank = global_rank |
| 79 | |
| 80 | def get_persistent(self, loader_cfg: DataLoaderStageCfg) -> bool | None: |
| 81 | return None if loader_cfg.num_workers == 0 else loader_cfg.persistent_workers |
| 82 | |
| 83 | def get_generator(self, loader_cfg: DataLoaderStageCfg) -> torch.Generator | None: |
| 84 | if loader_cfg.seed is None: |
| 85 | return None |
| 86 | generator = Generator() |
| 87 | generator.manual_seed(loader_cfg.seed + self.global_rank) |
| 88 | return generator |
| 89 | |
| 90 | def train_dataloader(self): |
| 91 | datasets = get_dataset(self.dataset_cfgs, "train", self.step_tracker) |
| 92 | data_loaders = [] |
| 93 | for dataset in datasets: |
| 94 | dataset = self.dataset_shim(dataset, "train") |
| 95 | data_loaders.append( |
| 96 | DataLoader( |
| 97 | dataset, |
| 98 | self.data_loader_cfg.train.batch_size, |
| 99 | shuffle=not isinstance(dataset, IterableDataset), |
| 100 | num_workers=self.data_loader_cfg.train.num_workers, |
| 101 | generator=self.get_generator(self.data_loader_cfg.train), |
| 102 | worker_init_fn=worker_init_fn, |
| 103 | persistent_workers=self.get_persistent(self.data_loader_cfg.train), |
| 104 | ) |
| 105 | ) |
| 106 | return data_loaders if len(data_loaders) > 1 else data_loaders[0] |
| 107 | |
| 108 | def val_dataloader(self): |
| 109 | datasets = get_dataset(self.dataset_cfgs, "val", self.step_tracker) |
| 110 | data_loaders = [] |
| 111 | for dataset in datasets: |
| 112 | dataset = self.dataset_shim(dataset, "val") |
| 113 | data_loaders.append( |
| 114 | DataLoader( |
| 115 | ValidationWrapper(dataset, 1), |