DataLoader中sampler参数介绍

大家好，欢迎来到IT知识分享网。

DataLoader | Sampler | DataSet 关系

• Sampler : 提供数据集中元素的索引
• DataSet : 根据 Sampler 提供的索引来检索数据
• DataLoader : 批量加载数据用于后续的训练和测试

Sampler

class Sampler(object):
    r"""Base class for all Samplers. Every Sampler subclass has to provide an __iter__ method, providing a way to iterate over indices of dataset elements, and a __len__ method that returns the length of the returned iterators. """

    def __init__(self, data_source):
        pass

    def __iter__(self):
        raise NotImplementedError

    def __len__(self):
        raise NotImplementedError

PyTorch官网已经实现了多种 Sampler :

SequentialSampler

若 shuffle=False ，且未指定 sampler ，默认使用

class SequentialSampler(Sampler):
    r"""Samples elements sequentially, always in the same order. Arguments: data_source (Dataset): dataset to sample from """

    def __init__(self, data_source):
        self.data_source = data_source

    def __iter__(self):
        return iter(range(len(self.data_source)))

    def __len__(self):
        return len(self.data_source)

RandomSampler

若 shuffle=True ，且未指定 sampler ，默认使用

class RandomSampler(Sampler):
    r"""Samples elements randomly. If without replacement, then sample from a shuffled dataset. If with replacement, then user can specify ``num_samples`` to draw. Arguments: data_source (Dataset): dataset to sample from replacement (bool): samples are drawn with replacement if ``True``, default=``False`` num_samples (int): number of samples to draw, default=`len(dataset)`. This argument is supposed to be specified only when `replacement` is ``True``. """

    def __init__(self, data_source):
        self.data_source = data_source


    def __iter__(self):
        n = len(self.data_source)
        return iter(torch.randperm(n).tolist())

    def __len__(self):
        return self.num_samples

BatchSampler

like sampler, but returns a batch of indices at a time. Mutually exclusive with batch_size, shuffle, sampler, and drop_last

• 在 DataLoader 中设置 batch_sampler=batch_sampler 的时候，上面四个参数都必须是默认值。也很好理解，每次采样返回一个batch，那么 batch_size 肯定为 1

class BatchSampler(Sampler):
    r"""Wraps another sampler to yield a mini-batch of indices. Args: sampler (Sampler): Base sampler. batch_size (int): Size of mini-batch. drop_last (bool): If ``True``, the sampler will drop the last batch if its size would be less than ``batch_size`` Example: >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False)) [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] >>> list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True)) [[0, 1, 2], [3, 4, 5], [6, 7, 8]] """

    def __init__(self, sampler, batch_size, drop_last):
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last

    def __iter__(self):
        batch = []
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size:
                yield batch
                batch = []
        if len(batch) > 0 and not self.drop_last:
            yield batch

    def __len__(self):
        if self.drop_last:
            return len(self.sampler) // self.batch_size
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size

• 可以看到在构造 BatchSampler 实例的时候，需要传入一个sampler作为实参

最佳实践

最近看到一篇推文，分享了一个使模型训练速度提升20%的Trick–BlockShuffle 。fork了原作者的代码，并自定义了 batch_sampler ，源码见：TransformersWsz/BlockShuffleTest

参考自：

• 一个使模型训练速度提升20%的Trick–BlockShuffle
• Pytorch DataLoader详解
• torch.utils.data — PyTorch 1.10.1 documentation
• pytorch中用Mnist数据集dataloader 自定义batchsampler – 代码先锋网 (codeleading.com)
• pytorch 实现一个自定义的dataloader，每个batch都可以实现类别数量均衡 (tqwba.com)
• 一文弄懂Pytorch的DataLoader, DataSet, Sampler之间的关系