【更多、更及时内容欢迎留意微信公众号： 小窗幽记机器学习 】
0.引言
Pytorch 创建用以输入到模型的数据的一般流程如下：
创建一个 Dataset 对象，实现__getitem__()和__len__()这两个方法
创建一个 DataLoader 对象，该对象可以对上述Dataset对象进行迭代
遍历DataLoader对象，将样本和标签加载到模型中进行训练
在上述流程中会涉及 Dataset 、 Dataloader 、Sampler 和 TensorDataset，以下将逐一介绍。
1. Dataset
Dataset 是一个抽象类，所有自定义的 datasets 都需要继承该类，并且重载__getitem()__方法和__len__()方法 。__getitem()__方法的作用是接收一个索引，返回索引对应的样本和标签，这需要根据真实数据具体实现的逻辑。__len__()方法是返回所有样本的数量。
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
Data = np.array([[1, 2], [3, 4],[5, 6], [7, 8]])
Label = np.array([[0], [1], [0], [2]])
# 创建子类
class CustomDataset(Dataset):
# 初始化，定义数据内容和标签
def __init__(self, Data, Label):
self.Data = Data
self.Label = Label
# 返回数据集大小
def __len__(self):
return len(self.Data)
# 得到数据内容和标签
def __getitem__(self, index):
data = torch.Tensor(self.Data[index])
label = torch.IntTensor(self.Label[index])
return data, label
dataset = CustomDataset(Data, Label)
print(dataset)
print('dataset大小为：', dataset.__len__())
print(dataset.__getitem__(0))
print(dataset[0])
运行结果如下：
<__main__.test_datasets.
dataset大小为： 4
(tensor([1., 2.]), tensor([0], dtype=torch.int32))
(tensor([1., 2.]), tensor([0], dtype=torch.int32))
1.2 延伸
其实有2种类型的 Dataset，一种就是上述这种，名为map-style datasets；另一种是iterable-style datasets。一个iterable-style的dataset实例需要继承IterableDataset类并实现__iter__()方法。这种类型的datasets 特别适用于随机读取代价大甚至不可能的情况，以及batch size取决于获取的数据。例如，读取数据库，远程服务器或者实时日志等数据的时候，可使用该样式，一般时序数据不使用这种样式。
2. DataLoader
torch.utils.data.DataLoader是PyTorch中加载数据集的核心。DataLoader 返回的是可迭代的数据装载器(DataLoader)，其初始化的参数设置如下。
DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
batch_sampler=None, num_workers=0, collate_fn=None,
pin_memory=False, drop_last=False, timeout=0,
worker_init_fn=None, *, prefetch_factor=2,
persistent_workers=False)
在上述定义的CustomDataset基础上使用DataLoader对其进行遍历：
# 创建DataLoader迭代器
dataloader = DataLoader(dataset, batch_size=2, shuffle=False, num_workers=4)
for i, item in enumerate(dataloader):
print('i:', i)
data, label = item
print('data:', data)
print('label:', label)
运行结果：
i: 0
data: tensor([[1., 2.],
[3., 4.]])
label: tensor([[0],
[1]], dtype=torch.int32)
i: 1
data: tensor([[5., 6.],
[7., 8.]])
label: tensor([[0],
[2]], dtype=torch.int32)
3.Sampler
在DataLoader的参数初始化中有两种sampler：sampler和batch_sampler，都默认为None。前者的作用是生成一系列的index，而batch_sampler则是将sampler生成的indices打包分组，得到一个又一个batch的index。生成的index是遍历Dataset所需的索引。例如下面示例中，BatchSampler将SequentialSampler生成的index按照指定的batch size分组。
a = list(BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=False))
print(a)
运行结果如下：
[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]
以下进一步介绍几种较为常见的Sampler：SequentialSampler，RandomSampler 和 BatchSampler。
3.1 SequentialSampler
SequentialSampler初始化形式：torch.utils.data.SequentialSampler(data_source)。SequentialSampler按顺序对数据集采样。其原理是首先在初始化的时候拿到数据集data_source，之后在__iter__方法中首先得到一个和data_source一样长度的range可迭代器。每次只会返回一个索引值。SequentialSampler为DataLoader提供了顺序遍历dataset的方式。
示例代码：
a = [1, 5, 7, 9, 10086]
b = torch.utils.data.SequentialSampler(a)
for x in b:
print(x)
运行结果：
0
1
2
3
4
3.2 RandomSampler
RandomSampler初始化形式：torch.utils.data.RandomSampler(data_source, replacement=False, num_samples=None, generator=None)。RandomSampler初始化参数除了data_source还有以下2个。
num_samples: 指定采样的数量，默认是所有。
replacement: 默认是False，若为True，则表示可以重复采样，即同一个样本可以重复采样，这样可能导致有的样本采样不到。所以此时可以设置num_samples来增加采样数量使得每个样本都可能被采样到。
示例代码：
torch.manual_seed(42) # 固定住 seed, 否则每次都会生成不同的结果indexs
a = [1, 5, 7, 9, 10086]
b = torch.utils.data.RandomSampler(a)
for x in b:
print(x)
运行结果：
1
3
2
4
0
3.3 WeightedSampler
WeightedSampler是根据给定的权重进行采样。WeightedRandomSampler初始化形式：torch.utils.data.WeightedRandomSampler(weights, num_samples, replacement=True, generator=None)，各参数说明如下：
weights (sequence) – 权重列表, 无需权重列表的和为1
num_samples (int) – 抽取的样本数
replacement (bool) – 与RandomSampler中的一样
generator (Generator) – 用于采样的发生器Generator
示例代码：
torch.manual_seed(4)
a= list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True))
print(a)
b = list(WeightedRandomSampler([0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=False))
print(b)
运行结果：
[4, 4, 1, 5, 5]
[4, 2, 1, 5, 3]
3.4 SubsetRandomSampler
SubsetRandomSampler是从给定的索引列表中抽样元素，该过程不重复采样。SubsetRandomSampler初始化形式：torch.utils.data.SubsetRandomSampler(indices, generator=None)。有如下参数：
indices (sequence) – 索引列表
generator (Generator) – 用于采样的发生器Generator
这个采样器常见的使用场景是将训练集划分成训练集和验证集，示例如下：
batch_size = 2
validation_split = .2
shuffle_dataset = True
random_seed = 42
Data = torch.arange(20)
Data = Data.reshape(10, 2)
Label = torch.arange(10)
Label = Label.reshape(10,1)
# 创建子类
class CustomDataset(Dataset):
# 初始化，定义数据内容和标签
def __init__(self, Data, Label):
self.Data = Data
self.Label = Label
# 返回数据集大小
def __len__(self):
return len(self.Data)
# 得到数据内容和标签
def __getitem__(self, index):
data = torch.LongTensor(self.Data[index])
label = torch.LongTensor(self.Label[index]) #torch.IntTensor
return data, label
dataset = CustomDataset(Data, Label)
dataset_size = len(dataset)
indices = list(range(dataset_size))
split = int(np.floor(validation_split * dataset_size))
if shuffle_dataset:
np.random.seed(random_seed)
np.random.shuffle(indices)
train_indices, val_indices = indices[split:], indices[:split]
# Creating PT data samplers and loaders:
train_sampler = SubsetRandomSampler(train_indices)
valid_sampler = SubsetRandomSampler(val_indices)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
sampler=train_sampler)
validation_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
sampler=valid_sampler)
# Usage Example:
print("train data:")
for batch_index, (data, labels) in enumerate(train_loader):
print(data, labels)
print("\nvalidation data:")
for batch_index, (data, labels) in enumerate(validation_loader):
print(data, labels)
运行结果如下：
train data:
tensor([[ 4, 5],
[10, 11]]) tensor([[2],
[5]])
tensor([[18, 19],
[ 6, 7]]) tensor([[9],
[3]])
tensor([[14, 15],
[12, 13]]) tensor([[7],
[6]])
tensor([[0, 1],
[8, 9]]) tensor([[0],
[4]])
validation data:
tensor([[16, 17],
[ 2, 3]]) tensor([[8],
[1]])
PS：上述使用SubsetRandomSampler的时候关键是获取索引列表，可以使用np.random.choice生成：
np.random.choice(indices, dataset_size)
#numpy.random.choice(a, size=None, replace=True, p=None)
#从a(只要是ndarray都可以，但必须是一维的)中随机抽取数字，并组成指定大小(size)的数组
#replace:True表示可以取相同数字，False表示不可以取相同数字
#数组p：与数组a相对应，表示取数组a中每个元素的概率，默认为选取每个元素的概率相同。
3.5 DistributedSampler
DistributedSampler用于在多机多卡情况下分布式训练数据的读取是一个问题，不同的卡读取到的数据应该是不同的。dataparallel的做法是直接将batch切分到不同的卡，这种方法对于多机来说不可取，因为多机之间直接进行数据传输会严重影响效率。于是有了利用sampler确保dataloader只会load到整个数据集的一个特定子集的做法。DistributedSampler就是做这件事的，可以约束数据只加载数据集的子集。它为每一个子进程划分出一部分数据集，以避免不同进程之间数据重复。DistributedSampler一般与torch.nn.parallel.DistributedDataParallel搭配使用。在这种情况下，每个进程都可以将DistributedSampler实例作为DataLoader的sampler，并加载专属于它的原始数据集的子集。
DistributedSampler初始化形式：torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=None, rank=None, shuffle=True, seed=0, drop_last=False)。具体参数定义如下：
dataset – 待采样的Dataset.
num_replicas (int, optional) – 分布式训练过程中使用的进程数，一般是使用world_size，world size指进程总数，在这里就是我们使用的卡数。
rank (int, optional) – 指当前进程序号。
shuffle (bool, optional) – 如果为True (默认), sampler则会shuffle索引。
seed (int, optional) – 用于对sampler进行shuffle的seed ，其前提是shuffle=True. 该seed数字在分布式组中的所有进程中是相同的。默认值为0。
drop_last (bool, optional) – 默认值为False。如果为True, sampler 丢弃尾部的数据，使其在各副本上均匀地可分。如果为False，sampler中将添加额外的索引，使数据在多个副本之间均匀分割。
注意：在分布式模式下，创建DataLoader迭代器之前需要在每个epoch开始时先调用set_epoch()方法，从而令shuffling在多个epoch中生效。否则，总是会使用相同的顺序。示例代码：
sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None), sampler=sampler)
for epoch in range(start_epoch, n_epochs):
if is_distributed:
sampler.set_epoch(epoch)
train(loader)
4. TensorDataset
TensorDataset 可以用来对 tensor 进行打包，其功能类似 python 中的 zip，将输入的tensors捆绑在一起组成元祖。该类通过每一个 tensor 的第一个维度进行索引。因此，该类中的 tensor 第一维度必须相等。Pytorch 中 TensorDataset 类的定义如下：
class TensorDataset(Dataset[Tuple[Tensor, ...]]):
r"""Dataset wrapping tensors.
Each sample will be retrieved by indexing tensors along the first dimension.
Arguments:
*tensors (Tensor): tensors that have the same size of the first dimension.
"""
tensors: Tuple[Tensor, ...]
def __init__(self, *tensors: Tensor) -> None:
assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors), "Size mismatch between tensors"
self.tensors = tensors
def __getitem__(self, index):
return tuple(tensor[index] for tensor in self.tensors)
def __len__(self):
return self.tensors[0].size(0)
通过代码可以看出TensorDataset是Dataset的子类，已经重载了__len__和__getitem__方法。__getitem__表示每个tensor取相同的索引，然后将这个结果组成一个元组。
示例代码：
data = torch.arange(12)
label = torch.arange(12,0,-1)
print("data=", data)
print("label=", label)
data_label = TensorDataset(data, label)
print("data_label len=", data_label.__len__())
print("data_label[0]=",data_label[0])
data_2 = data.reshape(3,4)
label_2 = torch.arange(3)
print("data2=", data_2)
print("label2=", label_2)
data_label_2 = TensorDataset(data_2, label_2)
print("data_label_2 len=", data_label_2.__len__())
print("data_label_2[1]=",data_label_2[1])
运行结果：
data= tensor([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])
label= tensor([12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1])
data_label len= 12
data_label[0]= (tensor(0), tensor(12))
data2= tensor([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]])
label2= tensor([0, 1, 2])
data_label_2 len= 3
data_label_2[1]= (tensor([4, 5, 6, 7]), tensor(1))