(self, stage=None)
| 115 | all_urls, [num_train, num_test, num_val]) # , generator=torch.Generator().manual_seed(self.seed) |
| 116 | |
| 117 | def setup(self, stage=None): |
| 118 | if 'train' in self.datasets: |
| 119 | self.datasets['train'] = ImageEmbeddingDataset( |
| 120 | self.train_urls, |
| 121 | self.tokenizer_t5, |
| 122 | self.tokenizer_t5_en, |
| 123 | self.tokenizer_en, |
| 124 | shuffle_shards=self.shuffle_train, |
| 125 | resample=self.resample_train, |
| 126 | hr_size=self.hr_size, |
| 127 | handler=wds.handlers.warn_and_continue, |
| 128 | center_crop=self.center_crop, |
| 129 | size=self.resolution, |
| 130 | ) |
| 131 | |
| 132 | if self.shuffle_num is not None and self.shuffle_num > 0: |
| 133 | self.datasets['train'].shuffle(self.shuffle_num) |
| 134 | |
| 135 | def _train_dataloader(self, preprocess=None, process_group=None, group_rank=None): |
| 136 |
no test coverage detected