简介这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者提供在PyTorch框架下实现Unet多类别语义分割的完整代码工程可用于医学影像、遥感图像等场景的像素级分类任务。压缩包共46个文件约69KB以19个py脚本为核心涵盖模型定义、数据加载与增强、损失函数、学习率调度、评价指标与训练保存等模块另有24个pyc编译文件、2个txt说明与1个json配置目录按dataloaders、modeling、utils等分层组织结构清晰便于二次开发。目前已有15274人学习下载。读者可据此掌握编码器与解码器搭建、跳跃连接、多通道输出层设计、交叉熵损失与IoU评估等关键环节并参考注意力机制、多尺度训练等优化思路快速迁移到自有数据集。1. 多类别语义分割落地从标注图到 Unet 推理的完整链路手头有一批自己标注的多类别分割数据想用 PyTorch 跑通 Unet结果卡在“标签图是彩色 RGB模型却要单通道类别索引”这一步的人不在少数。这个资源解决的正是这件事它不是一个只跑 Cityscapes 或 VOC 的演示脚本而是一套能直接对接“自己数据集”的多类别语义分割实现覆盖数据读取、标签映射、Unet 搭建、训练循环到推理可视化。适合已经会写基础 PyTorch 训练代码、但被多类别标签处理和多尺度特征融合卡住的从业者。下面按“数据怎么进、模型怎么搭、坑在哪、怎么验证”的顺序拆开讲每一步都能直接抄进自己的工程里。2. 数据管线把彩色标注图转成可训练的类别索引多类别分割和二分分割最大的区别在于标签不再是 0/1而是一张和原图同尺寸的索引图。很多人第一次跑多类别时直接把 RGB 标注图喂进去结果 loss 不降、预测全黑问题就出在这一步。2.1 标签映射表的建立与调色板对齐标注工具LabelMe、Labelme 衍生工具或自研标注器导出的 PNG 通常是调色板模式或 RGB 模式每个类别对应一种颜色。训练前必须建立一张“颜色 → 类别索引”的映射表并且保证映射顺序在整个训练和推理阶段一致。常见做法是先用脚本统计标注图里出现过的所有唯一颜色再人工指定索引。下面这段代码就是干这个的import numpy as np from PIL import Image from collections import Counter # 统计标注图中出现的所有颜色及其像素数 def collect_colors(mask_paths): counter Counter() for p in mask_paths: arr np.array(Image.open(p).convert(RGB)) # 把 H×W×3 压成 (H*W)×3再转成元组统计 pixels arr.reshape(-1, 3) counter.update(map(tuple, pixels)) return counter # 人工确认后写死映射表顺序即类别索引 COLOR_TO_INDEX { (0, 0, 0): 0, # 背景 (128, 0, 0): 1, # 类别一 (0, 128, 0): 2, # 类别二 (128, 128, 0): 3, # 类别三 (0, 0, 128): 4, # 类别四 } def rgb_mask_to_index(mask_rgb): 把 H×W×3 的 RGB 标注图转成 H×W 的类别索引图 h, w, _ mask_rgb.shape index_map np.zeros((h, w), dtypenp.uint8) for color, idx in COLOR_TO_INDEX.items(): # 逐通道比对避免浮点误差 match np.all(mask_rgb np.array(color, dtypenp.uint8), axis-1) index_map[match] idx return index_map逻辑说明collect_colors用来摸清标注里到底有哪些颜色避免映射表漏类rgb_mask_to_index用逐通道全等比对而不是距离阈值是因为标注图颜色是离散的用阈值反而容易把相邻类别混在一起。参数上COLOR_TO_INDEX的 value 必须从 0 连续排到num_classes-1中间不能跳号否则CrossEntropyLoss的ignore_index和后续 one-hot 都会错位。注意如果标注图是调色板模式P 模式convert(RGB)会按调色板展开颜色值可能和你在标注工具里看到的不完全一致建议统一先转 RGB 再统计。2.2 Dataset 与 DataLoader 的写法数据管线的第二段是把原图和索引图配对送进网络。这里有两个容易翻车的点一是原图归一化参数要和预训练权重匹配二是索引图不能做插值增强否则会出现 1.5 这种不存在的类别。import torch from torch.utils.data import Dataset, DataLoader import torchvision.transforms as T import cv2 class SegDataset(Dataset): def __init__(self, img_paths, mask_paths, img_size256): self.img_paths img_paths self.mask_paths mask_paths self.img_size img_size # 原图归一化均值方差按 ImageNet 统计量 self.img_tf T.Compose([ T.ToTensor(), T.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img cv2.imread(self.img_paths[idx])[:, :, ::-1] # BGR - RGB img cv2.resize(img, (self.img_size, self.img_size), interpolationcv2.INTER_LINEAR) mask cv2.imread(self.mask_paths[idx]) mask cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) mask cv2.resize(mask, (self.img_size, self.img_size), interpolationcv2.INTER_NEAREST) # 关键最近邻 mask_idx rgb_mask_to_index(mask) return self.img_tf(img), torch.from_numpy(mask_idx).long()逻辑说明原图用INTER_LINEAR缩放没问题但 mask 必须用INTER_NEAREST否则插值会产生映射表里不存在的颜色rgb_mask_to_index会把它归到背景等于白标。torch.from_numpy(mask_idx).long()是因为CrossEntropyLoss要求 target 是LongTensor。img_size建议先统一到 256 或 512Unet 对输入尺寸没有硬性要求但必须是 16 的倍数否则下采样再上采样会对不齐。3. Unet 结构多类别输出通道与跳跃连接的对齐Unet 本身不复杂但多类别场景下输出通道数、上采样方式和跳跃连接的特征拼接顺序是三个必须盯住的地方。3.1 编码器-解码器与输出通道设置标准 Unet 的编码器每下采样一次通道翻倍解码器每上采样一次通道减半最后用 1×1 卷积把通道数压到num_classes。多类别时输出层不加 Softmax因为CrossEntropyLoss内部已经带了 log_softmax。import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net 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): return self.net(x) class Unet(nn.Module): def __init__(self, in_ch3, num_classes5): super().__init__() self.down1 DoubleConv(in_ch, 64) self.down2 DoubleConv(64, 128) self.down3 DoubleConv(128, 256) self.down4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(512, 1024) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.conv4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.conv3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.conv2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.conv1 DoubleConv(128, 64) self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): d1 self.down1(x) d2 self.down2(self.pool(d1)) d3 self.down3(self.pool(d2)) d4 self.down4(self.pool(d3)) b self.bottleneck(self.pool(d4)) u4 self.conv4(torch.cat([self.up4(b), d4], dim1)) u3 self.conv3(torch.cat([self.up3(u4), d3], dim1)) u2 self.conv2(torch.cat([self.up2(u3), d2], dim1)) u1 self.conv1(torch.cat([self.up1(u2), d1], dim1)) return self.out(u1) # 输出 [B, num_classes, H, W]逻辑说明num_classes必须和COLOR_TO_INDEX的类别数完全一致多一个少一个都会在 loss 计算时报维度错误。torch.cat的拼接顺序是“上采样结果在前、编码器特征在后”这个顺序不影响正确性但必须和DoubleConv的输入通道数对应改结构时容易在这里对不上。ConvTranspose2d的stride2保证每次上采样尺寸翻倍和前面的MaxPool2d(2)对称。3.2 损失函数与类别不平衡处理多类别分割里背景像素往往占 70% 以上直接上CrossEntropyLoss会让模型倾向于全预测背景。常见做法是加类别权重或者换成 Dice Loss 与 CE 的组合。# 按像素频率倒数设置权重频率越低的类别权重越高 class_weights torch.tensor([0.1, 1.0, 1.5, 2.0, 2.5]).float().cuda() criterion nn.CrossEntropyLoss(weightclass_weights, ignore_index255) # 训练循环核心片段 for imgs, masks in dataloader: imgs, masks imgs.cuda(), masks.cuda() preds model(imgs) # [B, C, H, W] loss criterion(preds, masks) # masks: [B, H, W] optimizer.zero_grad() loss.backward() optimizer.step()逻辑说明ignore_index255用来忽略标注里“不确定”或“边界”像素前提是你在rgb_mask_to_index里把这类像素映射成 255。class_weights的具体数值要根据自己数据集的像素统计来调上面只是示例不是通用值。如果类别数超过 10 且分布极不均建议再叠一个 Dice Loss公式是1 - 2*|pred∩gt| / (|pred||gt|)按类别算完取平均。4. 训练与推理避坑与常见问题排查这一章集中说几个我实际踩过的坑每个都按“现象 → 原因 → 解决”写方便对照排查。4.1 常见问题排查现象一loss 一直停在ln(num_classes)附近不降。原因标签映射表没对上模型看到的 target 全是同一个类别或者rgb_mask_to_index把所有颜色都归到了背景。 解决在 Dataset 里随机抽几张 mask打印np.unique(mask_idx)确认类别索引分布正常再检查COLOR_TO_INDEX的 key 是否和标注图实际 RGB 值完全一致。现象二训练 loss 正常下降但推理结果全是背景。原因验证集和训练集的归一化参数不一致或者推理时忘了把输出做argmax。 解决推理阶段复用训练时的img_tf输出用torch.argmax(preds, dim1)取类别索引再映射回颜色可视化。现象三显存爆了batch size 降到 1 还是 OOM。原因Unet 在 512×512 输入下bottleneck 通道到 1024中间特征图占用很大。 解决把输入降到 256×256或者把DoubleConv里的通道数整体减半也可以用混合精度训练torch.cuda.amp能省将近一半显存。现象四某些类别 IoU 始终为 0。原因该类别像素太少被CrossEntropyLoss的平均机制淹没了。 解决提高该类别权重或者在采样时对含该类的图做 oversampling也可以先单独统计每个类别的像素占比占比低于 1% 的类别考虑合并或重新标注。现象五上采样后边缘出现棋盘格伪影。原因ConvTranspose2d的 stride 和 kernel 不匹配时会产生不均匀重叠。 解决把ConvTranspose2d换成nn.Upsample(scale_factor2, modebilinear)再接一个 3×3 卷积伪影会明显减轻。4.2 训练参数与验证指标训练轮数不用一上来就设几百先跑 30 轮看 loss 曲线和验证集 mIoU 的趋势。优化器用 Adam 起步学习率 1e-3如果 loss 震荡就降到 1e-4。验证指标建议自己实现一个混淆矩阵按类别算 IoU 再平均比直接看像素准确率靠谱得多。def compute_miou(preds, targets, num_classes): preds/targets: [B, H, W] 的类别索引 ious [] for cls in range(num_classes): pred_mask (preds cls) target_mask (targets cls) inter (pred_mask target_mask).sum().item() union (pred_mask | target_mask).sum().item() if union 0: continue # 该类别在本 batch 没出现跳过 ious.append(inter / union) return sum(ious) / len(ious) if ious else 0.0逻辑说明union 0时跳过而不是记 0是为了避免某个 batch 里没出现的类别拉低整体 mIoU。这个指标在每个 epoch 结束后在验证集上跑一遍比 loss 更能反映真实分割质量。5. 进阶技巧让多类别 Unet 真正可用的几个习惯跑通训练只是第一步真正让这套代码在项目里可用还得处理推理可视化和模型导出这两件事。推理可视化我一般会写一个index_to_color函数把argmax后的索引图映射回和标注一致的调色板叠在原图上做半透明显示这样一眼就能看出哪类分错了。INDEX_TO_COLOR {v: k for k, v in COLOR_TO_INDEX.items()} def visualize(img_tensor, pred_tensor): img_tensor: [3,H,W] 归一化后的图pred_tensor: [H,W] 类别索引 img img_tensor.cpu().numpy().transpose(1, 2, 0) img (img * [0.229, 0.224, 0.225] [0.485, 0.456, 0.406]) * 255 img img.astype(np.uint8) color_mask np.zeros_like(img) for idx, color in INDEX_TO_COLOR.items(): color_mask[pred_tensor.cpu().numpy() idx] color # 半透明叠加alpha 控制标注层透明度 overlay (img * 0.5 color_mask * 0.5).astype(np.uint8) return overlay模型导出方面如果后续要部署到推理引擎建议在导出前把argmax之前的部分单独包一层输出保持[B, num_classes, H, W]的 logits把后处理留给部署侧。导出时用torch.onnx.export输入尺寸固定成[1, 3, 256, 256]动态轴只开 batch 维避免某些推理引擎对动态 H/W 支持不好。还有一个容易被忽略的点类别数超过 20 时COLOR_TO_INDEX用字典逐类比对会变慢可以改成用np.unique加return_inverse一次性完成映射或者把颜色编码成整数后用查表法。数据量大的时候Dataset 里的cv2.imread也会成为瓶颈可以换成PIL加多进程num_workers但要注意num_workers 0时rgb_mask_to_index里的全局字典在子进程里是拷贝的不会共享状态这点不影响正确性但改映射表后要重启 DataLoader。从那以后我每次接新的多类别分割任务都强制先跑一遍“标签映射自检”随机抽 20 张 mask打印唯一颜色和映射后的索引分布确认没有漏类、没有跳号、没有 255 之外的异常值再开始训练。这个习惯帮我省下了至少三次通宵排查 loss 不降的时间。希望帮到你。本文还有配套的精品资源点击获取