Pytorch Dataloader 详解

用 Pytorch 框架训练深度学习模型，怎么能避得开 Dataloader 的使用呢？这个类的各个参数是什么意思？如何用好 Dataloader？本文的初衷，就是试着去回答这些问题。

---

基本设置

• dataset：torch.Dataset 子类的一个实例。
• batch_size：批量大小，默认为 1
• shuffle：指定是否打乱数据集。默认为 False

在单进程下（num_workers=0），根据是否需要数据加载的可复现性，可以通过 shuffle 参数控制 Dataloader 的随机性。我们之后会看到，多进程下需要更复杂的设置来保证可复现性。

---

进阶设置

• sampler：定义如何从数据集中采样数据。当 shuffle=False 时，Pytorch 默认使用 SequentialSampler，从数据集中顺序取出数据；当 shuffle=False 时，Pytorch 会使用 RandomSampler，随机取数据。
  我们可以传入自定义的 Sampler 子类的实例。
  参考：data shuffling in Pytorch dataloader；Samplers of Dataloader

---

• num_workers：默认为 0，利用主进程进行数据加载；设置它为一个大于零的整数时，Pytorch 会开启多个子进程，同时进行数据加载。
  那么设置为多少合适呢？一般来说，把它设置为 CPU 内核的个数，是大家常用的准则。

from multiprocessing import cpu_count
n_cpu = cpu_count() # 检查 CPU 内核个数

---

• collate_fn：定义数据处理和合并成 batch 的方式。
  它有什么应用场景呢？比如我们想把一系列长度不相等的张量合并为一个 batch——输入的长度不同的 RNN 网络会有这个需求。直接 stack 是行不通的，它要求所有的张量的大小一致。这时可以利用 collate_fn ，把当前 batch 的样本按照当前 batch 的最大长度进行 padding（可以使用辅助函数 pad_sequence），这样就可以组合成 batch 了。

from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence

def collate_fn(train_data):
    train_data.sort(key=lambda data: len(data), reverse=True) # 按照序列长度进行排序
    data_length = [len(data) for data in train_data] 
    train_data = pad_sequence(train_data, batch_first=True, padding_value=0) # 利用 rnn_utils.pad_sequence 进行 padding
    return train_data, data_length

注：这里之所以要先把数据按照长度从大到小排序，是因为我们想载入 batch 数据后，用 pack_padded_sequence() 将填充过的变长序列“压紧”。不过这些是后话。关于如何较好地处理 RNN 变长输入，可以参考文章 pytorch 如何在 LSTM 中输入可变长的序列、pytorch RNN 变长输入序列问题

Pytorch 的官网教程 中也有一个例子：

from torch.nn.utils.rnn import pad_sequence

# function to collate data samples into batch tensors
def collate_fn(batch):
    src_batch, tgt_batch = [], []
    for src_sample, tgt_sample in batch:
        src_batch.append(src_sample)
        tgt_batch.append(tgt_sample)

    src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)
    tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)
    return src_batch, tgt_batch

train_dataloader = DataLoader(train_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)

利用 Dataloader 中的 collate_fn 函数，可以动态将序列 padding 到当前 batch 中的最大长度，而不是将所有训练数据的序列都 padding 到统一的长度。这样可以加速训练，节省计算。这个细节对于大语言模型的训练还是蛮重要的。

---

• pin_memory：默认为 False。设置为 True 时，可以加速 CPU 向 GPU 传输数据的速度。可以与 non_blocking 搭配使用。详见我的另一篇博客：Pytorch 中 pin_memory 和 non_blocking 用法

---

• drop_last：默认为 False。设置为 True 时，如果数据集的总量不能被 batch_size 整除，那么最后一个不完整的 batch 会被丢掉，保证所有 batch 都是同样大小。

---

• worker_init_fn：默认为 None。它需要是一个可调用 (Callable) 的对象，比如 function。它会在每个子进程中被调用，并且把每个子进程相应的 worker id ($[0,\text{num\_workers}-1]$ 之间的整数) 作为输入。
• generator：默认为 None。它会被用来产生 base_seed。

我们可以利用上面两个参数控制多进程下 Dataloader 的随机性。默认情况下：

By default, each worker will have its PyTorch seed set to base_seed + worker_id, where base_seed is a long generated by main process using its RNG or a specified generator.
In worker_init_fn, you may access the PyTorch seed set for each worker with either torch.utils.data.get_worker_info().seed or torch.initial_seed(), and use it to seed other libraries before data loading.

通过下面这种方式，我们可以保证多进程下数据加载的可复现性——运行多次，数据会以相同的顺序和方式被加载、训练。参考：Pytorch: REPRODUCIBILITY

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    numpy.random.seed(worker_seed)
    random.seed(worker_seed)

g = torch.Generator()
g.manual_seed(0)

DataLoader(
    train_dataset,
    batch_size=batch_size,
    num_workers=num_workers,
    worker_init_fn=seed_worker,
    generator=g,
)

---

官方文档：torch.utils.data.DataLoader