简介这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者提供在PyTorch框架下实现Unet多类别语义分割的完整代码工程可用于医学影像、遥感图像等场景的像素级分类任务。压缩包共46个文件以19个py脚本和24个pyc缓存文件为主另含2个txt与1个json配置整体约69KB涵盖数据加载、自定义变换、模型定义、损失函数、学习率调度、指标计算与训练保存等模块目录按dataloaders、modeling、utils等分层组织便于按功能定位代码。目前已有15274人学习下载说明该实现方案在社区中具有较高参考价值。读者可据此快速搭建可运行的多类别分割流程理解输出通道数与类别数的对应关系、交叉熵损失与IoU等评估指标的落地方式并在此基础上替换编码器、引入注意力机制或调整数据增强策略完成从单类别到多类别的迁移实践。1. 从一堆标注图到能跑的分割模型这个 Unet 项目到底解决了什么如果你手里有一批自己标注的多类别语义分割数据打开 PyTorch 却不知道从哪下手——数据集怎么组织、标签怎么编码、Unet 怎么改输出通道、训练完怎么可视化预测结果这个项目就是冲这些事来的。它把「PyTorch Unet 自建多类别数据集」这条链路完整串了一遍从数据读取、模型搭建、训练循环到推理验证都有对应实现适合刚接触语义分割的开发者照着跑通也适合已经会二分类 Unet、想扩到多类别的人拿来改。多类别和单类别最大的区别不在模型结构而在标签处理、损失函数和评估指标这三处很多人第一次做多类别分割翻车就翻在这里。下面按「项目能干什么 → 怎么跑起来 → 核心代码怎么改」的顺序拆开讲。2. 项目结构与技术选型为什么是 Unet 加这套目录2.1 整体目录与模块职责这个项目走的是典型的「脚本式训练 模块化模型」路线没有套复杂的框架好处是每一块都能单独拎出来改。常见做法是下面这种目录组织unet-multiclass/ ├── datasets/ │ ├── images/ # 原图jpg 或 png │ └── masks/ # 标注掩码png像素值即类别 id ├── src/ │ ├── dataset.py # Dataset 与 DataLoader 封装 │ ├── unet.py # Unet 网络定义 │ ├── train.py # 训练入口 │ ├── predict.py # 单张/批量推理 │ └── utils.py # 指标、可视化、权重保存 ├── checkpoints/ # 训练权重输出 └── requirements.txtimages和masks的文件名一一对应这是最省心的做法读取时按文件名配对不用额外维护一份索引表。masks里每个像素的值就是类别编号0 是背景1、2、3……依次是各个前景类。这一点必须先确认因为很多标注工具导出的掩码是彩色图直接读进来会变成三通道训练时对不上。src下按职责拆文件dataset.py只管数据unet.py只管网络train.py负责把两者接起来。这种拆法的好处是换数据集时只动dataset.py换网络时只动unet.py互不干扰。2.2 技术栈与依赖核心依赖就三样PyTorch 负责张量计算和自动求导torchvision 用来做图像读取和基础变换numpy 和 Pillow 处理掩码和可视化。评估指标如果不想自己写可以加一个专门的分割指标库但自己实现 IoU 也就十几行依赖越少越不容易出环境问题。pip install torch torchvision numpy pillow提示PyTorch 版本和 CUDA 版本要对应装之前先确认显卡驱动支持的 CUDA 版本否则torch.cuda.is_available()会返回 False训练直接掉到 CPU 上速度差几十倍。选 Unet 而不是 DeepLab、SegFormer 这类新结构是因为 Unet 的编码器-解码器加跳跃连接在小数据集上表现稳定参数量可控改起来直观。多类别场景下Unet 的输出层从单通道换成num_classes通道其余结构基本不用动这也是它适合入门的原因。2.3 数据流从磁盘到张量的完整路径一条数据从磁盘到送进网络中间经过这几步读原图 → 归一化 → 转张量读掩码 → 保持单通道 → 转 LongTensor两者做同样的几何变换翻转、裁剪。这里最容易出错的是掩码的几何变换必须和原图完全同步否则标签和图像错位训练出来的模型就是玄学。import torch from torch.utils.data import Dataset from PIL import Image import numpy as np import os class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transformNone): self.img_dir img_dir self.mask_dir mask_dir self.transform transform # 按文件名配对要求 images 和 masks 文件名一致 self.names sorted(os.listdir(img_dir)) def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img Image.open(os.path.join(self.img_dir, name)).convert(RGB) mask Image.open(os.path.join(self.mask_dir, name)).convert(L) # 单通道 img np.array(img, dtypenp.float32) / 255.0 mask np.array(mask, dtypenp.int64) # 类别 id必须是整型 if self.transform: img, mask self.transform(img, mask) img torch.from_numpy(img).permute(2, 0, 1).float() mask torch.from_numpy(mask).long() return img, mask这段代码里有两个关键点。第一mask用convert(L)强制转单通道避免彩色掩码带来的通道错乱第二mask转成int64也就是 LongTensor因为后面交叉熵损失要求标签是整型类别索引用 float 会直接报错。img归一化到 0 到 1 之间是分割任务里最常见的做法比用 ImageNet 均值方差更省事效果也够用。几何变换部分如果自己写增强翻转和裁剪要同时作用在img和mask上def joint_transform(img, mask): # 随机水平翻转图像和掩码必须同步 if np.random.rand() 0.5: img np.fliplr(img).copy() mask np.fliplr(mask).copy() return img, masknp.fliplr之后加.copy()是为了让数组内存连续否则转张量时可能报负步长错误这个坑不少人踩过。3. 模型搭建与训练多类别 Unet 的关键改动点3.1 Unet 结构拆解与输出通道设置Unet 分三部分下采样编码器、上采样解码器、跳跃连接。编码器每下一层通道翻倍解码器每上一层通道减半跳跃连接把编码器同层的特征拼到解码器上补回下采样丢掉的细节。多类别改造只发生在最后一层卷积输出通道从 1 改成num_classes。import torch.nn as nn import torch 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_classes4): super().__init__() self.down1 DoubleConv(in_ch, 64) self.down2 DoubleConv(64, 128) self.down3 DoubleConv(128, 256) self.pool nn.MaxPool2d(2) self.up1 nn.ConvTranspose2d(256, 128, 2, stride2) self.conv1 DoubleConv(256, 128) self.up2 nn.ConvTranspose2d(128, 64, 2, stride2) self.conv2 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)) u1 self.up1(d3) u1 self.conv1(torch.cat([u1, d2], dim1)) # 跳跃连接 u2 self.up2(u1) u2 self.conv2(torch.cat([u2, d1], dim1)) return self.out(u2)num_classes这个参数是整套代码里最需要盯住的地方。它必须等于你数据集里实际类别数含背景。比如背景加三类目标就是 4。设小了会漏类设大了会多出永远学不会的空类评估指标被拉低。self.out用 1x1 卷积把通道压到类别数输出的是每个像素在每个类别上的 logits不需要手动 softmax因为交叉熵损失内部会做。3.2 损失函数与类别不平衡处理多类别分割默认用交叉熵损失它把每个像素当成一个分类样本。问题在于背景像素往往占绝大多数模型会倾向于全预测背景也能拿到很低的 loss但前景 IoU 惨不忍睹。常见做法是给每个类别加权重import torch.nn as nn # 假设背景多、前景少给前景更高权重 class_weights torch.tensor([0.2, 1.0, 1.0, 1.0]).cuda() criterion nn.CrossEntropyLoss(weightclass_weights)权重怎么定没有标准答案一个实用办法是统计训练集里每个类别的像素占比权重取占比的倒数再归一化。如果不想算先全设 1 跑一版看哪个类 IoU 特别低再单独调高它的权重。另一个思路是换 Dice Loss 或交叉熵加 Dice 的组合对小目标更友好但收敛会慢一些需要多训几轮。3.3 训练循环与指标监控训练循环本身不复杂关键是每个 epoch 结束后算一次验证集上的 mIoU用它来判断模型有没有真的在学而不是只看 loss。def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0 for img, mask in loader: img, mask img.to(device), mask.to(device) optimizer.zero_grad() pred model(img) # [B, num_classes, H, W] loss criterion(pred, mask) # mask: [B, H, W] loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(loader) torch.no_grad() def evaluate(model, loader, num_classes, device): model.eval() inter torch.zeros(num_classes) union torch.zeros(num_classes) for img, mask in loader: img, mask img.to(device), mask.to(device) pred model(img).argmax(dim1) # 取每个像素最大概率的类别 for c in range(num_classes): p (pred c) t (mask c) inter[c] (p t).sum().item() union[c] (p | t).sum().item() iou inter / (union 1e-6) return iou.mean().item(), iouargmax(dim1)是把[B, num_classes, H, W]的预测压成[B, H, W]的类别图这一步是推理和评估的标配。IoU 计算里加1e-6是防止某个类别在当前 batch 里没出现导致除零。返回的iou是每个类别的 IoU 数组打印出来能直接看出哪个类拖后腿。训练主循环里优化器一般用 Adam学习率 1e-3 起步跑不动再降到 1e-4。权重保存按验证 mIoU 最高的那一轮存而不是存最后一轮这是血泪经验——最后一轮往往已经过拟合了。4. 推理、可视化与常见问题排查4.1 单张推理与结果可视化训练完最想看到的就是模型在自己图片上画出来的分割结果。推理脚本要做三件事加载权重、前向传播、把类别图映射成彩色图叠加到原图上。import numpy as np import torch from PIL import Image # 每个类别对应一种颜色背景用黑色 palette np.array([ [0, 0, 0], # 类别 0 背景 [255, 0, 0], # 类别 1 [0, 255, 0], # 类别 2 [0, 0, 255], # 类别 3 ], dtypenp.uint8) def predict(model, img_path, device, num_classes4): model.eval() img Image.open(img_path).convert(RGB) arr np.array(img, dtypenp.float32) / 255.0 tensor torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).float().to(device) with torch.no_grad(): pred model(tensor).argmax(dim1).squeeze(0).cpu().numpy() color palette[pred] # 类别图转彩色 overlay (arr * 255 * 0.5 color * 0.5).astype(np.uint8) return overlaypalette的长度必须和num_classes一致否则索引越界。叠加时用 0.5 的透明度既能看清原图纹理又能看出分割区域比纯色掩码直观得多。如果发现预测结果全是背景先别急着改模型检查一下输入归一化方式和训练时是否一致这是最常见的推理翻车点。4.2 避坑与常见问题排查现象一训练 loss 一直不降或者降到某个值就卡住。原因多半是标签值超出了num_classes范围比如掩码里有像素值 5但num_classes设的是 4交叉熵遇到非法类别索引会算出异常梯度。解决办法是训练前统计一遍所有掩码的唯一值确认最大值小于num_classes超出的要么重新映射要么修正标注。现象二验证集 mIoU 很高但预测图一塌糊涂。这通常是数据泄漏训练集和验证集用了同一批图或者验证集划分时没打乱。检查划分逻辑确保两张集合的图片完全不重叠。另一个可能是验证时忘了model.eval()BatchNorm 还在用 batch 统计量结果不稳定。现象三某个类别 IoU 始终为 0。先看这个类在训练集里有多少像素如果总共就几百个像素模型根本学不到。解决办法是过采样含该类的图或者调高它的损失权重。也可能是标注时这个类的颜色和背景太接近标注本身就不可靠。现象四显存爆了batch size 只能设 1。Unet 在高分辨率输入下显存占用很大。常见做法是把输入裁剪成 256x256 或 512x512 的 patch 训练推理时再滑窗拼回原图。另一个办法是减小编码器通道基数比如从 64 起步改成 32 起步参数量降一半精度损失有限。现象五预测结果边缘毛糙、有锯齿。这是上采样用转置卷积的固有现象。可以换成双线性插值上采样加卷积边缘会平滑一些。也可以在损失里加一点边界权重但改动较大入门阶段先把主体跑通更重要。4.3 训练配置参数速查参数常见取值说明num_classes数据集类别数含背景设错直接导致训练失败batch_size4 到 16看显存太小梯度不稳lr1e-3 到 1e-4Adam 起步 1e-3epoch50 到 200看验证 mIoU 是否还在涨input_size256 或 512越大越吃显存class_weights按像素占比倒数类别不平衡时必调这张表里的值不是死的但num_classes和class_weights这两个必须根据自己数据算抄别人的数值大概率出问题。5. 二次开发把这份代码改成你自己的分割任务拿到这套代码后真正花时间的不是跑通而是改成适配自己数据的版本。我一般按这个顺序动刀先改dataset.py里的读取逻辑再确认num_classes然后调损失权重最后才考虑换网络结构。第一步替换数据读取。如果你的掩码不是「像素值即类别 id」的格式而是每个类一个二值图那要在__getitem__里把它们合并成一张类别图def merge_masks(mask_list): # mask_list: 每个类一张二值图背景为 0 merged np.zeros_like(mask_list[0], dtypenp.int64) for i, m in enumerate(mask_list, start1): merged[m 0] i # 类别从 1 开始0 留给背景 return merged这段逻辑的关键是类别编号从 1 开始0 永远留给背景这样和num_classes含背景的约定一致。合并顺序决定了类别 id改的时候要同步改palette。第二步验证类别分布。写个小脚本统计训练集里每个类别的像素总数如果最大类和最小类差了几个数量级class_weights就必须调否则小类永远学不出来。import numpy as np from PIL import Image import os counts np.zeros(num_classes) for name in os.listdir(datasets/masks): m np.array(Image.open(os.path.join(datasets/masks, name)).convert(L)) for c in range(num_classes): counts[c] (m c).sum() print(counts / counts.sum()) # 每个类别的像素占比第三步换评估口径。如果任务里某些类特别重要mIoU 会被大量背景稀释可以改成只统计前景类的平均 IoU或者用 F1 分数。改评估函数时注意别把背景算进去否则数字好看但没意义。第四步想提升精度再考虑换编码器。把 Unet 的编码器换成预训练的 ResNet 或 EfficientNet用 ImageNet 权重初始化小数据集上通常能涨几个点。改动方式是保留解码器把编码器替换成torchvision.models里的骨干网络取中间层特征做跳跃连接。这一步改动量不小建议先把基础版跑出基线再动。有个习惯我每次都强制走一遍训练前先拿一张训练图喂进未训练的模型确认输出 shape 是[1, num_classes, H, W]再拿一张验证图确认标签最大值小于num_classes。这两步花不了一分钟但能挡掉后面几小时的无效训练。从那以后我每次开新分割任务都先跑这个检查希望帮到你。本文还有配套的精品资源点击获取