torch.utils.data

这个package包含数据读取预处理的一些类，下面详细说明其中两个最常用的类：class torch.utils.data.Dataset 和 class torch.utils.data.DataLoader。然后，我们再给出一个代码范例。

1. 生成数据集

class torch.utils.data.Dataset

表示Dataset的抽象类。

我们可以自己定义数据子类继承这个抽象类。所有子类应该重写__len__和__getitem__，前者提供了数据集的大小，后者支持整数索引，范围从0到len(self)。

• __getitem__(self, index)：index是数据集路径列表的索引。一般来说，当数据集很大的时候，我们需要写一个函数专门生成用来存储数据集路径的列表。
• __len__：数据集路径列表的长度。

2. 加载数据集

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=, pin_memory=False, drop_last=False)

参数：

• dataset (Dataset) – 加载数据的数据集，由torch.utils.data.Dataset的子类产生。
• batch_size (int, optional) – 每个batch加载多少个样本(默认: 1)。
• shuffle (bool, optional) – 设置为True时会在每个epoch重新打乱数据(默认: False).
• sampler (Sampler, optional) – 定义从数据集中提取样本的策略。如果指定，则忽略shuffle参数。
• num_workers (int, optional) – 用多少个子进程加载数据。0表示数据将在主进程中加载(默认: 0)
• collate_fn (callable, optional) –
• pin_memory (bool, optional) –
• drop_last (bool, optional) – 如果数据集大小不能被batch size整除，则设置为True后可删除最后一个不完整的batch。如果设为False并且数据集的大小不能被batch size整除，则最后一个batch将更小。(默认: False)

总结：

1. 通过torch.utils.data.Dataset，我们可以定义一个数据集类，用于获取数据（__getitem__的作用）以及数据集的长度（__len__的作用）。
2. 通过torch.utils.data.DataLoader，我们可以定义一个数据加载器，它是可迭代的，通过enumerate(DataLoader的返回)可以获得一个迭代器，逐步遍历批量数据。比如：

train_data=torch.utils.data.DataLoader(…)
for i, (input, target) in enumerate(train_data):

3. 数据集划分

torch.utils.data.random_split(dataset, lengths, generator=<torch._C.Generator object>)

随机将一个数据集分割成给定长度的不重叠的新数据集。可选择固定生成器以获得可复现的结果（效果同设置随机种子）。

参数：

• dataset (Dataset) – 要划分的数据集。
• lengths (sequence) – 要划分的长度。
• generator (Generator) – 用于随机排列的生成器。

给出示例代码：

import torch
from torch.utils.data import random_split
dataset = range(10)
train_dataset, test_dataset = random_split(
    dataset=dataset,
    lengths=[7, 3],
    generator=torch.Generator().manual_seed(0)
)
print(list(train_dataset))
print(list(test_dataset))

输出：

[4, 1, 7, 5, 3, 9, 0]
[8, 6, 2]

torch.Generator().manual_seed(0)和torch.manual_seed(0)的效果相同，我们验证一下。

import torch
from torch.utils.data import random_split
dataset = range(10)
torch.manual_seed(0)
train_dataset, test_dataset = random_split(
    dataset=dataset,
    lengths=[7, 3]
)
print(list(train_dataset))
print(list(test_dataset))

输出：

[4, 1, 7, 5, 3, 9, 0]
[8, 6, 2]