PyTorch Unet多类别语义分割实战：从数据管线到mIoU提升的完整方案-育师

简介：这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者，提供在PyTorch框架下实现Unet多类别语义分割的完整代码工程，可应用于医学影像、遥感图像等需要像素级分类的场景。压缩包共46个文件，以19个py源码与24个pyc编译文件为主，另含少量txt与json配置，整体约69KB，涵盖数据加载与增强、Unet建模、损失函数、学习率调度、评价指标、训练与可视化等模块，目录划分清晰，便于按功能定位代码。目前已有15274人学习下载，说明其在同类教程中具有较高参考价值。读者可据此搭建从数据预处理、模型训练到IoU等指标评估的完整流程，并借助曲线记录与分割结果对比图直观排查问题，快速将Unet迁移到自己的多类别数据集上。

1. 从一张“跑不通”的训练日志说起：这套 PyTorch Unet 多类别分割方案到底解决什么

如果你手头有一份自己标注的多类别数据集，兴冲冲把 Unet 代码 clone 下来，改完num_classes一跑，大概率会遇到两种结局：要么 loss 死活不降，要么预测图全糊成一类。问题往往不在 Unet 本身，而在“多类别”这三个字带来的连锁反应——标签编码、损失函数、评估指标、类别不平衡，每一环都和二分类不一样。这套 PyTorch 下实现 Unet 多类别语义分割的资源，核心价值就是把这些环节串成一条能复现的链路：从自定义 Dataset 读多通道 mask，到多类别 CrossEntropy + Dice 组合损失，再到按类别统计的 mIoU 验证。它适合已经跑通过二分类分割、想把手头多类别数据真正训起来的从业者，也适合需要一套可改可扩的分割基线的人。下面我按自己拆包复现的顺序，把关键参数和翻车点讲透。

2. 多类别 Unet 的数据管线：Dataset 与标签编码怎么搭

多类别分割翻车，八成翻在数据这一层。二分类时 mask 是 0/1，直接乘个 255 就能看；多类别时 mask 里是 0、1、2、3……这些整数，一旦被当成灰度值做归一化或插值，标签就废了。所以数据管线的第一原则是：图像走浮点归一化，mask 走整数、绝不插值。

2.1 自定义 Dataset 的骨架与同步增强

我一般会写一个继承Dataset的类，把图像和 mask 成对读进来，增强用 albumentations 同步做。关键是 mask 的InterpolationMode.NEAREST，以及不要对 mask 做 Normalize。

import os import numpy as np import torch from torch.utils.data import Dataset import albumentations as A from albumentations.pytorch import ToTensorV2 from PIL import Image class MultiClassSegDataset(Dataset): def __init__(self, img_dir, mask_dir, img_size=256, train=True): self.img_dir = img_dir self.mask_dir = mask_dir self.names = sorted(os.listdir(img_dir)) # 训练增强：图像和 mask 必须同步 if train: self.tf = A.Compose([ A.Resize(img_size, img_size), A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.3), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), # 只对图像 ToTensorV2(), ]) else: self.tf = A.Compose([ A.Resize(img_size, img_size), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ]) def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = np.array(Image.open(os.path.join(self.img_dir, name)).convert("RGB")) mask = np.array(Image.open(os.path.join(self.mask_dir, name)).convert("L")) # mask 保持整数类别，不参与归一化 aug = self.tf(image=img, mask=mask) img_t = aug["image"] # float32, C,H,W mask_t = aug["mask"].long() # int64, H,W return img_t, mask_t

逻辑说明：A.Normalize只作用在image上，mask 传进去只是为了跟着做几何变换（Resize、Flip），它的像素值不会被改。mask_t转成long()是给CrossEntropyLoss用的，后者要求 target 是 int64 且取值在[0, num_classes-1]。参数上，img_size建议先固定 256 或 512，多类别下显存吃紧时优先降 batch 而不是降分辨率，因为小目标类别对分辨率很敏感。

2.2 标签编码：从彩色 mask 到类别索引

很多人拿到的标注是彩色 PNG，比如背景黑、类别一红、类别二绿。这种 mask 不能直接喂网络，必须先映射成 0/1/2 的索引图。常见做法是维护一张颜色到类号的查找表。

# 颜色 -> 类别索引，按你的标注规范改 COLOR_MAP = { (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 类别一 (0, 255, 0): 2, # 类别二 (0, 0, 255): 3, # 类别三 } def rgb_mask_to_index(mask_rgb): """把 H,W,3 的彩色 mask 转成 H,W 的类别索引""" h, w, _ = mask_rgb.shape index = np.zeros((h, w), dtype=np.int64) for color, cls in COLOR_MAP.items(): match = np.all(mask_rgb == np.array(color), axis=-1) index[match] = cls return index

逻辑说明：逐颜色做全图匹配再赋值，简单可靠，适合类别数不多的场景。参数上要注意COLOR_MAP必须和标注工具里的调色板完全一致，差一个像素值就整类丢失。如果类别多，可以改成向量化的最近邻匹配，但多数项目十几类以内，这个循环够用。转换完建议随机抽几张可视化确认，别等训完才发现某一类全是背景。

2.3 类别不平衡：采样与权重的取舍

多类别数据集里，背景往往占 80% 以上，小类别可能不到 1%。直接训，网络会倾向于全预测背景，accuracy 看着很高，mIoU 惨不忍睹。常见做法有两类：一是WeightedRandomSampler按类别出现频率反比采样，二是给损失函数加类别权重。我一般先上损失权重，改动小、可控。

from torch.utils.data import WeightedRandomSampler def build_sampler(dataset, num_classes): # 统计每张图里出现过的类别，给稀有类别更高的采样权重 weights = [] for i in range(len(dataset)): _, mask = dataset[i] present = torch.unique(mask) # 出现类别越稀有，权重越高 w = sum(1.0 / (num_classes - int(c)) for c in present) weights.append(w) weights = torch.tensor(weights, dtype=torch.double) return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)

逻辑说明：这里用了一个简化的稀有度权重，类别号越大（通常越稀有）权重越高。实际项目里更稳的做法是先统计全量像素的类别频率，再取反比。参数上replacement=True表示有放回采样，稀有样本会被重复抽到。注意采样器和损失权重不要同时拉满，否则容易过拟合小类别，我一般二选一，或者采样器为主、损失权重为辅。

3. Unet 网络与多类别损失：结构改动和 loss 组合

数据通了，接下来是模型和损失。Unet 的编码器-解码器加跳跃连接结构本身不用大改，多类别的关键在输出通道数和损失函数。这一章把网络输出、损失组合、以及训练循环里容易忽略的细节讲清楚。

3.1 输出通道与上采样：num_classes 怎么落

Unet 最后输出的是每个像素的类别 logits，通道数等于类别数（含背景）。常见错误是输出通道写成 1 然后用 Sigmoid，那是二分类的写法。多类别必须是num_classes通道，且推理时走argmax而不是阈值。

import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.net = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.net(x) class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=4): 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, stride=2) self.up_conv1 = DoubleConv(256, 128) self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2) self.up_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.up_conv1(torch.cat([u1, d2], dim=1)) u2 = self.up2(u1) u2 = self.up_conv2(torch.cat([u2, d1], dim=1)) return self.out(u2) # 返回 logits，不接 softmax

逻辑说明：self.out用 1x1 卷积把 64 通道压到num_classes，输出是 logits，不接 softmax，因为CrossEntropyLoss内部自带 log_softmax。参数上num_classes必须等于你标签索引的最大值加一，比如标签是 0/1/2/3，那就是 4。跳跃连接用torch.cat拼接，注意通道数对齐，up_conv1的输入是 128（上采样）+128（编码器）=256。

3.2 CrossEntropy + Dice 组合损失

多类别下单用 CrossEntropy 容易被背景主导，单用 Dice 又可能训练不稳。我一般用两者加权组合，CrossEntropy 负责像素级分类，Dice 负责类别区域重叠。

class DiceLoss(nn.Module): def __init__(self, num_classes, smooth=1e-6): super().__init__() self.num_classes = num_classes self.smooth = smooth def forward(self, logits, target): # logits: B,C,H,W target: B,H,W probs = torch.softmax(logits, dim=1) target_onehot = torch.nn.functional.one_hot( target, self.num_classes).permute(0, 3, 1, 2).float() dims = (0, 2, 3) inter = (probs * target_onehot).sum(dims) union = probs.sum(dims) + target_onehot.sum(dims) dice = (2 * inter + self.smooth) / (union + self.smooth) return 1 - dice.mean() class ComboLoss(nn.Module): def __init__(self, num_classes, ce_weight=1.0, dice_weight=1.0): super().__init__() self.ce = nn.CrossEntropyLoss() self.dice = DiceLoss(num_classes) self.ce_w = ce_weight self.dice_w = dice_weight def forward(self, logits, target): return self.ce_w * self.ce(logits, target) + self.dice_w * self.dice(logits, target)

逻辑说明：DiceLoss先把 logits 过 softmax 得到概率，再把 target 转 one-hot，按类别算 Dice 后取平均。smooth防止除零。ComboLoss里两个权重是核心调参点，我一般从ce_weight=1.0, dice_weight=1.0起步，如果小类别还是学不出来，把 dice 权重提到 2.0 试试。注意CrossEntropyLoss的 target 是B,H,W的 int64，别传成 one-hot。

3.3 训练循环与验证指标

训练循环里除了常规的前向反向，多类别要特别关注验证指标。accuracy 在多类别不平衡下没有参考价值，必须看 mIoU 和每类 IoU。

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() logits = model(img) loss = criterion(logits, mask) 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() # 混淆矩阵：行是真实类，列是预测类 conf = torch.zeros(num_classes, num_classes, dtype=torch.long) for img, mask in loader: img, mask = img.to(device), mask.to(device) pred = model(img).argmax(dim=1) # 关键：argmax 取类别 for t, p in zip(mask.view(-1), pred.view(-1)): conf[t.long(), p.long()] += 1 ious = [] for c in range(num_classes): tp = conf[c, c].item() fp = conf[:, c].sum().item() - tp fn = conf[c, :].sum().item() - tp iou = tp / (tp + fp + fn + 1e-6) ious.append(iou) return sum(ious) / num_classes, ious

逻辑说明：evaluate用混淆矩阵累计所有像素，再逐类算 IoU，最后取平均得到 mIoU。argmax(dim=1)是多类别推理的标准动作，别用阈值。参数上num_classes要和模型输出一致。返回的ious列表能直接看出哪一类拖后腿，比一个 mIoU 数字有用得多。训练时建议每个 epoch 存一次验证结果，方便回溯。

4. 多类别分割避坑清单：五个我踩过的坑

这一章全是血泪经验，每条按现象、原因、解决写。多类别分割的坑大多不在模型，而在数据和评估的细节里。

4.1 现象：loss 一直不降，预测全是一类

原因：mask 被当成灰度图做了归一化或双线性插值，类别索引被插值成了小数，CrossEntropyLoss拿到非法 target 后行为异常。解决：确认 mask 全程走NEAREST插值、不参与 Normalize，并在 Dataset 里断言mask.max() < num_classes。

4.2 现象：训练 accuracy 很高，mIoU 却只有 0.1

原因：背景占比过高，网络学会了全预测背景。解决：上WeightedRandomSampler或给损失加类别权重，同时把评估指标从 accuracy 换成 mIoU，别被虚高的 accuracy 骗了。

4.3 现象：某一类 IoU 始终为 0

原因：该类在标签里的索引和COLOR_MAP对不上，或者该类样本太少根本没被采到。解决：先可视化几张 mask 确认类别索引正确，再统计每类像素占比，占比低于千分之一的类别考虑合并或过采样。

4.4 现象：验证集 loss 正常，但预测图边缘全是噪点

原因：上采样用了双线性插值再 argmax，边界处概率被平滑，argmax 后出现锯齿。解决：推理时保持 logits 分辨率，最后一步再 argmax；如果显存允许，别在中间层做插值。

4.5 现象：换一批数据后 mIoU 暴跌

原因：新数据的类别定义或颜色映射和训练时不一致，或者图像归一化的 mean/std 没跟着换。解决：把COLOR_MAP和归一化参数做成配置文件，训练和推理共用同一份，别硬编码在代码里。

5. 进阶技巧：把 mIoU 从 0.6 推到 0.75 的几个实操动作

基线跑通后，想再往上提，靠的不是换更深的网络，而是把细节抠到位。我一般按下面这个顺序调，每步都验证一次 mIoU，别一次改一堆。

第一，损失权重网格搜索。ce_weight和dice_weight的组合对结果影响很大，我会固定其他条件，试(1,1)、(1,2)、(2,1)、(1,3)几组，选验证集 mIoU 最高的。别在训练集上选，会过拟合。

第二，学习率调度。多类别分割对学习率敏感，我一般用CosineAnnealingLR，初始 lr 设 1e-3，配合 warmup 前 5 个 epoch。如果 loss 前期震荡，把初始 lr 降到 3e-4。

from torch.optim.lr_scheduler import CosineAnnealingLR optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4) scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)

逻辑说明：AdamW的 weight_decay 比 Adam 更规范，T_max设成总 epoch 数，eta_min是学习率下限。每个 epoch 后调scheduler.step()。参数上weight_decay从 1e-4 起步，过拟合明显时加到 1e-3。

第三，测试时增强（TTA）。推理时对图像做水平翻转，两次预测取平均再 argmax，通常能涨 1 到 2 个点 mIoU，代价是推理时间翻倍。

@torch.no_grad() def predict_tta(model, img, device): model.eval() img = img.to(device) logits = model(img) logits_flip = model(torch.flip(img, dims=[3])) logits_flip = torch.flip(logits_flip, dims=[3]) return ((logits + logits_flip) / 2).argmax(dim=1)

逻辑说明：翻转预测后再翻回来，和原预测平均，能抵消一部分位置偏差。注意只对水平翻转做，垂直翻转对某些数据集反而有害。

第四，难例挖掘。把验证集里 IoU 最低的几类挑出来，单独看它们的预测图，往往是标注边界模糊或类别定义重叠。这时候与其调模型，不如回去修标注，收益更直接。