| 131 | yield self.batchs_in_one_epoch[batch_tuple_id] |
| 132 | |
| 133 | def iter(self): |
| 134 | if self.shuffle: |
| 135 | if self.seed is not None: |
| 136 | random.seed(self.seed) |
| 137 | else: |
| 138 | random.seed(self.epoch) |
| 139 | if not self.ds_width: |
| 140 | random.shuffle(self.img_indices) |
| 141 | random.shuffle(self.img_batch_pairs) |
| 142 | indices_rank_i = self.img_indices[ |
| 143 | self.rank:len(self.img_indices):self.num_replicas] |
| 144 | else: |
| 145 | indices_rank_i = self.img_indices[ |
| 146 | self.rank:len(self.img_indices):self.num_replicas] |
| 147 | |
| 148 | start_index = 0 |
| 149 | batchs_in_one_epoch = [] |
| 150 | for batch_tuple in self.batch_list: |
| 151 | curr_w, curr_h, curr_bsz = batch_tuple |
| 152 | end_index = min(start_index + curr_bsz, self.n_samples_per_replica) |
| 153 | batch_ids = indices_rank_i[start_index:end_index] |
| 154 | n_batch_samples = len(batch_ids) |
| 155 | if n_batch_samples != curr_bsz: |
| 156 | batch_ids += indices_rank_i[:(curr_bsz - n_batch_samples)] |
| 157 | start_index += curr_bsz |
| 158 | |
| 159 | if len(batch_ids) > 0: |
| 160 | if self.ds_width: |
| 161 | wh_ratio_current = self.wh_ratio[ |
| 162 | self.wh_ratio_sort[batch_ids]] |
| 163 | ratio_current = wh_ratio_current.mean() |
| 164 | ratio_current = ratio_current if ratio_current * curr_h < self.max_w else self.max_w / curr_h |
| 165 | else: |
| 166 | ratio_current = None |
| 167 | batch = [(curr_w, curr_h, b_id, ratio_current) |
| 168 | for b_id in batch_ids] |
| 169 | # yield batch |
| 170 | batchs_in_one_epoch.append(batch) |
| 171 | return batchs_in_one_epoch |
| 172 | |
| 173 | def set_epoch(self, epoch: int): |
| 174 | self.epoch = epoch |