用 Pytorch 框架训练深度学习模型，怎么能避得开 Dataloader 的使用呢？这个类的各个参数是什么意思？如何用好 Dataloader？本文的初衷，就是试着去回答这些问题。
基本设置
-
dataset：torch.Dataset 子类的一个实例。
-
batch_size：批量大小，默认为 1
-
shuffle：指定是否打乱数据集。默认为 False
在单进程下（num_workers=0），根据是否需要数据加载的可复现性，可以通过 shuffle 参数控制 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 到统一的长度。这样可以加速训练，节省计算。这个细节对于大语言模型的训练还是蛮重要的。
-
drop_last：默认为 False。设置为 True 时，如果数据集的总量不能被 batch_size 整除，那么最后一个不完整的 batch 会被丢掉，保证所有 batch 都是同样大小。
-
worker_init_fn：默认为 None。它需要是一个可调用 (Callable) 的对象，比如 function。它会在每个子进程中被调用，并且把每个子进程相应的 worker id (
[
0
,
num_workers
−
1
]
[0, \text{num\_workers} - 1]
[0,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
