PyTorch笔记之 Dataset 和 Dataloader
简介
在 PyTorch 中，我们的数据集往往会用一个类去表示，在训练时用 Dataloader 产生一个 batch 的数据
比如官方例子中对 CIFAR10 图像数据集进行分类，就有用到这样的操作，具体代码如下所示
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2) testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)
简单说，用 一个类 抽象地表示数据集，而 Dataloader 作为迭代器，每次产生一个 batch 大小的数据，节省内存
Dataset
Dataset 是 PyTorch 中用来表示数据集的一个抽象类，我们的数据集可以用这个类来表示，至少覆写下面两个方法即可
这返回数据前可以进行适当的数据处理，比如将原文用一串数字序列表示
- __len__：数据集大小
- __getitem__：实现这个方法后，可以通过下标的方式（ dataset[i] ）的来取得第 $i$ 个数据
下面我们来为编写一个类表示一个情感二分类数据集，继续用苏神整理的数据集
https://github.com/bojone/bert4keras/tree/master/examples/datasets
数据集没有表头，只有2列，一列是评论（文本），另一列是标签，以制表符进行分隔
from torch.utils.data import Dataset, DataLoader import pandas as pd class SentimentDataset(Dataset): def __init__(self, path_to_file): self.dataset = pd.read_csv(path_to_file, sep="\t", names=["text", "label"]) def __len__(self): return len(self.dataset) def __getitem__(self, idx): text = self.dataset.loc[idx, "text"] label = self.dataset.loc[idx, "label"] sample = {"text": text, "label": label} return sample
Dataloader
基本使用
Dataloader 就是一个迭代器，最基本的使用就是传入一个 Dataset 对象，它就会根据参数 batch_size 的值生成一个 batch 的数据
if __name__ == "__main__": sentiment_dataset = SentimentDataset("sentiment.test.data") sentiment_dataloader = DataLoader(sentiment_dataset, batch_size=4, shuffle=True, num_workers=2) for idx, batch_samples in enumerate(sentiment_dataloader): text_batchs, text_labels = batch_samples["text"], batch_samples["label"] print(text_batchs)
Sampler
PyTorch 提供了 Sampler 模块，用来对数据进行采样，可以在 DataLoader 的通过 sampler 参数调用
一般我们的加载训练集的 dataloader ，shuffle参数都会设置为True ，这时候使用了一个默认的采样器——RandomSampler
当 shuffle 设置为 False 时，默认使用的是 SequencetialSampler，其实就是按顺序取出数据集中的元素
在 PyTorch 中默认实现了以下 Sampler，如果我们要使用别的 Sampler， shuffle 要设置为 False
- SequentialSampler
- RandomSampler
- WeightedSampler
- SubsetRandomSampler
SubsetRandomSampler 常用来将数据集划分为训练集和测试集，比如这里就训练集和测试集按7:3 进行分割
n_train = len(sentiment_train_set) split = n_train // 3
indices = list(range(n_train)) train_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices[split:]) valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(indices[:split])
train_loader = DataLoader(sentiment_train_set, sampler=train_sampler, shuffle=False) valid_loader = DataLoader(sentiment_train_set, sampler=valid_sampler, shuffle=False)
具体推荐下面的博文，讲得挺详细的
一文弄懂Pytorch的DataLoader, DataSet, Sampler之间的关系
https://www.cnblogs.com/marsggbo/p/11541054.html
Pytorch Sampler详解
collate_fn
可以用来进行一些数据处理，比如在文本任务中，一般由于文本长度不一致，我们需要进行截断或者填充。对于图片，我们则希望它们有同样的尺寸
我么可以编写一个函数，然后用这个参数调用它，下面是一个简单的例子，我们把文本截断成只有10个字符
def truncate(data_list): """传进一个batch_size大小的数据""" for data in data_list: text = data["text"] data["text"]=text[:10] return data_list test_loader = DataLoader(sentiment_train_set, batch_size=batch_size, shuffle=True, num_workers=2, collate_fn=truncate)
我们可以看看返回的内容是否已经经过截断了
for i in test_loader: print(i) break
这时候返回的是一个列表而不是字典了，其中一个 batch 的返回结果如下，我们可以看到这里一个样本放在了一个字典中
[{'text': '看了一个通宵，实在是', 'label': 1}, 。。。, {'text': '看了携程的其他用户评', 'label': 0}]
下面是没有使用 collate_fn 的返回结果，它会将数据和标签分开，存放在一起，如下所示
{
'text':['3月1号订的,3月15号还没到货 客服每天说下个工作日能到货已经连续5天了 我无语。想早点儿看这本书的人还是去陶宝或卓越上订吧,尤其是广东省的朋友.当当送货太没保证了.',。。。, '非常纯朴的故事,但包含了主人公坎坷的一生,活着就是痛苦,不得不佩服生命的韧性'],
'label': tensor([1, 。。。, 0])
}
