Args: batch_size: typical batch_size collate_fn: collate function used to combine samples from dataset drop_last: whether to drop the rest of the data that cannot form a batch num_workers:
(
self,
batch_size,
collate_fn,
drop_last=False,
num_workers=0,
batch_wrt_length=True,
shuffle=True,
max_total_samples_per_epoch=-1,
max_batch_combined_size=-1,
batch_sampler_type="sort",
)
| 47 | """ |
| 48 | |
| 49 | def __init__( |
| 50 | self, |
| 51 | batch_size, |
| 52 | collate_fn, |
| 53 | drop_last=False, |
| 54 | num_workers=0, |
| 55 | batch_wrt_length=True, |
| 56 | shuffle=True, |
| 57 | max_total_samples_per_epoch=-1, |
| 58 | max_batch_combined_size=-1, |
| 59 | batch_sampler_type="sort", |
| 60 | ): |
| 61 | """ |
| 62 | Args: |
| 63 | batch_size: |
| 64 | typical batch_size |
| 65 | collate_fn: |
| 66 | collate function used to combine samples from dataset |
| 67 | drop_last: |
| 68 | whether to drop the rest of the data that cannot form a batch |
| 69 | num_workers: |
| 70 | number of loading threads to use |
| 71 | batch_wrt_length: |
| 72 | whether to batch samples of similar sequence lengths |
| 73 | shuffle: |
| 74 | whether to shuffle data within each bucket |
| 75 | max_total_samples_per_epoch: |
| 76 | maximum total number of samples per epoch. -1: ignored |
| 77 | max_batch_combined_size: |
| 78 | limitation on batch_size * seq_len. -1: ignored |
| 79 | batch_sampler_type: |
| 80 | type of the batch sampler [bucket | sort] |
| 81 | """ |
| 82 | |
| 83 | self.batch_size = batch_size |
| 84 | self.collate_fn = collate_fn |
| 85 | self.drop_last = drop_last |
| 86 | self.num_workers = num_workers |
| 87 | self.batch_wrt_length = batch_wrt_length |
| 88 | self.shuffle = shuffle |
| 89 | self.max_total_samples_per_epoch = max_total_samples_per_epoch |
| 90 | self.max_batch_combined_size = max_batch_combined_size |
| 91 | self.batch_sampler_type = batch_sampler_type |
| 92 | assert self.batch_sampler_type in {"bucket", "sort"} |
| 93 | self.ready = False # record the state whether combined_datasets is ready to be used |
| 94 | |
| 95 | self.datasets: T.List[Dataset] = [] # contains each of the dataset |
| 96 | self.dataset_ratios: T.List[ |
| 97 | float |
| 98 | ] = [] # contains the ratio of the samples in each dataset to be used in every epoch |
| 99 | self.combined_dataset: ConcatDataset = None |
| 100 | self.combined_seq_lens: T.List[int] = None |
| 101 | self.selected_global_idxs: T.List[int] = None # index (of the combined dataset) that is chosen |
| 102 | if self.batch_sampler_type == "bucket": |
| 103 | self.batch_sampler = BucketSampler( |
| 104 | batch_size=self.batch_size, |
| 105 | drop_last=self.drop_last, |
| 106 | shuffle=self.shuffle, |
nothing calls this directly
no test coverage detected