前言：为什么不能一直依赖 ImageFolder？

在 PyTorch 入门阶段，我们经常使用自带的 torchvision.datasets.ImageFolder 来加载数据。但它有一个极其严苛的前提条件：数据必须按类别分好子文件夹存放。

但在实际的工程项目或开源数据集中（如图像分类、目标检测等），数据往往是**“大杂烩”**形式的：几千上万张图片全部混放在一个大文件夹中，另附一个标注文件（如 train.txt 或 val.json）。

面对这种“一图一表”的结构，ImageFolder 就不再适用了。此时，我们需要根据标注文件手写自己的 Dataset。本文将以花卉分类任务为例，带你从零手写自定义的 Dataset 与 DataLoader。

一、 数据集目录结构与标注格式

在开始写代码前，我们先明确一下常见的数据集摆放格式。假设我们的花卉数据集（flower_data）结构如下：

flower_data/
├── train_images/          # 所有训练集图片都混放在这里
│   ├── image_001.jpg
│   ├── image_002.jpg
│   └── ...
├── train.txt              # 训练集标注文件
└── val.txt                # 验证集标注文件

打开 train.txt，里面记录了每张图片的名称和它对应的类别 ID（以空格分隔），格式如下：

image_001.jpg 0
image_002.jpg 1
image_003.jpg 2
...

我们接下来的目标，就是让程序自动解析这个 TXT 文件，并把图片送入深度学习的流水线中。

二、 DataLoader 的核心运转逻辑

在写代码之前，我们需要先理解 DataLoader 是如何工作的。可以将它想象成一个“大型京东物流仓库”：

  1. 洗牌（Shuffle）：训练开始前，先把所有货品（数据序列）的顺序打乱。

  2. 按需打包（Batching）：处于待命状态。当 GPU 算完一轮需要新数据时，它能迅速按设定的批次大小（如 Batch Size = 64），打包好 64 个数据送过去。

  3. 多进程提速（num_workers）：GPU 计算非常快，如果每次都用单线程读取图片，就会出现“算力等硬盘”的尴尬局面。DataLoader 通过开启多个子进程预加载数据，完美解决了这个问题。

三、 第一步：解析标注文件，构建路径与标签 List

不论标注是 TXT 还是 JSON，进入 DataLoader 之前的终极目标，是在内存中构建出两个长度相等的 List（列表）：

  • 图片路径 List：存有每张图片的完整路径。

  • 标签 List：存有每张图片对应的类别 ID。

代码实现：解析 TXT 标注

为了保持顺序并让逻辑最清晰，我们直接逐行读取 TXT，并将路径和标签追加到列表中：

import os

def load_annotations(txt_path, root_dir):
    """
    解析 TXT 标注文件并返回图片完整路径列表和标签列表
    TXT 格式示例: image_001.jpg 2
    """
    img_paths =[]
    labels =[]
    
    with open(txt_path, 'r', encoding='utf-8') as f:
        for line in f:
            # 去除换行符并按空格切分
            parts = line.strip().split()
            if len(parts) != 2:
                continue
                
            img_name = parts[0]
            label = int(parts[1])
            
            # 补全实际路径 (关键易错点)
            # TXT 中只有名字，程序无法直接读取，必须拼接前面的文件夹路径
            img_path = os.path.join(root_dir, img_name)
            img_paths.append(img_path)
            labels.append(label)
            
    return img_paths, labels

四、 第二步：重写 Dataset 核心类

有了这两个 List，我们需要继承 torch.utils.data.Dataset，并重写三个必备的魔法方法。这就相当于建立一个属于我们自己的“取货逻辑”。

  • __init__：接收目录、标注路径和预处理操作。在这里调用刚才的解析函数。

  • __len__：返回数据集总大小。

  • __getitem__：最核心的函数。每次 DataLoader 索要数据时，会传进一个随机的索引（如 idx=761），我们依据索引去读图、取标签、做预处理。

代码实现：自定义 Dataset

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

class CustomFlowerDataset(Dataset):
    def __init__(self, root_dir, txt_path, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        # 调用辅助函数，生成两个核心 List
        self.img_paths, self.labels = load_annotations(txt_path, root_dir)
        
    def __len__(self):
        return len(self.img_paths)
    
    def __getitem__(self, idx):
        # 1. 拿着索引去找图片的完整路径
        img_path = self.img_paths[idx]
        
        # 2. 用 PIL 读取图像读入内存
        image = Image.open(img_path).convert('RGB')
        
        # 3. 拿着相同的索引去取标签
        label = self.labels[idx]
        
        # 4. 执行数据增强与预处理 (Transforms)
        if self.transform is not None:
            image = self.transform(image)
            
        # 5. 将标签转换为 LongTensor
        label = torch.tensor(label, dtype=torch.long)
        
        # 返回处理干净的一对数据 (图像张量, 标签张量)
        return image, label

五、 第三步：封装 DataLoader 与可视化校验

代码写完绝对不能直接丢进模型训练，因为数据加载极易出现维度错误或预处理 Bug。必须先抽取一个 Batch 进行数据可视化校验！

⚠️ 避坑指南：我们在这里引入了 num_workers=4 以开启多进程加载。

  1. 如果在 Windows 系统下运行脚本，必须将 DataLoader 的执行与迭代代码放在 if __name__ == '__main__': 保护块下，否则会引发多进程递归报错。

  2. 如果电脑配置较低，或者运行环境是 Jupyter Notebook，建议先把 num_workers 设置为 0，确认代码无误后再逐步调大。

代码实现：实例化与测试校验（完整可运行版）

为了方便大家复现，这里提供整合后的完整版本，大家可以直接复制运行（确保准备好了对应的测试图片和 TXT 文件）：

import os
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np

# ================= 1. 辅助函数 =================
def load_annotations(txt_path, root_dir):
    img_paths = []
    labels =[]
    with open(txt_path, 'r', encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) != 2:
                continue
            img_name = parts[0]
            label = int(parts[1])
            img_path = os.path.join(root_dir, img_name)
            img_paths.append(img_path)
            labels.append(label)
    return img_paths, labels

# ================= 2. Dataset 类 =================
class CustomFlowerDataset(Dataset):
    def __init__(self, root_dir, txt_path, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.img_paths, self.labels = load_annotations(txt_path, root_dir)

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

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        image = Image.open(img_path).convert('RGB')
        label = self.labels[idx]

        if self.transform is not None:
            image = self.transform(image)

        label = torch.tensor(label, dtype=torch.long)
        return image, label

# ================= 3. 执行与校验 =================
if __name__ == '__main__':
    # 定义预处理: Resize、ToTensor()转为[0,1]、Normalize()做ImageNet标准化
    data_transforms = transforms.Compose([
        transforms.Resize((64, 64)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    # 实例化 Dataset
    train_dataset = CustomFlowerDataset(
        root_dir='./flower_data/train_images',
        txt_path='./flower_data/train.txt',
        transform=data_transforms
    )

    # 实例化 DataLoader (Windows 下必须在 __main__ 中运行)
    train_loader = DataLoader(
        train_dataset,
        batch_size=64,
        shuffle=True,
        num_workers=4  # 开启 4 个子进程预加载数据
    )

    # 抽取一个 Batch 的数据进行校验
    inputs, labels = next(iter(train_loader))
    print(f"Batch 数据形状: {inputs.shape}") # 预期:[64, 3, 64, 64]
    print(f"Batch 标签形状: {labels.shape}") # 预期:[64]

    # 取出第一张图片进行反向处理展示
    img_tensor = inputs[0]
    
    # 维度变换: PyTorch [C, H, W] -> matplotlib[H, W, C]
    img_np = img_tensor.cpu().numpy().transpose((1, 2, 0))

    # 逆标准化: X_原图 = X_tensor * std + mean
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img_np = std * img_np + mean
    img_np = np.clip(img_np, 0, 1) # 限制像素范围在 0-1 之间

    # 画图与标签核对
    plt.imshow(img_np)
    plt.title(f"Label ID: {labels[0].item()}")
    plt.axis('off')
    plt.show()

总结

通过以上步骤，我们彻底打通了自定义数据集读取的任督二脉。本质上，构建自定义数据流就是完成了一次数据解耦：

  • 解析函数：负责地址与标签映射。

  • Dataset：负责单张读取与预处理。

  • DataLoader：负责多进程打包与分发。

掌握了这一套完整模板，以后无论面对的是 JSON 格式的检测框、PNG 格式的语义分割掩码，还是文本序列数据，只要修改解析函数提取正确的标签，就能写出适配复杂项目的数据流水线！

# pytorch # 分类 # 人工智能
Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

「2026国际AI黑客松」选手招募

avatar AtomGit开源社区
cover

南京大学、京东等开源 VideoJev： LVBench 73.4%，每题约 1.5 秒

avatar AtomGit开源社区
cover

1024 程序员节｜启元机器人 Developer Playground 上海开发者见面会圆满落幕

avatar AtomGit开源社区