The class assembles batches by sorting the sequence length of samples. It returns indices of the data in the dataset within a batch.
| 7 | |
| 8 | |
| 9 | class SortSampler: |
| 10 | """ |
| 11 | The class assembles batches by sorting the sequence length of samples. |
| 12 | It returns indices of the data in the dataset within a batch. |
| 13 | """ |
| 14 | |
| 15 | def __init__(self, batch_size, max_batch_combined_size=-1): |
| 16 | """ |
| 17 | Args: |
| 18 | batch_size: |
| 19 | batch size |
| 20 | max_batch_combined_size: |
| 21 | limitation on batch_size * seq_len. -1: ignored |
| 22 | """ |
| 23 | self.batch_size = batch_size |
| 24 | self.max_batch_combined_size = max_batch_combined_size |
| 25 | self.dataset_lengths = None |
| 26 | self.total_batches = None |
| 27 | |
| 28 | def set_dataset_lengths(self, dataset_lengths, global_idxs=None): |
| 29 | """ |
| 30 | Args: |
| 31 | dataset_lengths: |
| 32 | a list containing the length of each data |
| 33 | global_idxs: |
| 34 | global index if the dataset is combined and dataset_length are sub-selected |
| 35 | from the combined dataste. None: 0:N |
| 36 | """ |
| 37 | self.dataset_lengths = dataset_lengths |
| 38 | if isinstance(self.dataset_lengths, (list, tuple)): |
| 39 | self.dataset_lengths = np.array(self.dataset_lengths) |
| 40 | |
| 41 | self.global_idxs = global_idxs # only used when returned in a batch |
| 42 | if self.global_idxs is None: |
| 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 |