文档：Dataset - PyTorch documentation

数据集

数据集 Dataset

torch.utils.data.Dataset 是 PyTorch 所有数据集的抽象基类。根据该类的定义，Dataset 都含有以下两个函数：__getitem__ 和 __len__，也就是根据下标取数，和数据集总数据量。

根据 Python 官方文档的定义，当一个对象含有 __getitem__ 方法时，就成为一个 iterable（可迭代对象），可以使用 for 循环从里面遍历元素。此时，Python 会从下标 0 开始，将 0, 1, 2, ... 代入到 __getitem__ 函数中，直到报错。

因此，Dataset 的这一规定使得数据集总是可以按顺序取样本，也可以使用 for 循环进行赋值。

需要定义自己的数据集时，可以继承 Dataset 类定义一个派生类，只需要正确指定 __getitem__ 和 __len__ 方法即可。

迭代式数据集 IterableDataset

数据加载器

即 DataLoader，这是 pytorch 的数据集加载器，用于按一定的 batch 加载数据集。

使用方法如下：

torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0)

其中的参数为：

• dataset：读取数据集
• batch_size：一次读取的 batch 数
• shuffle：是否打乱顺序
• num_workers：异步读取的子进程数，为 0 则不开启子进程

定义 DataLoader 后，在训练过程种就可以每次取出指定 batch 数量的样本，便于训练。

聚合方法 collate_fn

创建 DataLoader 时，如果不指定聚合函数 collate_fn，PyTorch 将按照默认的方法将数据聚合起来。具体来说，默认聚合方法会接受元组或字典，然后将相同元组位置或者相同字典键的对象进行聚合。对于形状为 size 的 Tensor，它叠成形状为 (B, *size) 的张量，其中 B 为 batch 大小；对于单独的整数（常见于图像分类），默认聚合方法会创建形状为 (B,) 的整数型 Tensor；对于字符串，默认聚合方法会将它们构成一个字符串列表。

torchvision 数据集

MNIST

手写数字识别的经典数据集。

调用方法为：

dataset = torchvision.datasets.MNIST(root="./mnist", train=True,
    transform=torchvision.transforms.ToTensor(), download=True)

其中：

• root 为指定的数据存储目录
• train 为指定训练集或测试机，为 True 则为训练集
• transform 指定图像变换方法，若不指定则数据集会返回 PIL 图像。一般使用 ToTenser 变换为 PyTorch Tensor 类型。
• download=True 会在该数据集不存在时自动下载数据集。当未指定且 root 目录不存在该数据集时，该函数会报错。

数据集返回值：该数据集是 torch.utils.data.DataSet 的继承类，可以用下标取一个数据。每次返回一个二元组，元组的第一个元素为图像，第二个图像为一个整数——该图像的标签。若使用 ToTensor 转换图像，则第一个元素为 1 \times 28 \times 28 形状的 Tensor。