| 133 | self.bucket_idx_to_data_idxs[self.bucket_idxs[i]].append(i) |
| 134 | |
| 135 | def _batching(self): |
| 136 | # working in local index (0:N) |
| 137 | |
| 138 | # shuffle the data within each bucket separately |
| 139 | if self.shuffle: |
| 140 | for bid in self.bucket_idx_to_data_idxs.keys(): |
| 141 | np.random.shuffle(self.bucket_idx_to_data_idxs[bid]) |
| 142 | |
| 143 | # divide each bucket into batches |
| 144 | all_batches = [] |
| 145 | for bid in self.bucket_idx_to_data_idxs.keys(): |
| 146 | bucket_size = len(self.bucket_idx_to_data_idxs[bid]) |
| 147 | # num_rest = bucket_size % self.batch_size |
| 148 | n_batches = bucket_size // self.batch_size |
| 149 | for i in range(n_batches): |
| 150 | all_batches.append( |
| 151 | self.bucket_idx_to_data_idxs[bid][(i * self.batch_size) : ((i + 1) * self.batch_size)] |
| 152 | ) |
| 153 | |
| 154 | if self.drop_last is False: |
| 155 | rest = self.bucket_idx_to_data_idxs[bid][n_batches * self.batch_size :] |
| 156 | if len(rest) > 0: |
| 157 | all_batches.append(rest) |
| 158 | |
| 159 | # check if the batch contains too many elements |
| 160 | if self.max_batch_combined_size < 0: |
| 161 | self.all_batches = all_batches |
| 162 | else: |
| 163 | # reduce the batch_size if it is too large (batch * seq_len > self.max_batch_combined_size) |
| 164 | self.all_batches = all_batches |
| 165 | current_idx = 0 |
| 166 | total_removed = 0 |
| 167 | while current_idx < len(self.all_batches): |
| 168 | batch = self.all_batches[current_idx] |
| 169 | batch_size = len(batch) |
| 170 | seq_len = (self.dataset_lengths[batch]).max() |
| 171 | batch_combined_size = batch_size * seq_len |
| 172 | |
| 173 | if batch_combined_size <= self.max_batch_combined_size: |
| 174 | current_idx += 1 |
| 175 | else: |
| 176 | # if batch_size == 1, remove the sample |
| 177 | if batch_size == 1: |
| 178 | self.all_batches[current_idx] = None |
| 179 | total_removed += 1 |
| 180 | current_idx += 1 |
| 181 | continue |
| 182 | |
| 183 | # divide the batch into two part, add the new part to the end of all_batches |
| 184 | batch_size1 = batch_size // 2 |
| 185 | self.all_batches[current_idx] = batch[:batch_size1] |
| 186 | self.all_batches.append(batch[batch_size1:]) |
| 187 | |
| 188 | # remove [] from all_batches |
| 189 | self.all_batches = [batch for batch in self.all_batches if batch is not None] |
| 190 | log.debug(f"total_removed = {total_removed}") |
| 191 | |
| 192 | # shuffle the batches |