刚跑完FCN模型的第一个完整训练流程我先把数据加载这块的代码理清楚了。语义分割学习系列做到第七篇前几篇我们聊了分割任务的基本概念、PASCAL VOC的数据集结构也简单过了一遍FCN的网络骨架但说实话真正把数据从硬盘上搬到显存里的这部分代码才是决定整个训练能不能跑起来的关键。这次我就把FCN模型加载VOC格式数据集的代码逐段拆开讲清楚包括目录怎么读、标注图怎么转成类别索引、数据增强怎么做才不会把图像和掩码弄错位以及我在实际调试中踩过的坑。这篇文章适合正卡在“模型会搭、但数据一跑就报错”阶段的读者。如果你已经知道FCN的原理但自己写数据加载器时总觉得别扭或者想看看别人是怎么组织VOC读取代码的那这篇很值得耐心看完。我会把每个关键函数的目的、参数选择的原因、容易出错的地方都讲透代码可以直接抄去用也能帮你理解后续其他分割模型的数据加载是怎么设计的。1. 整体设计思路为什么数据加载要单独写一个类1.1 先从FCN的训练流程看数据加载的位置FCN的训练流程大概是这样的从硬盘读取原始图像和对应的标注图对它们做同样的预处理缩放、翻转、归一化然后打包成一个batch送进网络计算输出和标注之间的交叉熵损失反向传播更新权重。这里面最容易出问题的一环就是“同样的预处理”这六个字。图像是普通的RGB三通道图片标注图是PNG格式的索引色图两者根本不是同一种数据形态。标注图里每个像素的值不是颜色而是一个类别编号0表示背景1表示aeroplane2表示bicycle以此类推。如果我们像处理普通图片那样去处理标注图比如直接resize或者归一化类别编号就会被破坏模型学到的就是一堆错误标签。我在设计加载器时第一步就是把类和函数分开VOCDataset类负责单张样本的读取和预处理collate_fn或者Dataset的__getitem__负责把单个样本组装成batch。这样每个环节都能单独调试出问题也知道该查哪里。1.2 我为什么采用“P模式标注 numpy索引转换”方案第一次接触VOC数据的人十有八九会卡在标注图读取这一步。VOC的标注图虽然是.png后缀但它不是普通的RGB图而是带有调色板信息的索引色图。PIL读取这种图时用convert(P)能得到一个二维的、每个像素直接就是类别索引的数组用convert(RGB)反而会得到一堆三通道的颜色值还得自己去颜色表里查反。我在代码里选择的做法是用PIL的Image.open打开标注图先调用convert(P)再通过np.array直接转成numpy数组。这个数组的shape是(H, W)里面每个元素就是类别编号跟模型的输出通道数一一对应。这比先转RGB再查表要省事得多也是后来很多开源分割框架的统一做法。这块的核心逻辑就是标注图本质上是“每个像素一个整数标签”的二维矩阵我们只需要把它从PNG文件里原封不动地还原出来就行了。不要试图把它当彩色图去显示或resize更不要用ImageNet的均值方差去归一化否则标签就全乱了。2. 核心细节解析VOC目录结构、类别表与标注读取2.1 VOC数据集目录结构与读取顺序PASCAL VOC 2012的数据集下载下来后目录结构固定是这样的VOCdevkit/ └── VOC2012/ ├── JPEGImages/ # 原始图像jpg格式文件名如 2007_000027.jpg ├── SegmentationClass/ # 语义分割标注图png格式索引色模式 ├── ImageSets/ │ └── Segmentation/ # 划分文件train.txt, val.txt, trainval.txt └── ...读取顺序其实是反着的先读ImageSets/Segmentation/train.txt拿到所有的文件名前缀然后根据前缀分别去JPEGImages和SegmentationClass里拼接出图像路径和标注路径。这样做的原因是VOC的train.txt、val.txt是官方划分好的保证了训练集和验证集没有重叠我们自己随机划分的话很容易在数据量少的时候造成信息泄漏。我在代码里用一个load_image_ids函数去读txt文件把每一行的strip结果存进list里。有一点值得注意txt文件里每行只是不带扩展名的文件名不要想当然地拼上.jpg就完事因为标注图在另一个目录扩展名也不一样。2.2 类别编码表从颜色到类别ID的映射关系VOC语义分割一共涉及20个前景类别加1个背景类别。标注图里像素值0到20是类别索引对应的颜色在标注图的调色板里定义。常见类别列表如下索引类别名称备注0background背景1aeroplane飞机2bicycle自行车3bird鸟4boat船5bottle瓶子6bus公交车7car汽车8cat猫9chair椅子10cow牛......共20类索引到20很多教程里会给出一张20x3的RGB颜色表那是在可视化的时候用的。如果你需要把预测结果转成彩色图保存就得有这张表但训练的时候模型输出的是21个通道的概率图取argmax得到的就是索引根本不需要经过颜色转换。所以我在数据加载阶段刻意绕过了颜色表只在可视化阶段才引入它。2.3 图像与标注的同步预处理resize、翻转、裁剪必须一致这是整个数据加载器里最容易写错的部分。图像需要resize到固定尺寸标注图也必须resize到完全一样的尺寸图像做随机翻转标注图也要翻转。如果只处理图像忘了处理标注训练出来的模型基本等于废的。我在预处理函数里把图像和标注拼接成一个元组用同一个随机种子去控制翻转、裁剪等所有带随机性的操作。具体做法是先调用random.random()生成一个概率值如果大于0.5就同时翻转两个数组这样保证了两者永远执行相同的变换。裁剪同理裁剪区域的坐标只需算一份然后同时作用到图像和标注上。3. 实操过程VOC格式数据加载器的完整代码实现3.1 数据加载器主结构我习惯把整个加载器拆成三个部分VOCDataset类负责单个样本的读取和预处理build_transform函数负责组合预处理流程main函数里调用DataLoader完成batch组装。下面这段代码是完整可用的核心框架import os import random import numpy as np from PIL import Image import torch from torch.utils.data import Dataset, DataLoader from torchvision import transforms VOC_CLASSES [ background, aeroplane, bicycle, bird, boat, bottle, bus, car, cat, chair, cow, diningtable, dog, horse, motorbike, person, potted plant, sheep, sofa, train, tv/monitor ] class VOCDataset(Dataset): def __init__(self, root, image_settrain, transformNone): self.root root self.image_set image_set self.transform transform # 读取文件列表 self.image_ids self._load_image_ids() # 拼接目录路径 self.images_dir os.path.join(root, JPEGImages) self.masks_dir os.path.join(root, SegmentationClass) def _load_image_ids(self): # 根据 train.txt / val.txt 读取文件名前缀 split_file os.path.join( self.root, ImageSets, Segmentation, self.image_set .txt ) with open(split_file, r) as f: image_ids [line.strip() for line in f.readlines()] return image_ids def __len__(self): return len(self.image_ids) def __getitem__(self, idx): image_id self.image_ids[idx] # 图像路径 img_path os.path.join(self.images_dir, image_id .jpg) mask_path os.path.join(self.masks_dir, image_id .png) # 读取图像 image Image.open(img_path).convert(RGB) # 读取标注关键一步用 convert(P) 拿索引 mask Image.open(mask_path).convert(P) # 转为 numpy 数组 image np.array(image, dtypenp.float32) mask np.array(mask, dtypenp.int64) if self.transform is not None: image, mask self.transform(image, mask) # 转换为 tensor图像 HWC - CHW image torch.from_numpy(image).permute(2, 0, 1) mask torch.from_numpy(mask).long() return image, mask3.2 关键细节为什么标注要读成P模式而不是RGB模式如果你不小心写成了mask Image.open(mask_path).convert(RGB)你会发现np.array(mask)出来的是一个(H, W, 3)的数组里面每个像素是RGB颜色值。VOC标注图里颜色和类别不是等号关系比如类别0背景是黑色(0,0,0)类别1是红色(128,0,0)类别2是绿色(0,128,0)有些颜色还特别接近。如果拿这个三通道数组去算交叉熵损失模型会直接报错因为输出通道数是21而目标是3。所以这段代码的灵魂就是convert(P)。P模式表示调色板模式每个像素值是对应调色板颜色的索引而这个索引正好就是类别编号。VOC官方在制作标注图的时候就是这么安排的。你只要转成P模式再转numpy数组得到的就是(H, W)、取值0~20的类别索引图跟模型的21通道输出正好对齐。3.3 同步预处理Resize、翻转、Normalize到底该怎么写数据增强我这里选择的是随机水平翻转因为对于分割任务左右翻转不改变语义标签的合理性。resize部分我固定缩放到512x512这个尺寸对VOC原图来说会损失一部分小目标细节但能显著加快训练速度算是一个权衡。写同步预处理时我封装了一个自定义的transform函数它接收image和mask两个数组返回同样两个数组class RandomHorizontalFlip: def __call__(self, image, mask): if random.random() 0.5: image np.flip(image, axis1).copy() mask np.flip(mask, axis1).copy() return image, mask class Resize: def __init__(self, size): self.size size # 或 (H, W) def __call__(self, image, mask): # 这里自己实现双线性resize图像最近邻resize标注 image resize_image(image, self.size) # 用双线性 mask resize_mask(mask, self.size) # 用最近邻 return image, mask注意resize时有个区分图像用双线性插值没问题因为颜色值是连续的但标注图是离散的类别索引如果用双线性插值类别值会被插成小数比如1和2之间可能插出1.5合法像素瞬间变成无效标签。所以标注图resize必须用最近邻插值。实际项目里我一般用cython或者简单遍历来写这个resize避免引入OpenCV的padding边界问题。如果你不想自己写torchvision有一些v2的API支持同步变换但在写系列教程时我倾向于把所有操作显式写出来让每一步都看得见。3.4 归一化与数值类型图片转float32标签转int64图像输入网络前要除以255减去均值除以方差。VOC训练时我采用ImageNet的均值和方差因为FCN通常用预训练VGG16做backbone它的权重是基于ImageNet预训练的输入数据的分布也得对齐。归一化这块我会在transform函数里做先除以255再减均值除方差mean np.array([0.485, 0.456, 0.406]) std np.array([0.229, 0.224, 0.225]) image image / 255.0 image (image - mean) / std这里有个细节如果使用预训练的VGG16作为backbone归一化参数必须用ImageNet的均值和方差不能自己统计VOC数据集的均值和方差否则预训练权重失去意义。但全部权重随机初始化的时候自己算VOC的mean/std也完全可以。mask数组保持int64类型PyTorch的CrossEntropyLoss要求target是long类型int64正好对应long。float64反而不行CrossEntropyLoss内部会把target当作索引去gather用的就是long的接口。3.5 组装DataLoader与batch有了Dataset之后DataLoader的配置也有几个讲究。shuffleTrue用于训练集保证每个epoch样本顺序不同避免模型记住顺序。num_workers根据机器配置可以设成4到8对CPU预处理有加速效果但Windows下num_workers0有时候会有线程相关的问题我在实际使用中一般mac或Linux设4Windows保守一点设2。train_dataset VOCDataset( root./data/VOCdevkit/VOC2012, image_settrain, transformtrain_transform ) train_loader DataLoader( train_dataset, batch_size8, shuffleTrue, num_workers4, drop_lastTrue )drop_lastTrue意味着如果样本数不是batch_size的整数倍最后一个不完整的batch会丢掉。拿VOC train集1464张图除以batch size 8正好整除倒不会触发但别的数据集的样本数啥情况都有drop_last能避免最后一个小batch带来的BN层统计量抖动。4. 常见问题与排查技巧4.1 标注图读出来全黑或者值全是0出现这种情况绝大多数时候是因为convert没写对。很多人习惯性地用convert(RGB)去打开标注图然后看到图像显示出来和原图不一样的地方都是黑的numpy数组里0和255之外的值几乎不存在。这其实是标注被当成RGB读进来后颜色值发散到三通道硬编码的类别索引没了。换成convert(P)基本就解决了。我个人排查时习惯在Dataset里临时打印一下mask.min()和mask.max()如果能输出0和20说明读取正确如果全是0那大概率是读错了分割标注文件检查一下SegmentationClass目录下的文件是不是真的是训练集对应的那部分。4.2 运行时Shape不匹配图像是CHW但标注还是HW或反过来FCN输入是四维的BCHW标注是三维的BHW。如果你在__getitem__里忘记给图像做permute图像就还是HWC送进模型后卷积层的通道维度全乱了。这里给图像permute(2,0,1)标注本身是(H,W)不需要增加通道维度因为CrossEntropyLoss期望的target是(B,H,W)。如果你一次性把image和mask都送进DataLoader记得collate_fn不用额外处理因为这两个维度已经对齐了图像是(3,H,W)mask是(H,W)batch维度由DataLoader自动加上。4.3 训练loss一直不降怀疑数据对不上数据对不上有几种可能一种是图像和标注resize时用了不同的插值方式标注被双线性插值后出现小数类别loss变成NaN另一种是随机翻转只处理了图像标注没翻导致标签和内容左右颠倒。第二种尤其阴险因为loss还会降但mIoU永远上不去。我的检查方法是每个epoch开始前随机打印一个batch的可视化结果。把图像、标注和模型第一次前向的预测图叠在一起看。图像里的人是朝左的标注里同一位置如果也是朝左的像素说明数据管线是同步的。如果人和标签方向相反百分之百是翻转没对齐。4.4 显存比预想中占得多很多VOC的标注图是PNG索引色本身存储很小但转成numpy数组后每个像素是int64占8个字节。如果直接把这些int64的mask喂给GPU显存会白白浪费几倍。解决方案是在__getitem__返回前转成torch.long再交给DataLoaderPyTorch里long也是8字节但比int64的numpy数组更自动。更高效的做法是数据在Dataset里就直接保持uint8或int16等进入训练循环再转long不过具体效果还得看机子。我实际测试中如果把mask保持成int8类型np.uint8每张图大概能省掉大部分标注内存。代价是需要在训练代码里手动转类型有时候容易忘。所以我默认方案还是int64转long省心第一。4.5 如何快速验证数据加载器是否正常写代码最忌一股脑全部跑起来出错了再调。我习惯先单独实例化Dataset取一个样本跑一下__getitem__看看输出的shape和数值范围是否符合预期dataset VOCDataset(root./data/VOCdevkit/VOC2012, image_settrain, transformtrain_transform) img, mask dataset[0] print(img.shape, img.dtype, mask.shape, mask.dtype) print(mask.min(), mask.max()) print(mask.unique())如果img.shape是(3,512,512)mask.shape是(512,512)mask.min是0mask.max是20mask.unique()输出0到20之间的整数这套加载器基本可以放心交给DataLoader去跑了。5. 可视化辅助把标注图转成彩色图方便观察平时调试时把标注图直接显示成灰度图人眼很难分辨出哪些像素是飞机哪些是自行车。这时候需要把P模式的索引图转换回RGB彩色图。我写了一个把索引mask映射回颜色的函数def mask_to_color(mask): color_map np.zeros((mask.shape[0], mask.shape[1], 3), dtypenp.uint8) for cls_idx in np.unique(mask): color_map[mask cls_idx] VOC_COLORS[cls_idx] return color_mapVOC_COLORS可以提前准备一个21行3列的数组跟前面类别表一一对应。有了这个函数配合matplotlib或者PIL就能同时显示原图、标注彩色图和预测彩色图一眼看出FCN学到了什么、没学到什么。这步对理解语义分割任务本身和排查数据问题比看loss曲线直观得多。我在调试数据加载器时一定会把可视化函数写进去。6. 加载器扩展与后续改进6.1 从VOC迁移到自定义数据集数据加载器的设计思路其实是通用的。如果自己采集了一批图片用LabelMe或者精灵标注助手之类的工具标注导出格式可能也是VOC风格那这个加载器基本不用改。如果标注格式是COCO的json格式那就需要额外写一个从json解析掩码并转成索引图的模块但后续的Dataset逻辑完全一致。我实际项目里经常这样干先把各种格式的数据全部预处理成“JPEGImages SegmentationClass”这种VOC风格目录之后不管训练什么模型加载器代码都不用动了。这个约定会给整个工程节省大量时间。6.2 从FCN扩展到DeepLab、U-NetFCN的输入输出格式在后面所有分割模型里都是同一套图像是CHW的RGB标注是HW的索引图类别数跟模型输出通道一致。所以这里写好的VOCDataset换成DeepLab或者U-Net时可以直接复用只需要修改类别数或归一化参数。数据增强部分如果后面想训练更强的模型建议把随机翻转扩展成随机缩放、随机裁剪、色彩抖动等操作。注意加入随机裁剪时必须保证裁剪区域同时包含前景物体否则有些样本裁出来全是背景模型会过拟合到背景上。6.3 模型训练中我常用的数据加载参数总结从FCN训练的实际经验来看以下参数组合比较适中图像resize到512x512训练集batch size设8初始学习率设1e-3使用SGDmomentum加载器num_workers设4。如果你的显卡显存比较紧张可以改成resize到384x384batch size设4训练速度减低但显存占用明显下降。如果加载成为时间瓶颈可以考虑把图像和标注预先全部加载到内存里尤其当数据集不大VOC才1499张训练图完全放得下。有时候训练一个epoch只有几分钟但数据加载占了三分之一时间这时用内存缓存能显著提速。写到这里数据加载器的代码和调试思路基本都讲完了。我第一次跑通这套代码时费了不少劲现在回头看核心就一句话图像和标注必须走同一条预处理流水线标注永远保持整数索引。只要这个原则不变形后面不管接FCN还是其他分割模型数据这块都不会再成为负担。