pytorch数据读取之Dataset与DataLoader

1. 先前处理数据集的代码经常比较混乱并且难以维护

2. 数据集处理代码应该和训练代码解耦合，从而达到模块化和更好的可读性

因此，pytorch提出了两个数据处理类：DataLoader与Dataset

from torch.utils.data import Dataset, DataLoader

class SimpleDataset(Dataset):
    def __init__(self, size):
        self.x = [i for i in range(size)]

    def __getitem__(self, index):
        return self.x[index]

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

dataset = SimpleDataset(322)
# 将dataset对象传入DataLoader类
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False)
# 此处不执行，生成器

for idx, data in enumerate(dataloader):
    print(idx)
    print(len(data))

image.gif

当处理自己的数据时，需要继承Dataset类，重写init, len, getitem函数。

init：读取数据集

len：获取整个数据集的样本个数

getitem：获取样本中的第index个样本

当你需要加载数据时，将dataset对象传入DataLoader类

batch_size：将数据集分成大小为batch_size的序列

shuffle：是否打乱数据。如果打乱，获取样本并不是按照顺序的，多次执行dataloader获取的数据内容也不一样的

drop_last：将数据集按照batch_size大小分批，如果有剩余，是否丢掉剩余的

上面每条样本只有一个数据，当一条样本对应于多个数据时，需要用到字典：

from torch.utils.data import Dataset, DataLoader

class SimpleDataset(Dataset):
    def __init__(self, size):
        self.x = [i for i in range(size)]
        self.y = [i + 1 for i in range(size)]

    def __getitem__(self, index):
        sample = {"x": self.x[index], "y": self.y[index]}
        return sample

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

dataset = SimpleDataset(322)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False)

for idx, data in enumerate(dataloader):
    print(idx)
    print(len(data))

image.gif

对于非向量化的数据同样可处理：

from torch.utils.data import Dataset, DataLoader

class SimpleDataset(Dataset):
    def __init__(self, size):
        self.x = ["I love you" for i in range(size // 2)] + ["I love you very much" for i in range(size // 2, size)]
        self.y = [i + 1 for i in range(size)]

    def __getitem__(self, index):
        sample = {"x": self.x[index], "y": self.y[index]}
        return sample

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

dataset = SimpleDataset(322)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False)

for idx, data in enumerate(dataloader):
    print(idx)
    print(len(data))

image.gif

读取文本后再向量化有一个好处，就是进行padding时，只需要考虑当前batch中最大长度，而不需要考虑整个数据集的最大长度。

并且如果数据集太大，无法在SimpleDataset的init函数中全部向量化，则可以只读原始数据，通过dataloader分批之后再向量化。

但是以下处理是错误的：

from torch.utils.data import Dataset, DataLoader

class SimpleDataset(Dataset):
    def __init__(self, size):
        self.x = ["I love you".split() for i in range(size // 2)] + \
                 ["I love you very much".split() for i in range(size // 2, size)]
        # self.x 里面的元素为列表，自动解析成每个样本对应的多个元素；
        # 但是因为列表的长度不一致，后面列表中的very,much取不到，因为有的样本不包含这两个元素
        self.y = [i + 1 for i in range(size)]

    def __getitem__(self, index):
        sample = {"x": self.x[index], "y": self.y[index]}
        return sample

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

dataset = SimpleDataset(322)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False)

for idx, data in enumerate(dataloader):
    print(idx)
    print(len(data))

image.gif