PyTorch-Dataset\DataLoader数据加载

1. 引言

在深度学习项目中，数据加载是训练流程的第一步，也是至关重要的一环。PyTorch 提供了强大而灵活的数据加载机制，帮助开发者高效地管理和预处理数据。本文将带你从零开始，深入理解 PyTorch 中的数据加载原理，并通过实战代码掌握 Dataset 与 DataLoader 的核心用法。

2. PyTorch 数据加载核心概念

PyTorch 中加载数据主要依赖两个核心组件：Dataset 和 DataLoader。

Dataset：提供一种方式去获取每个数据及其对应的 label，同时告诉我们总共有多少个数据。它负责定义数据的索引方式与读取逻辑。
DataLoader：为后面的网络提供不同的数据形式，它将数据一批一批地进行打包，支持 shuffle、并行加载等特性，方便模型训练时按批次取用。

简单来说，Dataset 解决「数据怎么读」的问题，DataLoader 解决「数据怎么给」的问题。

3. 常用数据集的两种组织形式

在实际项目中，数据集通常以以下两种形式组织：

3.1 文件夹名称作为 label

这是最常见的一种形式，文件夹的名称就是数据的标签（label）。例如：

复制代码

Data/FirstTypeData/train/
├── ants/    # 蚂蚁图片，label 为 "ants"
└── bees/    # 蜜蜂图片，label 为 "bees"

每个子文件夹内的图片都属于该类别，读取时只需遍历文件夹即可。

3.2 文本文件记录 label

第二种形式中，label 以文本格式存储。文本文件名与图片名称对应，文本内容即为该图片的 label。例如：

复制代码

Data/SecondTypeData/
├── images/
│   ├── 001.jpg
│   └── 002.jpg
└── labels/
    ├── 001.txt   # 内容为 "cat"
    └── 002.txt   # 内容为 "dog"

这种形式适合 label 信息较为复杂或需要额外描述的场景。

4. 认识 Dataset 抽象类

在动手写自定义数据集之前，先来认识 PyTorch 内置的 Dataset 抽象类。通过 help() 可以查看其完整定义：

python 复制代码

from torch.utils.data import Dataset

help(Dataset)

从输出可以看到，Dataset 是一个抽象类，所有表示「键到数据样本映射」的数据集都应继承它。子类必须重写 __getitem__ 方法，用于根据索引获取对应的数据样本；同时可以重写 __len__ 方法，返回数据集的大小，供 Sampler 和 DataLoader 使用。

5. 路径直接加载数据

在正式封装数据集之前，先来看如何直接通过路径加载一张图片。这里使用 PIL 库：

python 复制代码

from PIL import Image

img_path = "Data/FirstTypeData/train/ants/0013035.jpg"
img = Image.open(img_path)
img.show()

这段代码通过图片路径直接打开并显示图片，是理解后续 Dataset 封装逻辑的基础。

6. 自定义 Dataset 加载数据

接下来，我们通过继承 Dataset 类，实现一个完整的自定义数据集。以下代码以「文件夹名称作为 label」的第一种数据形式为例：

python 复制代码

from torch.utils.data import Dataset
from PIL import Image
import os


class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        # 该魔术方法在创建实例对象时自动调用
        self.root_dir = root_dir          # 根目录
        self.label_dir = label_dir        # 标签目录（即类别名）
        self.path = os.path.join(self.root_dir, self.label_dir)  # 拼接完整路径
        self.img_path = os.listdir(self.path)  # 获取该路径下所有图片文件名

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img, label

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

6.1 代码逐行解析

__init__：初始化方法，接收根目录和标签目录，拼接出图片所在路径，并列出该路径下所有图片文件名。
__getitem__：根据索引 idx 返回对应的图片和标签。这是 Dataset 必须实现的核心方法。
__len__：返回数据集中的样本总数。

6.2 实例化并使用数据集

python 复制代码

root_dir = "Data/FirstTypeData/train"
ants_label_dir = "ants"
bees_label_dir = "bees"

ants_dataset = MyData(root_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_label_dir)

print(len(ants_dataset))   # 输出：124
print(len(bees_dataset))   # 输出：121

# 两个数据集可以直接相加合并
train_dataset = ants_dataset + bees_dataset
print(len(train_dataset))  # 输出：245

# 通过索引获取数据
img, label = train_dataset[200]
print("label：", label)    # 输出：label： bees
img.show()

运行结果：

text 复制代码

124
121
245
label： bees

可以看到，两个 Dataset 可以直接通过 + 运算符合并成一个更大的数据集，这是 PyTorch 提供的便捷特性。

7. 总结

本文从 PyTorch 数据加载的核心概念出发，详细介绍了 Dataset 与 DataLoader 的分工，梳理了两种常见的数据集组织形式，并通过完整的代码示例演示了如何自定义 Dataset 类来加载图片数据。掌握这些基础，你就为后续使用 DataLoader 进行批量训练打下了坚实的基础。

在接下来的学习中，可以进一步探索 DataLoader 的 batch_size、shuffle、num_workers 等参数，以及如何对图片进行 transform 预处理，让数据加载更加高效和规范。