| 43 | self.global_idxs = np.arange(len(self.dataset_lengths)) |
| 44 | |
| 45 | def _batching(self): |
| 46 | # working in local index (0:N) |
| 47 | total_samples = len(self.dataset_lengths) |
| 48 | |
| 49 | permute_idxs = np.random.permutation(total_samples) |
| 50 | self.dataset_lengths = [self.dataset_lengths[i] for i in permute_idxs] |
| 51 | # use the same permute_idxs to permute global idxs (important) |
| 52 | self.global_idxs = [self.global_idxs[i] for i in permute_idxs] |
| 53 | |
| 54 | # sort dataset length from short to long |
| 55 | sort_idxs = np.argsort(self.dataset_lengths, kind="stable") |
| 56 | |
| 57 | # divide each bucket into batches |
| 58 | all_batches = [] |
| 59 | index = 0 # index of sort_idxs |
| 60 | while index < total_samples: |
| 61 | current_batch_size = min(total_samples - index, self.batch_size) |
| 62 | |
| 63 | # Get the biggest sequence length on the batch (last one in batch, since sorted ascendingly) |
| 64 | max_seq_len = self.dataset_lengths[sort_idxs[index + current_batch_size - 1]] |
| 65 | |
| 66 | # Adaptively shrink the batch size if the sequence is too long to fit the memory |
| 67 | if self.max_batch_combined_size > 0: |
| 68 | batch_combined_size = current_batch_size * max_seq_len |
| 69 | while batch_combined_size > self.max_batch_combined_size: |
| 70 | if current_batch_size == 1: |
| 71 | break |
| 72 | current_batch_size = current_batch_size // 2 |
| 73 | max_seq_len = self.dataset_lengths[sort_idxs[index + current_batch_size - 1]] |
| 74 | batch_combined_size = current_batch_size * max_seq_len |
| 75 | if batch_combined_size > self.max_batch_combined_size: |
| 76 | # Even single sample won't fit the memory. Skip the input. |
| 77 | continue |
| 78 | |
| 79 | batch = sort_idxs[index : index + current_batch_size] |
| 80 | all_batches.append(batch) |
| 81 | index += current_batch_size |
| 82 | |
| 83 | self.all_batches = all_batches |
| 84 | self.total_batches = len(self.all_batches) |
| 85 | |
| 86 | def __len__(self): |
| 87 | return self.total_batches |