我们做深度学习训练时，加载和处理数据集往往要花大量的时间和精力，本文将结合torch.utils.data中的Dataset和DataLoader介绍如何自定义一个数据集加载模块，并给出代码模板。

• Dataset：负责根据index读取相应数据并执行预处理（负责处理索引index到样本sample映射的一个类）；
• Dataloader：最顶层的抽象，通过index找出一条数据出来。

深度学习中使用Dataset和Dataloader类的流程:

• 定义Dataset并实例化;
• 使用Dataloader加载数据；
• 循环迭代使用Dataloader加载的数据进行训练或者验证；

Part1Dataset构建数据集

自定义Dataset的基本模板：

from torch.utils.data import Dataset, DataLoader

 
class ExampleDataset(Dataset):

    def __init__(self, flag='train'):
        assert flag in ['train', 'test', 'valid']
        self.flag = flag
        # 也可以把数据作为一个参数传递给类，__init__(self, data)；
        # self.data = data
        self.data = self.__load_data__()
    
    def __getitem__(self, index):
        # 根据索引返回数据
        # data = self.preprocess(self.data[index]) # 如果需要预处理数据的话
        return self.data[index]
    
    def __len__(self):
        # 返回数据的长度
        return len(self.data)
    
    # 以下可以不在此定义。
    
    # 如果不是直接传入数据data，这里定义一个加载数据的方法
    def __load_data__(self, csv_paths: list):
        # 假如从 csv_paths 中加载数据，可能要遍历文件夹读取文件等，这里忽略
        # 可以拆分训练和验证集并返回train_X, train_Y, valid_X, valid_Y
        pass
      
    def preprocess(self, data):
        # 将data 做一些预处理
        pass

在我们继承Dataset构建自定义数据集时，一般要有这三个函魔法函数，并根据自己的数据集做响应的修改即可：

• def _ _ init_ _ :

初始化，把数据作为一个参数传给类

• def _ _ getitem_ _:

根据索引获取样本对(x,y) 索引为(0,len(dataset)-1)，根据数据集长度从0开始的索引序列；模型通过这个函数获取一对样本对

• def _ _ len_ _:

表示数据集的长度,最终训练时用到的数据集的样本个数。

构建好了自定义数据集，如何加载数据做训练和验证呢？

Part2创建迭代器

Dataloader 代码示例

上面我们创建了自定义数据集 ExampleDataset，并根据flag设定训练集、测试集、验证集，这里我们创建DataLoader迭代器，并加载数据:

# 数据集实例
train_dataset = ExampleDataset(flag='train')
valid_dataset = ExampleDataset(flag='valid')

# 数据迭代器
train_dataloader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
valid_dataloader = DataLoader(dataset=valid_dataset, batch_size=64, shuffle=True)

加载数据后，数据和标签是tuple元组的形式，我们可以使用enumerate访问可遍历的数组对象即可:

# 使用enumerate访问可遍历的数组对象
for step, (input, target) in enumerate(train_dataloader):
    print('step is :', step)
    # data, label = input, target
    print('data is {}, label is {}'.format(data, label))

# 或者这样使用
for idx, item in enumerate(train_dataloader):
    print('idx:', idx)
    data, label = item
    print('data:', data)
    print('label:', label)