Learn AI with Python · 课时

自定义数据集与 DataLoaders

torch.utils.data.Dataset、__len__/__getitem__、DataLoader、变换和数据增强。

自定义数据集与 DataLoaders 是 CoddyKit 上的免费 Learn AI with Python 课时。这是第 2 节课，共 4 节。你可以在下方免费阅读本课时的完整内容 — 然后在浏览器中使用内置代码编辑器和全天候 AI 导师进行实践。这是 Learn AI with Python 学习路径的一部分，你的进度在网页和 CoddyKit 应用中同步。 Learn AI with Python 课程共包含 4 节课。

向模型提供数据

训练需要一种高效读取、变换和批处理数据的方式。PyTorch 提供了两种抽象：数据集（知道如何获取一个样本）和 DataLoader（将样本组成批次并打乱顺序）。

from torch.utils.data import Dataset, DataLoader

数据集接口

自定义的 Dataset 子类必须实现两个方法：__len__（样本数量）和 __getitem__（返回指定索引处的样本）。PyTorch 会调用它们来提取数据。

实现 len

__len__ 会告诉 PyTorch 数据集的大小，使其知道有多少个索引以及一个训练周期包含多少个批次。

class ImageDataset(Dataset):
    def __init__(self, paths, labels):
        self.paths = paths
        self.labels = labels

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

实现 getitem

__getitem__ 根据给定索引加载并返回一个样本及其标签。您会在这里打开图像文件，并将其转换为 tensor。

from PIL import Image

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        label = self.labels[idx]
        return img, label

为什么需要变换？

原始图像的尺寸和像素范围各不相同。变换可以将它们标准化：调整到固定尺寸、转换为 tensor，并归一化像素值，使模型能够稳定训练。

from torchvision import transforms

transforms.Compose

transforms.Compose 会将多个变换按顺序串联成一个处理流程。典型的链式处理流程是先执行 Resize，再执行 ToTensor，最后执行 Normalize。

tf = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

理解变换

Resize 会固定空间尺寸；ToTensor 会将 PIL 图像转换为 tensor，并将像素缩放到 [0,1]；Normalize 会对每个通道进行平移和缩放，使其具有零均值和单位方差，从而加快收敛。

在数据集中应用变换

将变换传入数据集，并在 __getitem__ 中应用它，这样每个样本在获取时都能经过一致的预处理。

class ImageDataset(Dataset):
    def __init__(self, paths, labels, transform):
        self.paths, self.labels, self.transform = paths, labels, transform

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        return self.transform(img), self.labels[idx]

封装到 DataLoader 中

DataLoader 会将数据集转换为由批次组成的可迭代对象。设置 batch_size 可以控制每个 step 的样本数，设置 shuffle=True 可以在每个训练周期随机打乱顺序（这对训练很重要）。

dataset = ImageDataset(paths, labels, tf)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

使用 num_workers 提速

num_workers 会启动并行子进程，在 GPU 训练的同时加载并变换数据，从而隐藏输入输出延迟。设置为 4 之类的值通常可以让 GPU 持续获得数据，而不是停下来等待。

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)

遍历批次

遍历 DataLoader 以获取批量 tensor。每次迭代都会生成 (images, labels)，其中 images 的形状为 [batch_size, channels, H, W]，可以直接输入模型。

for images, labels in loader:
    print(images.shape)  # torch.Size([32, 3, 224, 224])
    break

快速检查

检验您对数据处理流程的理解。

回顾：数据集与 DataLoaders

您使用 __len__ 和 __getitem__ 构建了自定义 Dataset，使用 transforms.Compose（Resize、ToTensor、Normalize）预处理图像，并将其封装到带有 batch_size、shuffle 和 num_workers 的 DataLoader 中，从而高效地向模型提供批次。

常见问题解答