Introduction to PyTorch系列官方教程中文版-2 DATASETS&DATALOADER

Chapter2 数据集和数据集加载器

  处理数据样本的代码可能会变得混乱且难以维护，我们理想地希望我们的数据集代码与我们的模型训练代码解耦合，以获得更好的可读性和模块化。Pytorch提供了两个数据原语：torch.utils.data.DataLoader和torch.utils.data.Dataset，可以更好地使用预加载数据集，以及自己的数据。
  Pytorch提供了许多预加载的数据集，例如FashionMNIST等。

1 加载数据集

  以下是如何从TorchVision加载Fashion-MNIST数据集的示例。Fashion-MNIST图像数据集，由60000个训练示例和10000个测试示例组成。每个示例都包含一个28×28灰度图像和来自10个类别之一的相关标签。
  我们使用以下参数加载FashionMNIST数据集：

参数	功能
root	存储训练/测试数据集的路径
train	指定训练或测试数据集(train=True:表示当前为训练数据集)
download=True	如果数据不可用，则从Internet下载数据存到root
transform、target_transfrom	指定特征和标签转换

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())

2 迭代和可视化数据集

  可以手动索引Datasets列表： training_data[index]。使用matplotlib来可视化训练数据集中的一些样本。

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()

3 创建一个自定义数据集

  自定义数据类必须实现三个函数：init、len、getitem。下面是一个实现样例：FashionMNIST图像存储在一个目录中:img_dir，他们的标签单独存储在一个CSV文件中：annotations_file。

在接下来的部分中，我们将分解每个函数中发生的事情。

import os
import panads as pd
from torchvision.io import read_image
"""
The labels.csv file looks like:

tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9
"""

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

def __len__(self):	# 返回数据集datasets中样例数量
	return len(self.img_labels)

def __getitem__(self, idx):	# 返回索引idx指定的样例
	img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
	# self.img_labels.iloc[idx, 0]返回位于索引idx的那行数据的第一列(0)，也即文件名
	image = read_image(img_path)
	label = self.img_labels.iloc[idx, 1]	# self.img_labels.iloc[idx, 1]返回位于索引idx的那行数据的第二列(1)，也即文件名对应的标签
	if self.transform:
		image = self.transform(image)
	if self.target_transform:
		label = self.target_transform(label)
	return image, label

init

  __init__函数在实例化Dataset对象时运行一次。我们初始化包含图像、注释文件和这两种转换的目录（下一节将详细介绍）
  The labels.csv file looks like:

"""
tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9
"""

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

len

  __len__函数返回数据集中样例的数量

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

getitem

  __getitem__函数根据给定的索引idx从数据集中加载并返回样例。根据索引，该函数可以确定样例在磁盘中的存放位置，并使用read_image函数将其转换成tensor，从self中的csv数据中检索相应的标签。调用他们上的变换函数（如果适用），并返回元组中的张量图像和相应的标签。

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

4 准备好数据，使用DataLoaders加载数据进行训练

  Dataset每次检索数据集中的一个样本的特征和标签。当训练模型时，我们通常希望使用“小批量”样本，在每个epoch重新洗选数据(reshuffle)，以此来降低模型过拟合，同时可以使用Python的multiprocessing来加快数据检索。
  DataLoader是一个可迭代对象，它以一个简单的API为我们抽象了这种复杂性。

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)

5 通过DataLoader来迭代数据

  通过上面的操作，已经将数据加载到了DataLoader中，之后可以对数据集进行遍历。下面代码中的每一次迭代都会返回一个batch的数据，包括train_feature和train_labels(包含了batch_size=64的特征和对应标签)。因为我们指定了shuffle=True，因此当我们迭代了所有的batch之后，数据集会被重新打乱。

# 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}")
"""
Output:
Feature batch shape: torch.Size([64, 1, 28, 28])
Labels batch shape: torch.Size([64])
Label: 4
"""