Class DataLoader

modelscope/preprocessors/nlp/space/data_loader.py:47–92

Implement of DataLoader.

class DataLoader(object):
    """ Implement of DataLoader. """

    @classmethod
    def add_cmdline_argument(cls, group):
        group.add_argument('--shuffle', type=str2bool, default=True)
        group.add_argument('--sort_pool_size', type=int, default=0)
        return group

    def __init__(self,
                 dataset,
                 batch_size,
                 hparams,
                 collate_fn=None,
                 sampler=None,
                 is_test=False):
        self.dataset = dataset
        self.collate_fn = collate_fn
        self.gpu = hparams.gpu
        self.sort_pool_size = hparams.sort_pool_size

        if sampler is None:
            if hparams.shuffle and not is_test:
                sampler = RandomSampler(dataset)
            else:
                sampler = SequentialSampler(dataset)

        if self.sort_pool_size > 0 and not is_test:
            sampler = SortedSampler(sampler, self.sort_pool_size)

        def reader():
            for idx in sampler:
                yield idx

        drop_last = False if self.gpu <= 1 or is_test else True
        self.reader = batch(reader, batch_size=batch_size, drop_last=drop_last)
        self.num_batches = math.floor(len(dataset) / batch_size) if drop_last \
            else math.ceil(len(dataset) / batch_size)

    def __len__(self):
        return self.num_batches

    def __iter__(self):
        for batch_indices in self.reader():
            samples = [self.dataset[idx] for idx in batch_indices]
            yield self.collate_fn(samples)


class SequentialDataLoaderWrapper: