【深度学习项目实战·33】自定义数据集：把自己的图片接入PyTorch

文章目录

• • [01 本篇要解决什么问题](#01 本篇要解决什么问题)
  • [02 背景：先约定数据格式，再写读取代码](#02 背景：先约定数据格式，再写读取代码)
  • [03 实操：从CSV到一个Batch](#03 实操：从CSV到一个Batch)
  • • （1）准备可以复现的图片与标注表
    • （2）Dataset负责样本，DataLoader负责组批
    • （3）固定类别映射，统一张量约定
    • （4）把容易隐藏的问题提前暴露
  • [04 一个具体例子：模型训练之前先验收数据](#04 一个具体例子：模型训练之前先验收数据)
  • [05 常见问题与自测](#05 常见问题与自测)
  • [06 本篇回顾与下一步](#06 本篇回顾与下一步)

01 本篇要解决什么问题

前面的项目能跑通，往往是因为公开数据集已经替我们整理好了目录、标签和读取接口。真正换成自己的数据时，最先卡住的常常不是网络，而是"这一张图片究竟对应哪个标签"。

从本篇开始，我们连续做六件事：接入数据、固定划分、训练调参、保存实验、分析测试结果，最后做一个能选择图片的GUI。六篇共用同一套工程，避免每篇从头复制一个训练脚本。

本篇先完成最基础的闭环：**一张图片、一条标注、一个稳定的样本ID，最后变成模型能使用的Tensor。**关键词是Dataset、标注表、类别映射和数据检查。

02 背景：先约定数据格式，再写读取代码

为了让读者能下载小数据复现，我们使用Fashion-MNIST作者仓库中的服饰图片。原始数据是28×28灰度图、十个类别，官方训练和测试分别有60000、10000张。本工程先导出小子集，再以自定义图片数据集的方式读取。

这里的教学选择是：每类取200张开发图片、50张测试图片，开发集随后再划分。这样既保留真实分类任务，又能在CPU上完整走一遍工程流程。它不是官方全量基准实验。

标注表不只记录路径和标签，还记录样本来源。id是追踪单位，group是划分单位；二者职责不同。同一次采集生成多个切片时，每个切片应有不同ID，却可以共享同一个采集组。

字段	例子	工程用途
id	train_00012	在预测和错例里找到原样本
path	images/train_00012.png	相对数据根目录定位文件
label	Ankle_boot	原始类别名称
group	来源图像的内容指纹	相同来源不能跨集合
partition	dev / test	保留独立测试边界

03 实操：从CSV到一个Batch

（1）准备可以复现的图片与标注表

先按第03篇安装匹配的PyTorch与torchvision，再从工程根目录执行：

python -m pip install numpy pillow truststore
python examples/prepare_data.py

生成data/fashion_custom/images、manifest.csv、split.json和source.json。下载器核对作者公布的压缩文件MD5，导出时使用固定随机种子，并记录跳过的重复候选数；已有数据版本不会直接覆盖。

换成自己的图片时，不必重写网络。准备同样列名的CSV，把图片路径写成相对路径，再运行第34篇的划分入口。标签写中文也可以，但同一类别的拼写、空格和大小写需要统一；"正常"和"正常 "不应该被误当成两个类别。

（2）Dataset负责样本，DataLoader负责组批

PyTorch数据教程将单样本读取和批量加载分开。本工程的Dataset接收已经确定的记录列表，返回图像、整数目标和样本ID。下面是实际运行的类，保存时使用配套完整模块中的导入：

class ManifestDataset(Dataset):
    def __init__(self, root, rows, labels, training=False, augment=True):
        self.root = Path(root)
        self.rows = rows
        self.class_to_idx = {name: i for i, name in enumerate(labels)}
        self.transform = transform(training, augment)

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

    def __getitem__(self, index):
        row = self.rows[index]
        with Image.open(image_path(self.root, row['path'])) as im:
            image = ImageOps.exif_transpose(im).convert('L')
            tensor = self.transform(image)
        return {'image': tensor, 'target': self.class_to_idx[row['label']], 'id': row['id']}

__len__回答"有多少条样本"，__getitem__回答"第i条样本是什么"。先读取图像，再应用当前集合对应的变换。返回ID的价值在后面会体现：预测错误时能够回查原文件，而不用猜DataLoader当时的顺序。

（3）固定类别映射，统一张量约定

类别名只从开发数据得到，排序后保存进划分文件。训练、验证、测试以及GUI全部使用这一份顺序，不能在各目录里分别扫描并重新编号。这里的Ankle_boot会排在原始标签编号顺序之外，所以不能再拿官方整数编号直接解释模型输出。

本工程单张输入是float32 [1,28,28]，像素先除以255，再按固定均值0.5、标准差0.5变到约-1,1。这组固定常数没有用测试集拟合。批量输入是[N,1,28,28]，目标组批后是int64 [N]，适用于十分类交叉熵。

from torch.utils.data import DataLoader
from data import ManifestDataset, load_contract

rows, contract = load_contract('data/fashion_custom',
                               'data/fashion_custom/manifest.csv',
                               'data/fashion_custom/split.json')
training = [r for r in rows if contract['assignments'][r['id']] == 'train']
dataset = ManifestDataset('data/fashion_custom', training,
                          contract['labels'], training=True)
batch = next(iter(DataLoader(dataset, batch_size=8, shuffle=False)))
print(batch['image'].shape, batch['target'].dtype, batch['id'][:2])

（4）把容易隐藏的问题提前暴露

划分前逐张解码，发现坏图就明确报错；不在__getitem__里悄悄跳过，否则样本总数、类别分布和评价分母都会变化。缺失路径、重复ID、空标签及未知类别也需要处理。

图片必须位于指定数据根目录内。文件存在但内容后来被替换，也会导致实验含义发生变化，因此划分文件保存解码后内容指纹，加载时再次核对。大规模数据可以做预先审计和缓存，但不能把"程序没有报错"当作数据版本没有变。

本例的group是图像内容指纹，只能解决完全相同图片的分组；换成医学切片、视频帧或工业连拍时，应优先使用患者、视频或工件ID。内容相似却并不完全相同的图像不会被简单哈希识别出来。

04 一个具体例子：模型训练之前先验收数据

本批实际导出2500张图片，随后固定为1600张训练、400张验证、500张测试，每类分别是160、40、50张。小样本按类别抽取，因此不保留某个真实业务场景的自然类别频率。

我建议第一轮只检查三个层次：原始图片是否能看懂、标签是否对应正确、一个Batch的形状和范围是否符合网络。先抽看几个类别，再让前向和反向跑一个Batch；这比遇到loss不下降就更换模型更容易定位问题。

当你迁移自己的数据时，至少保留一组人工核对的样本：路径、原图、类别名、类别编号与转换后的图像。它能发现灰度与RGB混用、通道位置写反、黑白反转等仅看shape无法发现的问题。

如果以后把输入改成三通道224×224，Dataset、网络首层、预处理元数据和部署入口必须一起升级。本工程的28×28灰度约定是这套任务的输入协议，不能当作所有分类任务的通用模板。

05 常见问题与自测

问题1：目录已经按类别分好了，还需要CSV吗？

简单分类可以用ImageFolder。需要记录分组、采集来源、固定测试边界或逐样本分析时，标注表更方便。选择依据是任务信息，而不是哪种接口看起来高级。

问题2：标签需要先做one-hot吗？

本例使用类别索引，Dataset返回整数即可。不要把类别字符串或错误的浮点编号直接传给这一训练流程。

问题3：坏图为什么不直接忽略？

应先审计并形成明确的新数据版本。静默忽略会让每次实际使用的样本集合不透明，还可能只在某个worker里触发。

问题4：group能不能直接写label？

不可以。类别说明预测目标，group说明样本相关性；把整个类别当成一组，会让某类无法同时出现在训练和验证中。

问题5：能否在初始化Dataset时把所有图片读进内存？

小数据可以，但要明确内存开销。当前实现只保存记录并按需读取，更容易迁移；优化前应先测数据加载是否真的成为瓶颈。

06 本篇回顾与下一步

自定义数据集的完成标准是输入、标签和来源都能追溯，且训练和推理对同一输入有一致解释。Dataset接口只是最后承接这些约定的一层。

术语	本篇理解
样本ID	找回原始记录的稳定标识
类别映射	类别名与模型输出位置之间的约定
数据契约	路径、标签、形状、范围与版本的共同约定

下一篇：34｜数据增强与数据划分。完整工程与运行顺序见阅读说明。