| 118 | self.batchs_in_one_epoch_id = [i for i in range(self.length)] |
| 119 | |
| 120 | def create_batch(self): |
| 121 | batch_list = [] |
| 122 | for ratio in self.indices_rank_i_ratio_unique: |
| 123 | ratio_ids = np.where(self.indices_rank_i_ratio == ratio)[0] |
| 124 | ratio_ids = self.indices_rank_i_ori[ratio_ids] |
| 125 | if self.shuffle: |
| 126 | random.shuffle(ratio_ids) |
| 127 | num_ratio = ratio_ids.shape[0] |
| 128 | if ratio < 5: |
| 129 | batch_size_ratio = self.base_batch_size |
| 130 | else: |
| 131 | batch_size_ratio = min( |
| 132 | self.max_bs, |
| 133 | int( |
| 134 | max(1, (self.base_elements / |
| 135 | (self.base_im_h * ratio * self.base_im_h))))) |
| 136 | if num_ratio > batch_size_ratio: |
| 137 | batch_num_ratio = num_ratio // batch_size_ratio |
| 138 | print(self.rank, num_ratio, ratio * self.base_im_h, |
| 139 | batch_num_ratio, batch_size_ratio) |
| 140 | ratio_ids_full = ratio_ids[:batch_num_ratio * |
| 141 | batch_size_ratio].reshape( |
| 142 | batch_num_ratio, |
| 143 | batch_size_ratio, 1) |
| 144 | w = np.full_like(ratio_ids_full, ratio * self.base_im_h) |
| 145 | h = np.full_like(ratio_ids_full, self.base_im_h) |
| 146 | ra_wh = np.full_like(ratio_ids_full, ratio) |
| 147 | ratio_ids_full = np.concatenate([w, h, ratio_ids_full, ra_wh], |
| 148 | axis=-1) |
| 149 | batch_ratio = ratio_ids_full.tolist() |
| 150 | |
| 151 | if batch_num_ratio * batch_size_ratio < num_ratio: |
| 152 | drop = ratio_ids[batch_num_ratio * batch_size_ratio:] |
| 153 | if self.is_training: |
| 154 | drop_full = ratio_ids[:batch_size_ratio - ( |
| 155 | num_ratio - batch_num_ratio * batch_size_ratio)] |
| 156 | drop = np.append(drop_full, drop) |
| 157 | drop = drop.reshape(-1, 1) |
| 158 | w = np.full_like(drop, ratio * self.base_im_h) |
| 159 | h = np.full_like(drop, self.base_im_h) |
| 160 | ra_wh = np.full_like(drop, ratio) |
| 161 | |
| 162 | drop = np.concatenate([w, h, drop, ra_wh], axis=-1) |
| 163 | |
| 164 | batch_ratio.append(drop.tolist()) |
| 165 | batch_list += batch_ratio |
| 166 | else: |
| 167 | print(self.rank, num_ratio, ratio * self.base_im_h, |
| 168 | batch_size_ratio) |
| 169 | ratio_ids = ratio_ids.reshape(-1, 1) |
| 170 | w = np.full_like(ratio_ids, ratio * self.base_im_h) |
| 171 | h = np.full_like(ratio_ids, self.base_im_h) |
| 172 | ra_wh = np.full_like(ratio_ids, ratio) |
| 173 | |
| 174 | ratio_ids = np.concatenate([w, h, ratio_ids, ra_wh], axis=-1) |
| 175 | batch_list.append(ratio_ids.tolist()) |
| 176 | return batch_list |
| 177 | |