Returns data loader
(
dataset,
batch_size,
use_data_parallel=False,
shuffle=False,
use_multiprocess=True,
sequential_sampler=False,
)
| 151 | |
| 152 | |
| 153 | def batchify( |
| 154 | dataset, |
| 155 | batch_size, |
| 156 | use_data_parallel=False, |
| 157 | shuffle=False, |
| 158 | use_multiprocess=True, |
| 159 | sequential_sampler=False, |
| 160 | ): |
| 161 | """Returns data loader""" |
| 162 | if sequential_sampler: |
| 163 | batch_sampler = SequentialSampler(batch_size=batch_size, corpus_length=len(dataset)) |
| 164 | else: |
| 165 | batch_sampler = BucketsSampler(buckets=dataset.buckets, batch_size=batch_size, shuffle=shuffle) |
| 166 | loader = TextDataLoader( |
| 167 | dataset=dataset, |
| 168 | batch_sampler=batch_sampler, |
| 169 | collate_fn=dataset.collate_fn, |
| 170 | use_data_parallel=use_data_parallel, |
| 171 | use_multiprocess=use_multiprocess, |
| 172 | ) |
| 173 | |
| 174 | return loader |