PyTorch入门实践:PyTorch数据处理入门

数据集和数据加载器

学习目标

通过本课程，学员可以学会一些对数据集的处理，完成处理后进行自定义数据集的操作最后使用API管道读取的方式进行数据集的训练。

相关知识点

• 数据集和数据加载器

学习内容

1 数据集和数据加载器

PyTorch 提供了两个用于数据处理的基本组件：

• torch.utils.data.DataLoader
• torch.utils.data.Dataset

它们允许更加方便的使用预加载的数据集，也可以用于自定义数据。[1]

• Dataset：存储样本及其对应的标签。
• DataLoader：在 Dataset 的基础上封装了一个可迭代对象，能够更轻松地访问这些样本。

1.1 加载数据集

使用wget的方式进行数据集的下载并解压，命令如下:

!wget --no-check-certificate https://model-community-picture.obs.cn-north-4.myhuaweicloud.com/ascend-zone/notebook_datasets/2590e70a2bdb11f0825afa163edcddae/FashionMNIST.zip

!unzip FashionMNIST.zip

使用以下代码进行数据集的加载:

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
import torch_npu
device = 'npu:0' if hasattr(torch_npu, 'npu') and torch_npu.npu.is_available() else 'cpu'
print(f"Using {device} device")

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

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=False,
    transform=ToTensor()
)

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

1.3 创建自定义数据集

一个自定义的 Dataset 类必须实现以下三个函数：

__init__、__len__ 和 __getitem__。

下面来看一个实现示例：

该示例中，FashionMNIST 图像存储在一个目录 img_dir 中，而它们的标签则单独存储在一个 CSV 文件 annotations_file 中。

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.3.1 __init__函数

实例化 Dataset 对象时，__init__函数会运行一次。

在这个函数中，初始化存放图像的目录、标注文件（annotations file）以及两种数据变换操作。

labels.csv 文件的内容如下所示：

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

每一行包含一个图像文件名和对应的标签，其中标签是一个从 0 到 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

1.3.2 __len__函数

__len__函数返回数据集中样本的总数。

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

1.3.3 __getitem__函数

__getitem__函数用于根据给定的索引 idx，从数据集中加载并返回一个样本。

具体来说，它会根据索引确定图像在磁盘上的位置，使用 read_image 将其转换为张量（tensor），然后从 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.4 使用 DataLoader 读取数据进行训练

Dataset 每次检索我们数据集中的一个样本及其特征和标签。然而，在训练模型时，我们通常希望能够以“小批量”（minibatches）的形式传递样本，每轮epoch都重新打乱数据以减少模型的过拟合，并使用 Python 的 multiprocessing 来加速数据的读取过程。

DataLoader 是一个可迭代对象，它通过一个简单的 API 为我们抽象了这些复杂性。具体来说，DataLoader 可以：

• 将数据集分成若干个小批次（batches）进行处理。
• 在每个 epoch 开始时自动打乱数据（如果设置 shuffle=True）。
• 使用多线程或多个子进程并行加载数据（通过设置参数如 num_workers），从而加快数据加载速度

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.4.1 遍历 DataLoader

将数据集加载到 DataLoader 中，并可以根据需要遍历整个数据集。

下面的每一次迭代都会返回一个批次的 train_features 和 train_labels，其中分别包含 batch_size=64 个特征和对应的标签。

由于在定义 DataLoader 时设置了 shuffle=True，因此在每轮（epoch）遍历完所有批次之后，数据都会被自动打乱顺序，以帮助模型更好地泛化并减少过拟合。

对数据加载顺序进行更细粒度的控制，可以查看 PyTorch 的 Samplers 。

# 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}")