(self)
| 413 | self.datasets[k] = WrappedDataset(self.datasets[k]) |
| 414 | |
| 415 | def _train_dataloader(self): |
| 416 | is_iterable_dataset = isinstance(self.datasets['train'], Txt2ImgIterableBaseDataset) |
| 417 | if is_iterable_dataset or self.use_worker_init_fn: |
| 418 | init_fn = worker_init_fn |
| 419 | else: |
| 420 | init_fn = None |
| 421 | if "train2" in self.dataset_configs and self.dataset_configs["train2"]['params']["caption"] != '': |
| 422 | train_set = self.datasets["train"] |
| 423 | train2_set = self.datasets["train2"] |
| 424 | concat_dataset = ConcatDataset(train_set, train2_set) |
| 425 | return DataLoader(concat_dataset, batch_size=self.batch_size // 2, |
| 426 | num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True, |
| 427 | worker_init_fn=init_fn) |
| 428 | else: |
| 429 | return DataLoader(self.datasets["train"], batch_size=self.batch_size, |
| 430 | num_workers=self.num_workers, shuffle=False if is_iterable_dataset else True, |
| 431 | worker_init_fn=init_fn) |
| 432 | |
| 433 | def _val_dataloader(self, shuffle=False): |
| 434 | if isinstance(self.datasets['validation'], Txt2ImgIterableBaseDataset) or self.use_worker_init_fn: |
nothing calls this directly
no test coverage detected