[Pytorch]DataLoader数据读取机制

作用：

在实际项目中，如果数据量很大，考虑到内存有限、I/O 速度等问题，在训练过程中不可能一次性的将所有数据全部加载到内存中，也不能只用一个进程去加载，所以就需要多进程、迭代加载，而 DataLoader 就是基于这些需要被设计出来的。DataLoader 是一个迭代器，最基本的使用方法就是传入一个 Dataset 对象，它会根据参数 batch_size 的值生成一个 batch 的数据，节省内存的同时，它还可以实现多进程、数据打乱等处理。

pytorch的数据读取机制DataLoader包括两个子模块：

1. Sampler模块，主要是生成索引index
2. DataSet模块，主要是根据索引读取数据

Dataset是用来解决数据从哪里读取以及如何读取的问题。pytorch给定的Dataset是一个抽象类，所有自定义的Dataset都要继承它，并且复写__getitem__()和__len__()类方法，__getitem__()的作用是接受一个索引，返回一个样本或者标签。下面通过实例构造一个数据集：

import torch
from torch.utils.data import Dataset

class MyDataset(Dataset):
    # 构造函数
    def __init__(self, data_tensor, target_tensor):
        self.data_tensor = data_tensor
        self.target_tensor = target_tensor
    # 返回数据集大小
    def __len__(self):
        return self.data_tensor.size(0)
    # 返回索引的数据与标签
    def __getitem__(self, index):
        return self.data_tensor[index], self.target_tensor[index]

结合代码可以看到，我们定义了一个名字为 MyDataset 的数据集，在构造函数中，传入 Tensor 类型的数据与标签；在 __len__ 函数中，直接返回 Tensor 的大小；在 __getitem__ 函数中返回索引的数据与标签。

接下来看如何调用刚才定义的数据集。首先随机生成一个 10*3 维的数据 Tensor，然后生成 10 维的标签 Tensor，与数据 Tensor 相对应。利用这两个 Tensor，生成一个 MyDataset 的对象。查看数据集的大小可以直接用 len() 函数，索引调用数据可以直接使用下标。

# 生成数据
data_tensor = torch.randn(10, 3)
target_tensor = torch.randint(2, (10,)) # 标签是0或1

# 将数据封装成Dataset
my_dataset = MyDataset(data_tensor, target_tensor)

# 查看数据集大小
print('Dataset size:', len(my_dataset))
'''
输出：
Dataset size: 10
'''

# 使用索引调用数据
print('tensor_data[0]: ', my_dataset[0])
'''
输出:
tensor_data[0]:  (tensor([ 0.4931, -0.0697,  0.4171]), tensor(0))
'''

pytorch笔记5-数据读取机制DataLoader - 知乎 (zhihu.com) 更详细的解释