Pytorch的DataLoader输入输出(以文本为例)

本文不做太多原理介绍，直讲使用流畅。想看更多底层实现-〉传送门。

DataLoader简介

torch.utils.data.DataLoader是PyTorch中数据读取的一个重要接口，该接口定义在dataloader.py脚本中，只要是用PyTorch来训练模型基本都会用到该接口。本文介绍torch.utils.data.DataLoader与torch.utils.data.Dataset结合使用的方法。

torch.utils.data.DataLoader：接收torch.utils.data.Dataset作为输入，得到DataLoader，它是一个迭代器，方便我们去多线程地读取数据，并且可以实现batch以及shuffle的读取等。

torch.utils.data.Dataset：这是一个抽象类，所以我们需要对其进行派生，从而使用其派生类来创建数据集。最主要的两个函数实现为__Len__和__getitem__。
__init__：可以在这里设置加载的data和label。
__Len__：获取数据集大小
__getitem__：根据索引获取一条训练的数据和标签。

dataLoader的基本使用

输入数据格式

在使用torch.utils.data.DataLoader与torch.utils.data.Dataset前，需要对自己的数据读取或者做一些处理，比如我已经将我的文本数据读取到Dict里了，格式如下：(就放了两个例子，list里面存储多个Dict，一个Dict存的数据是我的一条样例)

train_pair =
[{'id': 'bb6b40-en', 'question': 'paddy:rice', 'choices': ['walnut:walnut crisp', 'cotton:cotton seed', 'watermelon:melon seeds', 'peanut:peanut butter'], 'text': ['question: paddy:rice. option: walnut:walnut crisp', 'question: paddy:rice. option: cotton:cotton seed', 'question: paddy:rice. option: watermelon:melon seeds', 'question: paddy:rice. option: peanut:peanut butter'], 'label': 1}, {'id': '1af9fe-en', 'question': 'principal:teacher', 'choices': ['police:thief', 'manager:staff', 'teacher:student', 'doctor:nurse'], 'text': ['question: principal:teacher. option: police:thief', 'question: principal:teacher. option: manager:staff', 'question: principal:teacher. option: teacher:student', 'question: principal:teacher. option: doctor:nurse'], 'label': 1
}]

构造DataSet

这里的功能是主要是设置加载的data和label，获取数据集大小并根据索引获取一条训练的数据和标签。是为使用DataLoader作准备。

class MyDataset(Dataset):def __init__(self, data_pairs): super().__init__()self.data = data_pairsdef __len__(self):return len(self.data)def __getitem__(self, index):return self.data[index]train_data = MyDataset(train_pair)

使用collate_fn在DataLoader基础上自定义自己的输出

这里实现的东西很简单，因为DataLoader会自动把多个Dict中的数据合并。但是，如果我上面的text字段存储了List，List就会被合并成元组，但是我在使用数据的时候希望所有的句子一起输入模型。这样我就可以自己定义一个函数，去控制合并的操作。

def my_collate(batch_line):batch_line = deepcopy(batch_line)text = []label = []for line in batch_line:text.extend(line['text']) #我只使用这两个字段，其他的可以不处理不输出label.append(line["label"])batch = {"text":text,"label":label,}return batchtrain_data_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True, collate_fn=my_collate)

这是自己控制合并操作的结果：

{'text': ['question: white pollution:biodegradation. option: industrial electricity:solar energy', 'question: white pollution:biodegradation. option: domestic water:reclaimed water recycling', 'question: white pollution:biodegradation. option: chinese herbal prescriptions:medical research','question: stone wall:earth wall. option: legal:illegal', 'question: stone wall:earth wall. option: riverway:waterway', 'question: stone wall:earth wall. option: new house:wedding room', 'question: kiln:ceramics. option: school:student', 'question: kiln:ceramics. option: oven:bread',], 'label': [3, 0]
}

如果不自己处理，直接使用：

train_data_loader = DataLoader(train_data, batch_size=args.batch_size, shuffle=True)

输出结果就会变成：

{   'id': ['bb6b40-en','1af9fe-en'], 'question': ['paddy:rice', 'principal:teacher'], 'choices':......(省略)'text': ('question: white pollution:biodegradation. option: industrial electricity:solar energy', 'question: white pollution:biodegradation. option: domestic water:reclaimed water recycling', 'question: white pollution:biodegradation. option: chinese herbal prescriptions:medical research','question: stone wall:earth wall. option: legal:illegal'), ('question: stone wall:earth wall. option: riverway:waterway', 'question: stone wall:earth wall. option: new house:wedding room', 'question: kiln:ceramics. option: school:student', 'question: kiln:ceramics. option: oven:bread'), 'label': [3, 0]
}

DataLoader的参数

dataset (Dataset) – 加载数据的数据集。
batch_size (int, optional) – 每个batch加载多少个样本(默认: 1)。
shuffle (bool, optional) – 设置为True时会在每个epoch重新打乱数据(默认: False).
sampler (Sampler, optional) – 定义从数据集中提取样本的策略，即生成index的方式，可以顺序也可以乱序
num_workers (int, optional) – 用多少个子进程加载数据。0表示数据将在主进程中加载(默认: 0)
collate_fn (callable, optional) –将一个batch的数据和标签进行合并操作。
pin_memory (bool, optional) –设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。
drop_last (bool, optional) – 如果数据集大小不能被batch size整除，则设置为True后可删除最后一个不完整的batch。如果设为False并且数据集的大小不能被batch size整除，则最后一个batch将更小。(默认: False)
timeout，是用来设置数据读取的超时时间的，但超过这个时间还没读取到数据的话就会报错。

参考材料

pytorch中的数据导入之DataLoader和Dataset的使用介绍

PyTorch源码解读之torch.utils.data.DataLoader