简介这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者聚焦在PyTorch框架下用Unet完成多类别语义分割任务可应用于医学影像、遥感图像等场景。压缩包共46个文件约69KB以19个py源码文件为核心涵盖模型定义、数据加载、损失函数、学习率调度、指标计算与训练保存等模块另有24个pyc缓存文件、2个txt说明与1个json配置结构紧凑便于直接运行与二次修改。目前已有15264人学习下载热度较高。资源围绕Unet的编码器—解码器结构、跳跃连接、多通道输出层设计以及交叉熵损失、Adam优化器、数据增强、IoU与混淆矩阵评估等关键环节给出可复用代码读者可据此快速搭建自己的多类别分割流程并在此基础上尝试更换编码器、引入注意力机制或调整学习率策略以提升性能。1. 从一张三通道标注图说起Unet 多类别语义分割到底在做什么你手里有一批自己拍的路面、遥感、医学或者工业质检图像用 LabelMe 或 ArcGIS 标出了「裂缝、坑槽、标线」三类目标导出成一张三通道 PNG。现在想用 Pytorch 跑一个 Unet让模型对每个像素输出这三类的概率——这就是多类别语义分割要解决的事。它和二分类分割最大的区别在于输出通道数从 1 变成 N损失函数从 BCE 换成 CrossEntropy标注图的像素值不再是 0/255 而是 0/1/2 的类别索引。很多人第一次跑 Unet 训练自己的数据集时模型 loss 一直卡在 0.6 不降八成是标注图没转成单通道索引图或者忽略了一个叫 ignore_index 的参数。这篇笔记按「数据怎么准备 → 模型怎么改 → 训练怎么调 → 坑在哪」的顺序把 Pytorch 下 Unet 多类别分割的完整链路拆开讲新手能照着复现熟手能看到边界条件。2. 多类别数据集怎么变成 Unet 能吃的张量2.1 标注格式的三种来源与统一策略自己做的数据集标注来源无非三类LabelMe 的 JSON、ArcGIS 导出的三通道彩色 mask、以及用脚本批量生成的灰度索引图。Unet 的多类别训练只认最后一种——单通道、像素值等于类别索引、背景为 0。彩色 mask 看着直观但直接喂进去会出大问题模型会把 RGB 三个通道当成三个独立特征而实际上它们只是同一个类别标签的编码。常见做法是写一个转换脚本把彩色 mask 按颜色映射表转成索引图。映射表要自己维护比如(0,0,0)→0、(255,0,0)→1、(0,255,0)→2。转换时用 numpy 的布尔索引比逐像素循环快几十倍一张 512×512 的图毫秒级完成。import numpy as np from PIL import Image # 颜色到类别索引的映射按自己的标注规范改 COLOR_MAP { (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 类别1 (0, 255, 0): 2, # 类别2 (0, 0, 255): 3, # 类别3 } def color_mask_to_index(mask_path, save_path): img np.array(Image.open(mask_path).convert(RGB)) index np.zeros(img.shape[:2], dtypenp.uint8) for color, idx in COLOR_MAP.items(): # 三个通道同时匹配才算命中避免颜色串扰 match np.all(img np.array(color), axis-1) index[match] idx Image.fromarray(index).save(save_path) return index这段代码的关键在np.all(..., axis-1)必须三个通道同时相等才判定为该类别。如果只用单通道比较红色 (255,0,0) 和黄色 (255,255,0) 会在第一个通道上撞车导致类别错乱。转换完建议抽 10 张图用np.unique检查像素值分布确认没有出现映射表之外的杂色。2.2 Dataset 与 DataLoader同步增强是翻车重灾区图像和 mask 必须做完全相同的几何变换这是多类别分割里最容易翻车的地方。用 torchvision 的transforms分别处理 image 和 mask随机裁剪和翻转的参数不一致训练几轮后 mask 和原图就错位了loss 会诡异地停在某个值不动。稳妥的写法是自己封装一个 Dataset在__getitem__里用同一个随机种子控制几何变换颜色变换只作用于原图。import torch from torch.utils.data import Dataset import random import numpy as np from PIL import Image class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, img_list, size512): self.img_dir img_dir self.mask_dir mask_dir self.img_list img_list self.size size def __len__(self): return len(self.img_list) def __getitem__(self, idx): name self.img_list[idx] img Image.open(f{self.img_dir}/{name}.jpg).convert(RGB) mask Image.open(f{self.mask_dir}/{name}.png) # 同步缩放mask 用最近邻避免插值出小数类别 img img.resize((self.size, self.size), Image.BILINEAR) mask mask.resize((self.size, self.size), Image.NEAREST) img np.array(img, dtypenp.float32) / 255.0 mask np.array(mask, dtypenp.int64) # 同步随机翻转用同一个随机数决定 if random.random() 0.5: img np.fliplr(img).copy() mask np.fliplr(mask).copy() # 只对原图做颜色抖动mask 不参与 if random.random() 0.5: img img * random.uniform(0.8, 1.2) img np.clip(img, 0, 1) img torch.from_numpy(img).permute(2, 0, 1).float() mask torch.from_numpy(mask).long() return img, mask参数上有三个点要盯住mask 缩放必须用NEAREST用双线性会插出 1.5 这种非整数类别CrossEntropy 直接报错mask 的 dtype 必须是int64uint8在部分 Pytorch 版本里会触发类型不匹配颜色抖动后要clip回 [0,1]否则归一化输入超出范围会让 BN 层统计量漂移。DataLoader 的num_workers在 Windows 上建议设 0 或 2设大了容易在 WSL 里卡死这是 pytorch 环境搭建 wsl 场景下的高频问题。2.3 类别不平衡先看分布再决定要不要加权多类别数据集里背景往往占 80% 以上直接训练模型会倾向于全预测背景mIoU 看着有 0.8 其实全是背景刷出来的。动手前先用脚本统计每个类别的像素占比import numpy as np from PIL import Image import glob counts np.zeros(4) # 假设4类 for p in glob.glob(masks/*.png): m np.array(Image.open(p)) for c in range(4): counts[c] (m c).sum() freq counts / counts.sum() print(类别频率:, freq) # 权重取频率倒数归一化后传给 CrossEntropyLoss weights 1.0 / (freq 1e-6) weights weights / weights.sum() * len(weights) print(建议权重:, weights)如果最小类频率低于 1%加权是必要的如果在 5% 以上先不加权跑一轮看混淆矩阵很多时候问题不在权重而在标注质量。权重别直接取倒数要归一化到均值为 1 附近否则 loss 量级变化太大会让学习率失效。3. Unet 输出层与损失函数多类别的两个必改点3.1 输出通道数、激活函数与上采样方式标准 Unet 的下采样是 4 次通道从 64 翻到 1024上采样回到原分辨率。做多类别时最后 1×1 卷积的输出通道要改成num_classes且不要接 Softmax。Pytorch 的CrossEntropyLoss内部已经包含 LogSoftmax外面再接一层 Softmax 会导致梯度被压两次训练极慢甚至不收敛。这是新手最常踩的坑没有之一。上采样方式上原版 Unet 用转置卷积容易出现棋盘格伪影。我一般把ConvTranspose2d换成nn.Upsample(scale_factor2, modebilinear, align_cornersTrue)加一个 3×3 卷积伪影明显减少参数量还略降。align_corners设 True 是为了让缩放后的像素对齐设 False 在边缘会有半像素偏移小目标分割上能看出差别。import torch.nn as nn class Up(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x, skip): x self.up(x) # 尺寸对齐防止奇数分辨率下 skip 和 up 差一个像素 diff_y skip.size(2) - x.size(2) diff_x skip.size(3) - x.size(3) x nn.functional.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) x torch.cat([skip, x], dim1) return self.conv(x)尺寸对齐那几行别省。输入 512×512 时四次下采样刚好整除但换成 500×500 就会出现 skip 和上采样特征差 1 像素torch.cat直接报维度错误。用F.pad补齐比裁剪更安全裁剪会丢边缘信息。3.2 CrossEntropyLoss 的 ignore_index 与权重传参CrossEntropyLoss有两个参数在多类别分割里必须显式设置。ignore_index用来屏蔽标注边界或未标注区域通常设 255标注时把不确定的像素涂成 255训练时自动跳过。weight传类别权重张量要和类别数等长且放在 GPU 上。import torch import torch.nn as nn num_classes 4 weights torch.tensor([0.1, 1.2, 1.5, 2.0]).cuda() # 按统计结果填 criterion nn.CrossEntropyLoss(weightweights, ignore_index255) # 前向logits 形状 [B, num_classes, H, W]target 形状 [B, H, W] logits model(imgs) # 不要接 softmax loss criterion(logits, masks)注意 target 的形状是[B, H, W]而不是 one-hot 的[B, C, H, W]。如果你从别的框架转过来习惯传 one-hot会报 Expected target size 之类的错。另外ignore_index255要求 mask 里未标注区域确实是 255如果标注工具用的是 0那背景和忽略区就混了得在 Dataset 里手动改。3.3 评价指标mIoU 和混淆矩阵怎么算才不骗自己训练时只看 loss 会被类别不平衡骗。我习惯在每个 epoch 结束后累计混淆矩阵再算 mIoU 和每类 IoU。混淆矩阵用bincount实现最快def update_confusion(conf_mat, pred, target, num_classes, ignore255): # pred: [B,H,W] 的 argmax 结果, target: [B,H,W] mask target ! ignore pred pred[mask] target target[mask] idx target * num_classes pred conf_mat torch.bincount(idx, minlengthnum_classes**2).reshape(num_classes, num_classes) return conf_mat def compute_miou(conf_mat): # 每类 IoU TP / (TP FP FN) tp conf_mat.diag() fp conf_mat.sum(0) - tp fn conf_mat.sum(1) - tp iou tp / (tp fp fn 1e-6) return iou, iou.mean()bincount的minlength必须设成num_classes**2否则某些类别没出现时混淆矩阵维度会缩水reshape 报错。算 mIoU 时背景类要不要算进去看你的任务——如果背景是有效类别就算如果只是「非目标」建议单独看前景类的 IoU背景刷出来的高分没有意义。4. 训练循环与调参从 loss 不降到 mIoU 上 0.74.1 优化器、学习率与 batch size 的搭配Unet 分割任务上Adam 比 SGD 收敛快初始学习率 1e-3 是安全起点配合ReduceLROnPlateau在 mIoU 停滞时降半。batch size 受显存限制512×512 输入下 8GB 显存大概能跑 batch 4用梯度累积凑等效 batch 16 也行。import torch.optim as optim from torch.optim.lr_scheduler import ReduceLROnPlateau optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4) scheduler ReduceLROnPlateau(optimizer, modemax, factor0.5, patience5) for epoch in range(100): model.train() for imgs, masks in train_loader: imgs, masks imgs.cuda(), masks.cuda() optimizer.zero_grad() logits model(imgs) loss criterion(logits, masks) loss.backward() # 梯度裁剪Unet 深了容易梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) optimizer.step() miou validate(model, val_loader) scheduler.step(miou) # modemax 时传指标 print(fepoch {epoch}, mIoU {miou:.4f}, lr {optimizer.param_groups[0][lr]:.6f})weight_decay设 1e-4 是轻量正则数据量小于 2000 张时可以加到 1e-3。梯度裁剪的max_norm1.0对 Unet 够用如果 loss 出现 NaN先降到 0.5 试试。ReduceLROnPlateau的mode要和监控指标方向一致监控 mIoU 用max监控 loss 用min设反了学习率会往错误方向调。4.2 数据增强的边界哪些增强会破坏 mask几何增强翻转、旋转、缩放、弹性形变对分割都安全只要 image 和 mask 同步。颜色增强亮度、对比度、HSV 抖动只作用于 image。有两类增强要慎用Cutout 和 MixUp。Cutout 在分类里好用但分割里挖掉的区域 mask 怎么填是个问题填 255 忽略会浪费样本填背景会引入噪声。MixUp 把两张图线性叠加mask 没法线性叠加多类别下直接失效。我一般用翻转 随机旋转 90 度的倍数 亮度抖动三件套简单可靠。弹性形变对医学和遥感有用但scipy.ndimage.map_coordinates对 mask 要用order0用默认的三次插值会插出小数类别。4.3 从 loss 曲线判断问题出在哪loss 不降分几种情况。前 10 个 epoch loss 在 0.6 附近震荡多半是学习率太大或者 mask 没转成索引图loss 缓慢下降但 mIoU 不动检查是不是背景类主导看混淆矩阵loss 突然变 NaN查梯度裁剪和输入归一化常见的是某张图全黑导致 BN 方差为 0。训练集 loss 降验证集不降加数据增强或者减模型深度Unet 在 2000 张以下的数据集上很容易过拟合。验证时记得model.eval()加torch.no_grad()忘了 eval 会让 BN 用 batch 统计量验证结果波动大。这两行是血泪经验我见过不止一个项目因为漏了 eval 导致验证 mIoU 比实际低 10 个点。5. 避坑与排查多类别 Unet 训练里最常见的 5 个翻车现场5.1 现象loss 卡在 0.69 附近不动预测全是背景原因标注图是 RGB 三通道Dataset 里没转索引CrossEntropyLoss拿到的 target 形状是[B, H, W, 3]或者被错误地当成 3 类模型学不到有效信号退化成全预测背景而背景占比高时 loss 恰好停在-ln(背景频率)附近。解决在 Dataset 的__getitem__里打印一次 mask 的shape和np.unique确认是[H, W]且像素值在[0, num_classes)范围内。是彩色图就先跑 2.1 的转换脚本。5.2 现象训练中途报 Expected target size [B, H, W], got [B, C, H, W]原因target 传成了 one-hot 编码。有些教程为了兼容 Dice Loss 会先转 one-hot但CrossEntropyLoss要的是类别索引。解决要么在 loss 前argmax回索引要么改用nn.BCEWithLogitsLoss配 one-hot。多类别下推荐前者CrossEntropy 数值更稳。5.3 现象mIoU 在 0.3 到 0.4 之间反复横跳验证结果不稳定原因验证时忘了model.eval()BN 层用当前 batch 的均值和方差batch 小的时候统计量噪声大指标自然抖。解决验证循环开头加model.eval()结束加model.train()整个验证包在torch.no_grad()里。另外验证集 batch size 可以设大一点BN 统计更稳。5.4 现象小目标类别 IoU 始终为 0混淆矩阵里该类一行全空原因小类像素占比太低加权不够或者下采样 4 次后小目标在特征图上只剩一两个像素skip connection 也救不回来。解决先确认标注里小类确实存在且像素数不为零然后提高该类权重到频率倒数的 2 到 3 倍如果还不行减少一次下采样Unet 改成 3 层或者用带空洞卷积的变体保持分辨率。7900xtx 这类大显存卡可以开更大的输入尺寸对小目标帮助明显。5.5 现象转 ONNX 后推理结果和 Pytorch 不一致原因nn.Upsample的align_corners在 ONNX 里默认行为和 Pytorch 不同或者动态轴没设对导致尺寸对不上。解决导出时显式指定dynamic_axesalign_corners在 Pytorch 和 ONNX 两边都设 True。导出后用onnxruntime跑同一张图逐像素比对差异应该在 1e-4 以内。pytorch 转 onnx 这一步别等到部署才做训练完就验一遍早发现早改。6. 把 mIoU 再抬 5 个点的三个具体技巧第一个技巧是在验证集上做 TTA测试时增强。推理时把图翻转、旋转 90 度各跑一次logits 平均后再 argmax。多类别分割上 TTA 通常能涨 2 到 3 个点代价是推理时间翻 4 倍。实现上不用改模型包一层函数就行def tta_predict(model, img): # img: [1, 3, H, W] preds [] for k in range(4): # 0/90/180/270 度旋转 x torch.rot90(img, k, dims[2, 3]) with torch.no_grad(): out model(x) out torch.rot90(out, -k, dims[2, 3]) preds.append(out) # 水平翻转再来一轮 x torch.flip(img, dims[3]) with torch.no_grad(): out model(x) preds.append(torch.flip(out, dims[3])) return torch.stack(preds).mean(0)注意旋转用torch.rot90而不是transforms后者会引入插值误差。平均的是 logits 不是 softmax 后的概率logits 平均再 softmax 数值更稳。第二个技巧是用 Lovasz Loss 或 Dice Loss 和 CrossEntropy 组合。CrossEntropy 是逐像素的对 IoU 这种区域指标是间接优化。Lovasz 直接优化 IoU 的凸 surrogate小类上提升明显。组合方式一般是loss 0.5 * CE 0.5 * Lovasz权重按验证集 mIoU 微调。Lovasz 对 ignore_index 的支持要自己处理把忽略像素的梯度置零。第三个技巧是冻结编码器做 warmup。如果用了预训练的 ResNet 或 EfficientNet 做编码器前 5 个 epoch 只训解码器学习率 1e-3之后再解冻全部用 1e-4 微调。这样能避免随机初始化的解码器在早期把预训练编码器的特征带偏。我自己的习惯是 warmup 阶段监控解码器 loss降到 0.3 以下再解冻比固定 epoch 数更靠谱。最后说个验证方法别只看 mIoU 一个数。把验证集里 IoU 最低的 10 张图挑出来可视化看是标注错了还是模型真不行。我做过的一个路面裂缝项目mIoU 卡在 0.62可视化后发现是标注把阴影也标成了裂缝重新标了 50 张后直接到 0.71。数据质量的天花板比模型结构高得多调参之前先把数据过一遍这是我踩了无数坑之后养成的习惯。希望帮到你。本文还有配套的精品资源点击获取