Datasets & DataLoaders

一句话解释

Dataset stores the samples and their corresponding labels, and DataLoader wraps an iterable around the Dataset to enable easy access to the samples.

Code for processing data samples can get messy and hard to maintain; we ideally want our dataset code to be decoupled from our model training code for better readability and modularity. Thus, PyTorch provides these two data primitives: torch.utils.data.DataLoader and torch.utils.data.Dataset that allow you to use pre-loaded datasets as well as your own data.

最小例子

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

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])
)

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)

小巧思

我加载数据集时这个root=“data”具体地址会是什么位置？需要我将数据集放置在项目的目录中吗？还是在环境中的目录下？

root=“data” 是一个相对路径，它相对于你运行 Python 命令时所在的当前工作目录

dataset可视化过程中的代码具体含义是？

# Iterating and Visualizing the Dataset
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)) # 创建一个宽 8 英寸、高 8 英寸的画布
cols, rows = 3, 3 #准备把画布划分为3 行 × 3 列 = 9 个位置
for i in range(1, cols * rows + 1): # 循环生成9个子图
    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() #显示完整画布

从 FashionMNIST 训练集中随机选出 9 张图片，以 3 × 3 网格显示，并在每张图片上方标注服装类别。 FashionMNIST 中的标签是 0～9 的数字。这个字典用于将数字转换为便于阅读的类别名称

labels_map[0]  # "T-Shirt"
labels_map[8]  # "Bag"

imshow()：显示图片 img.squeeze()：将图片变成 (28, 28) cmap=“gray”：使用灰度色彩映射

size=(1,)是什么意思

size=(1,)

size 是 torch.randint() 的关键字参数，(1,) 是一个只有一个元素的元组，用来指定输出 Tensor 的形状。

torch.randint(len(training_data), size=(1,))

意思是：

随机生成一个形状为 (1,)、只包含一个随机整数的 Tensor。 例如：

x = torch.randint(10, size=(1,))
print(x)
# tensor([7])
print(x.shape)
# torch.Size([1])

为什么 (1,) 后面有逗号？

在 Python 中，单元素元组必须保留逗号：

a = (1)
print(type(a))
# <class 'int'>

这里的 (1) 只是加了括号的整数。 而：

a = (1,)
print(type(a))
# <class 'tuple'>

才是单元素元组。 因此：

(1)  → 整数 1
(1,) → 只包含整数 1 的元组

size 决定输出形状

torch.randint(10, size=(1,))

生成一维 Tensor，包含 1 个随机数：

tensor([7])       shape=(1,)

torch.randint(10, size=(3,))

生成一维 Tensor，包含 3 个随机数：

tensor([7, 2, 9]) shape=(3,)

torch.randint(10, size=(2, 3))

生成一个 2 行 3 列的二维 Tensor：

tensor([
    [7, 2, 9],
    [1, 4, 6]
])

形状是：

torch.Size([2, 3])

.item() 的作用

原始结果是一个只包含一个元素的 Tensor：

sample_idx_tensor = torch.randint(
    len(training_data),
    size=(1,)
)
print(sample_idx_tensor)
# tensor([25137])

调用：

sample_idx = sample_idx_tensor.item()

会将它转换为普通 Python 整数：

print(sample_idx)
# 25137
print(type(sample_idx))
# <class 'int'>

所以完整代码：

sample_idx = torch.randint(
    len(training_data),
    size=(1,)
).item()

可以理解为：

生成一个随机整数 Tensor
        ↓
tensor([25137])
        ↓ .item()
普通 Python 整数 25137

之所以转换成普通整数，是因为它接下来要作为数据集索引：

img, label = training_data[sample_idx]

另外也可以直接生成零维 Tensor：

sample_idx = torch.randint(
    len(training_data),
    size=()
).item()

但教程使用 size=(1,) 更直观地表达“生成一个随机数”。

创建datasets时，发生的事

必须实现三个基础functions: init, len, getitem.

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

#创建一个名为 CustomImageDataset 继承 pytorch的Dataset的Python类，实现DataLoader能理解的数据访问方式（len，getitem）
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 = decode_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

创建dataset本质

定义一个 Python 类，让这个类遵守 PyTorch 数据集所要求的接口。 对于这种可以通过索引读取到的数据集，最重要的是实现：

__len__()
__getitem__()

用的时候：

创建对象：

dataset = CustomImageDataset(
    annotations_file="labels.csv",
    img_dir="images",
)

获取长度：

len(dataset) #用dataset.__len__()实现的

取某个对象：

dataset[10] # 用dataset.__getitem__(10)实现的

init

除了第一个参数 self，其他参数基本由类的设计者决定。 PyTorch 不关心构造参数是什么，它主要关心创建完成后的对象能否正确响应：

len(dataset)
dataset[index]

init 通常负责什么？ 通常负责：

• 读取较小的标签文件；
• 保存图片目录；
• 保存 transform；
• 建立文件列表；
• 检查必要路径是否存在。 例如：

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

通常不建议在 init 中一次性读取所有图片：

# 数据量大时不建议
self.images = [decode_image(path) for path in all_paths]

因为这可能一次占用大量内存。当前代码只在 init 中读取 CSV，在 getitem 被调用时才读取一张图片，内存利用更合理。

len

len 应该返回数据集中的样本数量，并且通常是非负整。

getitem

getitem 接收索引，返回该索引对应的一个样本。

def __getitem__(self, idx):
    # 从 CSV 第 idx 行读取图片文件名
    img_path = os.path.join(
        self.img_dir,
        self.img_labels.iloc[idx, 0],
    )

    # 加载这一张图片
    image = decode_image(img_path)

    # 从 CSV 第 idx 行读取标签
    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

getitem 不一定必须返回 (image, label)。也可以返回字典：

return {
    "image": image,
    "label": label,
    "path": img_path,
}

如果返回元组：

return image, label

训练循环通常写成：

for images, labels in dataloader:
    ...

如果返回字典：

return {"image": image, "label": label}

训练循环就要写成：

for batch in dataloader:
    images = batch["image"]
    labels = batch["label"]

编写 Dataset 时的实用注意事项

主要注意以下几点：

1. len 的数量必须与可用索引范围一致。
2. getitem 应该只读取一个样本，不要每次重新读取完整 CSV。
3. 返回的图片和标签格式要保持稳定。
4. 图片 transform 和标签 transform 要分清：

self.transform         # 处理图片
self.target_transform  # 处理标签

5. 同一个 batch 中的图片通常要具有相同形状，否则默认 DataLoader 无法直接堆叠。
6. 分类任务的标签最终通常需要是整数类型，例如：

label = int(self.img_labels.iloc[idx, 1])

7. 路径可以使用 Path 写得更清晰：

from pathlib import Path
self.img_dir = Path(img_dir)
img_path = self.img_dir / self.img_labels.iloc[idx, 0]

最简洁的总结是：

自定义 Dataset 是一个 Python 类；init 的参数和内部实现由我们设计，而 len 和 getitem 是为了让 PyTorch 知道“有多少数据”和“如何按索引读取一个样本”。self 相当于 Java 的 this，只不过 Python 要把它显式写在实例方法参数中。

图片 transform 和标签 transform是什么？有什么区别？

transform → 处理输入数据，例如图片 target_transform → 处理目标值，例如分类标签

图片 transform

图片 transform 用来把原始图片转换成适合模型输入的形式。 常见操作包括：

• 调整图片尺寸
• 转成浮点 Tensor
• 将像素值缩放到 [0, 1]
• 标准化
• 随机翻转、旋转、裁剪等数据增强

from torchvision.transforms import v2
import torch

image_transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])

使用：

dataset = CustomImageDataset(
    annotations_file="labels.csv",
    img_dir="images",
    transform=image_transform,
)

原图片可能为：

类型：uint8
形状：(3, 500, 600)
范围：0～255

经过transform后可能变成：

类型：float32
形状：(3, 224, 224)
经过标准化的数值

标签target_transform

target 是模型要预测的目标。在分类任务中通常就是标签。 例如原始 CSV 中可能用文字表示类别：

cat
dog
bird

但模型训练通常需要数字标签：

cat  → 0
dog  → 1
bird → 2

可以定义：

label_to_index = {
    "cat": 0,
    "dog": 1,
    "bird": 2,
}
target_transform = lambda label: label_to_index[label]

然后：

dataset = CustomImageDataset(
    annotations_file="labels.csv",
    img_dir="images",
    transform=image_transform,
    target_transform=target_transform,
)

读取样本时：

image, label = dataset[0]

得到的 label 就不再是 "cat"，而是：

0

也可以把标签转换为 Tensor：

target_transform = lambda label: torch.tensor(
    label,
    dtype=torch.long,
)

之所以经常使用 torch.long，是因为分类任务中的 nn.CrossEntropyLoss 通常要求类别标签是整数索引。

两者的区别

一个样本的处理流程是：

图片文件 → decode_image → transform → 模型输入
CSV标签  → 读取标签     → target_transform → 训练目标

为什么需要分开设置？

因为图片和标签的处理逻辑通常完全不同。 例如：

image_transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.RandomHorizontalFlip(),
    v2.ToDtype(torch.float32, scale=True),
])
target_transform = lambda label: torch.tensor(
    label,
    dtype=torch.long,
)

不能用 Resize 处理分类编号，也不能用标签的整数转换处理图片，因此需要分开。

不传 transform 会怎样？

构造方法中的默认值是：

transform=None
target_transform=None

所以：

if self.transform:

在没有传入 transform 时为 False，直接跳过处理。 例如：

dataset = CustomImageDataset(
    annotations_file="labels.csv",
    img_dir="images",
)

此时返回原始解码图片和原始标签。

一个完整例子

import torch
from torchvision.transforms import v2
image_transform = v2.Compose([
    v2.Resize((224, 224)),
    v2.ToDtype(torch.float32, scale=True),
])
target_transform = lambda label: torch.tensor(
    int(label),
    dtype=torch.long,
)
dataset = CustomImageDataset(
    annotations_file="labels.csv",
    img_dir="images",
    transform=image_transform,
    target_transform=target_transform,
)
image, label = dataset[0]
print(image.shape)
# torch.Size([3, 224, 224])
print(image.dtype)
# torch.float32
print(label.dtype)
# torch.int64

需要注意：对于目标检测、图像分割等任务，target 可能不只是一个类别编号，还可能包含：

• 边界框
• 分割掩码
• 关键点 这时某些图片变换也必须同步作用于 target。例如图片水平翻转后，边界框坐标也必须一起翻转。此时通常需要设计一个同时处理二者的变换：

image, target = joint_transform(image, target)

但对于普通图片分类，可以先简单理解为：

transform        → 把图片变成适合模型读取的输入
target_transform → 把标签变成适合计算损失的目标

仍不确定的问题

transform具体是什么？