(
self,
train_dataset,
val_dataset_out,
val_dataset_in,
test_dataset_out,
test_dataset_in,
global_batch_size,
num_workers,
num_nodes=1,
num_devices=1,
)
| 5 | |
| 6 | class ImageDataModule(L.LightningDataModule): |
| 7 | def __init__( |
| 8 | self, |
| 9 | train_dataset, |
| 10 | val_dataset_out, |
| 11 | val_dataset_in, |
| 12 | test_dataset_out, |
| 13 | test_dataset_in, |
| 14 | global_batch_size, |
| 15 | num_workers, |
| 16 | num_nodes=1, |
| 17 | num_devices=1, |
| 18 | ): |
| 19 | super().__init__() |
| 20 | self._builders = { |
| 21 | "train": train_dataset, |
| 22 | "val_out": val_dataset_out, |
| 23 | "val_in": val_dataset_in, |
| 24 | "test_out": test_dataset_out, |
| 25 | "test_in": test_dataset_in, |
| 26 | } |
| 27 | self.num_workers = num_workers |
| 28 | self.batch_size = global_batch_size // (num_nodes * num_devices) |
| 29 | print(f"Each GPU will receive {self.batch_size} images") |
| 30 | |
| 31 | @property |
| 32 | def num_classes(self): |
nothing calls this directly
no outgoing calls
no test coverage detected