| 342 | # --------------------------------------------------------------------------------------------------------------------- |
| 343 | |
| 344 | class BilliardSimDataModule: |
| 345 | |
| 346 | def __init__( |
| 347 | self, |
| 348 | batch_size: int, |
| 349 | train=None, |
| 350 | validation=None, |
| 351 | num_workers=4, |
| 352 | val_batch_size: int | None = None, |
| 353 | val_num_workers: int | None = None, |
| 354 | pre_collation_factor: int = 1, |
| 355 | ): |
| 356 | super().__init__() |
| 357 | self.batch_size = batch_size |
| 358 | self.num_workers = num_workers |
| 359 | self.train = train |
| 360 | self.validation = validation |
| 361 | self.val_batch_size = val_batch_size if val_batch_size is not None else batch_size |
| 362 | self.val_num_workers = val_num_workers if val_num_workers is not None else num_workers |
| 363 | self.pre_collation_factor = pre_collation_factor |
| 364 | if self.pre_collation_factor > 1: |
| 365 | assert self.batch_size % self.pre_collation_factor == 0 |
| 366 | |
| 367 | def make_loader( |
| 368 | self, |
| 369 | batch_size: int, |
| 370 | num_workers: int, |
| 371 | dataset_config: dict[str, Any], |
| 372 | worker_init: bool = True, |
| 373 | pre_collation_factor: int = 1, |
| 374 | ): |
| 375 | loader = torch.utils.data.DataLoader( |
| 376 | BilliardSimDataset(**dataset_config), |
| 377 | batch_size=batch_size // pre_collation_factor, |
| 378 | num_workers=num_workers, |
| 379 | worker_init_fn=worker_init_fn if worker_init else None, |
| 380 | collate_fn=dict_collation_fn, |
| 381 | pin_memory=True, |
| 382 | ) |
| 383 | if pre_collation_factor > 1: |
| 384 | loader = CollatingDataLoader(loader, pre_collation_factor, dict_join_collation_fn) |
| 385 | return loader |
| 386 | |
| 387 | def train_dataloader(self): |
| 388 | return self.make_loader( |
| 389 | **self.train, |
| 390 | batch_size=self.batch_size, |
| 391 | num_workers=self.num_workers, |
| 392 | pre_collation_factor=self.pre_collation_factor, |
| 393 | ) |
| 394 | |
| 395 | def val_dataloader(self): |
| 396 | return self.make_loader( |
| 397 | **self.validation, |
| 398 | batch_size=self.val_batch_size, |
| 399 | num_workers=self.val_num_workers |
| 400 | ) |
no outgoing calls
no test coverage detected