Pytorch中数据读取-Dataset、Dataloader 、TensorDataset 和 Sampler 的使用

0.引言

Pytorch 创建用以输入到模型的数据的一般流程如下：

1. 创建一个 Dataset 对象，实现__getitem__()和__len__()这两个方法

2. 创建一个 DataLoader 对象，该对象可以对上述Dataset对象进行迭代

3. 遍历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.<locals>.CustomDataset object at 0x7f4bf21d1128>
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))