PyTorch|Dataset与DataLoader使用、构建自定义数据集

---

一、Dataset与DataLoader

PyTorch提供的两个常用数据API：

• torch.utils.data.Dataset：用于处理单个训练样本，读取数据特征、size、标签等，并且包括数据转换等；
• torch.utils.data.DataLoader：DataLoader在Dataset周围重载一个可迭代对象，以便轻松访问样本。

官方案例： Fashion-MNIST数据集
torchvision：torch的一个视觉库，将torchvision中的datasets导入进来，就能获得其中的各种数据集

FashionMNIST图像存储在目录img_dir中，标签存储在CSV文件annotations_file中

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt


training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

对上述数据集进行可视化：

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

二、自定义Dataset类

• 构建自定义的Dataset类，需要继承TensorFlow的官方dataset类
• 自定义Dataset类必须实现三个函数:__init__，__len__和__getitem__

pytorch中的dataset类是在pytorch的torch下的utils之下的data文件夹里有一个dataset.py

（一）__init__函数

包含图像、注释文件和两个转换：

• annotations_file：标注文件
• img_dir：图像目录

def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
    self.img_labels = pd.read_csv(annotations_file) #标签存储在CSV文件annotations_file中
    self.img_dir = img_dir #FashionMNIST图像存储在目录img_dir中
    self.transform = transform #图像转换
    self.target_transform = target_transform
1
2
3
4
5

（二）__len__函数

返回数据集的样本数（就是img_labels的长度）

def __len__(self):
    return len(self.img_labels)
1
2

（三）__getitem_\函数

输入索引index，getitem函数从数据集中加载并返回对应index的一个样本：

def __getitem__(self, idx):
		#img_labels的第index行第0列标注了对应的照片文件名称
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    image = read_image(img_path) #使用read_image将图像转换为张量
    label = self.img_labels.iloc[idx, 1] #从self中的csv数据中检索相应的标签
    #调用转换函数
    if self.transform: 
        image = self.transform(image)
    if self.target_transform:
        label = self.target_transform(label)
    return image, label #返回张量图像和相应的标签
1
2
3
4
5
6
7
8
9
10
11

（四）全部代码

import os
import pandas as pd
from torchvision.io import read_image

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

三、将单个样本组成minibatch（DataLoader）

（一）PyTorch的DataLoader源码

1、DataLoader的参数

DataLoader通常是在torch.utils.data下

常用的参数有：

• dataset（数据集）：需要提取数据的数据集，Dataset对象
• batch_size（批大小）：每一次装载样本的个数，int型
• shuffle：是否打乱数据顺序
• sampler(Sampler, optional): 自定义从数据集中取样本的策略，如果指定这个参数，那么shuffle必须为False
• num_workers：进行数据加载时使用单个进程还是多进程进行加载，多进程意为加载速度更快，一般默认为0，表示使用主进程进行加载
• collate_fn (callable, optional): 将一个list的sample组成一个mini-batch的函数，一般用于对于一个batch进行后处理
• pin_memory (bool, optional)： 如果设置为True，那么data loader将会在返回它们之前将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中
• drop_last：当样本数不能被batchsize整除时， 是否舍弃最后一批数据

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)
1
2
3
4

2、init函数

主要做了三件事：构建sampler、构建batch_sampler、构建collate_fn

定义属性：

如果设置了自定义的sampler然后又设置了shuffle=true，这种情况是没有意义的：
（shuffle是官方自定义的一个随机sampler）

设置了batch_sampler的情况下，就不需要设置batch_size、shuffle、sampler和drop_last了：

如果没有设置sampler，则先判断数据集类型，如果使用的是map-style（else逻辑），就根据是否设置shuffle来选择pytorch内置的sampler：

设置了batch_size但是没有设置batch_sampler时，会使用内置的BatchSampler：

如果没有设置collate_fn，就判断auto_collation是否设置（auto_collation是根据batch_sampler是否是None来设置的，如果batch_sampler不是none，auto_collation就是true），default_collate是将batch作为输入，batch输出，并没有对数据做额外处理：

3、iter函数

iter函数返回的是get_iterator的值：

get_iterator根据num_workers的设置选择对应的内置DataLoaderIter：

所以可知，iter函数最终返回的是一个dataloaderiter对象，以SingleProcessDataLoaderIter为例，类里有next_data函数：

SingleProcessDataLoaderIter类是继承了BaseDataLoaderIter类，BaseDataLoaderIter类中的next函数就是使用了子类中的next_data：

（二）使用DataLoader遍历

根据上述源码分析，就可以对dataloader去迭代iter之后调用next函数来获得每一批次的数据：

• 通过DataLoader实现对于数据集的遍历，每次遍历会得到一个batch的数据，这里设置batch_size为64：

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)
1
2
3
4

• iter函数将train_dataloader变成一个迭代器，使用next函数可以以此从迭代器中生成一个一个的批次:

# Display image and label.
train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")
img = train_features[0].squeeze()
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"Label: {label}")
1
2
3
4
5
6
7
8
9

由于batch_size=64，因此最终返回的Feature batch shape以及Labels batch shape均为64。

---

参考：
PyTorch官方文档：Datasets & DataLoaders
5、深入剖析PyTorch DataLoader源码