PyTorch UNet多类别语义分割实战与避坑

简介：这份PyTorch Unet多类别语义分割项目面向图像分析方向的学习者与开发者，适用于医学影像、遥感图像等需要逐像素分类的场景，重点解决多类别输出通道设置、训练流程组织与结果评估等问题。压缩包共46个文件，约69KB，以19个py源码和24个pyc缓存文件为主，另含2个txt与1个json配置说明，源码覆盖数据加载、模型定义、损失函数、学习率调度、指标统计、权重计算及训练日志保存等环节，目录按功能拆分，便于按模块阅读和二次修改。目前已有15287人学习下载，说明其在Unet实战学习中具有一定参考价值。读者可从中获得一套多类别分割代码框架，理解编码器—解码器与跳跃连接的落地方式，并借助损失、IoU等评估模块和曲线记录排查训练问题，也可在现有结构上替换数据集或调整网络，用于课程实验与项目原型验证。

1. 从一个反直觉的失败说起

第一次拿到一份多类别标注数据时，很多人的第一反应是把num_classes从 2 改成 5，看 loss 在降就以为大功告成；结果推理图出来，整张图被预测成背景。PyTorch 下用 UNet 做自己数据集的语义分割，坑几乎都不在网络结构本身，而在掩码的编码方式、损失的类别权重、验证指标的口径这三处。这份资源围绕一套能直接跑的 UNet 多类别分割实现展开，从自定义数据集出发，把数据管线、网络、损失、训练与推理串成一条完整链路。适合已经能用 PyTorch 搭分类网络、想把语义分割真正落到自己标注数据上的开发者，也适合跑通过二分类分割、正被多类别问题卡住的人。

2. 项目骨架与选型：PyTorch + UNet 这套组合怎么搭才不返工

2.1 为什么多类别分割仍然值得先上 UNet

UNet 的结构不复杂：编码器做四次下采样拿语义，解码器做四次上采样恢复分辨率，每一级用一条 skip connection 把编码器的高分辨率特征拼到解码器对应层。这套设计的价值在多类别场景里体现得特别明显——浅层特征负责把类间边界勾出来，深层特征负责判断这块区域到底属于哪一类。自己做标注的数据集，误差往往集中在类与类挨着的地方，比如路面和人行道、病灶和正常组织、缺陷边缘和背景，skip connection 恰好把边界信息一路带到了输出层。

数据量是个绕不开的现实。自己标的数据集通常几百到几千张，ImageNet 预训练主干在小样本上确实有帮助，但原版 UNet 从零训在千张级别也能收敛，前提是增强策略和损失函数得跟上。参数量上，base_channels=64的原版 UNet 大约三千万参数，512×512 输入、batch_size=4、打开混合精度，一张 8G 显存的卡基本能吃满；只有 6G 的话把base_channels降到 32，或者把输入裁到 384，都属于常见做法。

选型上没必要一上来就追新。三种方案在自建多类别数据集上的实际差异大致如下。

我一般的判断是：类别数在 3 到 10 之间、数据量不到一万张、需要快速迭代看效果，UNet 是最省心的起点；等 mIoU 卡在一个平台上不去，再换主干或者上多尺度，收益才明显。

2.2 目录结构与依赖清单

工程目录按职责切分，避免所有东西堆在一个train.py里。下面这套结构是从几个实际项目里收敛出来的，改数据集时基本只动configs和datasets两处。

seg_project/ ├── configs/ │ └── unet_multiclass.yaml # 超参、路径、类别数集中在这里 ├── datasets/ │ ├── images/{train,val}/ # 原图，jpg/png 混放也行 │ └── masks/{train,val}/ # 灰度索引掩码，与图像同名 ├── src/ │ ├── data/ │ │ ├── dataset.py # Dataset 与 collate │ │ └── transforms.py # 图像与掩码同步增强 │ ├── models/ │ │ └── unet.py # DoubleConv / Up / UNet │ ├── losses/ │ │ └── seg_loss.py # CE + Dice 组合损失 │ ├── engine/ │ │ ├── train.py # 单轮训练 │ │ └── evaluate.py # 验证与混淆矩阵 │ └── utils/ │ ├── metrics.py # mIoU、per-class IoU │ └── visualize.py # 预测结果上色 ├── train.py # 入口，读配置、建模型、跑循环 ├── predict.py # 单图/批量推理 └── requirements.txt

依赖不需要多，下面是够用的最小集合。

# requirements.txt torch torchvision albumentations opencv-python-headless numpy pillow pyyaml tqdm tensorboard

albumentations是这套流程里最值得留的第三方库，它最大的价值是让图像和掩码走同一套随机参数，几何变换不会错位。opencv-python-headless在服务器上没有图形界面也能装，避免libGL报错这种和算法完全无关的翻车。

2.3 配置集中管理：把会变的东西全部外置

多类别分割最烦的是「换个数据集就要改一堆硬编码」。把类别数、忽略标签、图像尺寸、学习率这些全部提到 YAML 里，训练脚本只读配置，改数据集时不动代码。

# configs/unet_multiclass.yaml data: root: ./datasets num_classes: 5 # 背景 + 4 类前景，必须和掩码里的最大 id + 1 对齐 ignore_index: 255 # 未标注区域、边界过渡像素统一标成 255 image_size: 512 batch_size: 4 num_workers: 4 model: in_channels: 3 base_channels: 64 # 显存不够就降到 32 bilinear: false # false 用转置卷积，true 用双线性插值，后者参数更少 train: epochs: 120 lr: 0.0001 weight_decay: 0.00001 amp: true # 混合精度，8G 卡上基本必开 val_interval: 5 # 每 5 轮验一次，省时间 save_dir: ./runs/unet_mc

num_classes和掩码的关系必须先确认清楚：如果掩码里前景是 1、2、3、4，背景是 0，那num_classes=5；如果标注工具把背景标成 255，那背景得单独映射回 0，否则一训练就索引越界。ignore_index用 255 是个约定俗成的选择，好处是「未标注」和「背景」能彻底分开，坏处是掩码必须存成能放下 255 的格式，PNG 单通道刚好够用，JPEG 有损压缩会直接把 255 变成 254，这一点后面还会再提。

读取配置的入口很薄，只做三件事：加载 YAML、按device决定是否启用 AMP、把save_dir建出来。

# train.py 的配置加载部分 import argparse, os, yaml, torch def load_config(path): with open(path, "r", encoding="utf-8") as f: cfg = yaml.safe_load(f) # 路径统一转绝对路径，避免从不同目录启动时找不到数据 cfg["data"]["root"] = os.path.abspath(cfg["data"]["root"]) os.makedirs(cfg["train"]["save_dir"], exist_ok=True) return cfg if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", default="configs/unet_multiclass.yaml") args = parser.parse_args() cfg = load_config(args.config) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"device={device}, num_classes={cfg['data']['num_classes']}")

这段代码没什么技术含量，但把路径绝对化这一步能省掉大量「在项目根目录能跑、换个目录就报FileNotFoundError」的沟通成本。cfg一路往下传，函数签名里不再出现任何魔法数字，后面换数据集只改 YAML。

3. 数据管线拆解：多类别掩码从磁盘到 batch 的完整链路

3.1 掩码的三种存法，选错一种后面全乱

多类别掩码存磁盘的方式直接决定了Dataset怎么写。三种主流做法各有各的坑。

第二种是最容易翻车的地方，值得单独说。标注工具导出的调色板 PNG，用看图软件打开是彩色的，但用Image.open(...).convert("L")读，得到的是调色板索引经过亮度公式换算后的灰度值，和类别 id 完全对不上。正确做法是用np.array(Image.open(path))直接读，PIL 对 palette 模式会自动展开成索引数组；如果读到的是三维数组，说明工具把它存成了真彩色，那就得按「颜色 → 类别」的映射表反查，映射表写死在配置里或者单独存个 json。

这套实现统一要求掩码是灰度索引图，理由很简单：体积小、加载快、CrossEntropyLoss直接吃LongTensor的 H×W 标签，不需要任何中间转换。数据准备阶段多花十分钟确认格式，比训练到一半发现类别全错要划算得多。

3.2 Dataset 实现：成对校验、越界保护、dtype 对齐

Dataset要做的事比想象中多。除了读图，还得处理三件容易被忽略的事：图像和掩码文件名必须能一一对应、掩码里出现的类别 id 不能超过num_classes、返回的张量 dtype 必须符合损失函数的要求。

import os import numpy as np import torch from torch.utils.data import Dataset from PIL import Image class MultiClassSegDataset(Dataset): """多类别语义分割数据集：图像与灰度索引掩码同名存放。""" def __init__(self, image_dir, mask_dir, transform=None, num_classes=5, ignore_index=255): self.image_dir = image_dir self.mask_dir = mask_dir self.transform = transform self.num_classes = num_classes self.ignore_index = ignore_index # 只保留成对存在的样本，避免训练到一半才报 FileNotFoundError names = sorted(os.listdir(image_dir)) self.samples = [ n for n in names if os.path.exists(os.path.join(mask_dir, n)) ] dropped = len(names) - len(self.samples) if dropped: print(f"[warn] {dropped} 张图像缺少同名掩码，已跳过") def __len__(self): return len(self.samples) def __getitem__(self, idx): name = self.samples[idx] image = np.array( Image.open(os.path.join(self.image_dir, name)).convert("RGB") ) # 关键：mask 不 convert("L")，直接读原始像素值作为类别 id mask = np.array( Image.open(os.path.join(self.mask_dir, name)), dtype=np.uint8 ) # 越界类别值统一归到 ignore，防止 CE 越界直接崩 mask = np.where(mask < self.num_classes, mask, self.ignore_index) mask = mask.astype(np.int64) if self.transform is not None: out = self.transform(image=image, mask=mask) image, mask = out["image"], out["mask"] # 标准化后的图像已经是 CHW；未标准化时这里兜底转一次 if isinstance(image, np.ndarray): image = torch.from_numpy(image.transpose(2, 0, 1)).float() / 255.0 mask = torch.from_numpy(np.asarray(mask)).long() return image, mask

逻辑上分四步走。第一步扫描目录取交集，把没有同名掩码的样本提前剔除，同时在终端打一条警告，能快速发现标注漏项。第二步读原图并强制转 RGB，三通道输入是所有主流主干和预训练权重的默认预期，灰度图、四通道 RGBA 图混在一起是另一个常见事故源。第三步对掩码做越界裁剪，np.where(mask < num_classes, mask, ignore_index)这一行挡住了绝大多数「训练几轮后突然报IndexError: Target X is out of bounds」的情况，那些越界值通常来自标注工具的导出 bug 或者手工改图时手滑。第四步统一 dtype，图像转float32，掩码保持int64。

参数上，num_classes必须和 YAML 里完全一致，ignore_index默认 255；如果标注体系里 255 是一个有效类别，那这个值必须换掉，比如换成 1000，代价是掩码不能再存成 8 位 PNG，得用 16 位。这是一个很典型的取舍，没标过的区域和真实背景混在一起，模型学到的东西会明显变脏。

DataLoader的配置里，num_workers在 Linux 上给 4 到 8，Windows 上建议 0，多进程 fork 在 Windows 下的问题属于经典血泪经验。pin_memory=True配合non_blocking=True能省一点 H2D 拷贝时间，数据量不大的时候收益有限，但配置成本为零。

3.3 同步增强：几何变换必须同时作用在掩码上

增强对自建小数据集几乎是决定性的。但多类别分割的增强有个硬约束：几何变换的图像和掩码参数必须完全一致，且掩码的插值必须是最近邻，否则边界上会凭空出现 0.5 个类别这种不存在的值。

import albumentations as A from albumentations.pytorch import ToTensorV2 def build_train_transform(image_size=512): return A.Compose([ # 裁一块再缩放，比整图 resize 更能提升尺度鲁棒性 A.RandomResizedCrop( height=image_size, width=image_size, scale=(0.5, 1.0), ratio=(0.75, 1.33), p=1.0, ), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.2), A.RandomRotate90(p=0.5), # 颜色扰动只作用于图像，albumentations 内部会自动区分 A.ColorJitter( brightness=0.2, contrast=0.2, saturation=0.2, hue=0.05, p=0.5 ), A.GaussNoise(p=0.2), A.Normalize( mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), ), ToTensorV2(), ]) def build_val_transform(image_size=512): return A.Compose([ A.Resize(height=image_size, width=image_size), A.Normalize( mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), ), ToTensorV2(), ])

A.Compose在调用时同时接收image和mask两个关键字参数，内部会为几何类变换采样一组随机参数，两个输入共用，这就是它比手写torchvision.transforms更省事的原因。RandomResizedCrop对 mask 默认走最近邻插值，不会产生浮点类别值；如果自己写cv2.resize，interpolation必须显式传cv2.INTER_NEAREST，传默认的双线性会让边界出现类别 3 和 4 之间的中间值，再取整就变成了随机噪声。

验证集只做Resize和Normalize，不能放任何随机变换。这里有个容易忽略的点：训练时用的归一化均值方差必须和验证、推理完全一致，否则模型在验证集上的表现会莫名其妙掉几个点，而且很难定位。把这两组参数抽成模块级常量，训练和推理两边 import 同一份，比复制粘贴稳。

不同版本的 albumentations 在RandomResizedCrop的参数名上略有差异，遇到TypeError: __init__() got an unexpected keyword argument时直接看函数签名，别硬猜。

3.4 类别权重：让稀有类别的梯度不被淹没

多类别数据集天然不均衡，背景可能占 80%，最小的类只有 1%。这时候直接上CrossEntropyLoss，模型很快学会「全预测背景」，loss 还降得挺好看。解决办法是给损失传类别权重，权重按像素频率的反比算，再做上限裁剪防止极端类别主导梯度。

from PIL import Image def compute_class_weights(mask_dir, num_classes, max_ratio=20.0): """按像素频率反比计算类别权重，并裁剪上限。""" counts = np.zeros(num_classes, dtype=np.float64) for name in os.listdir(mask_dir): m = np.array(Image.open(os.path.join(mask_dir, name)), dtype=np.uint8) valid = m < num_classes counts += np.bincount( m[valid].ravel(), minlength=num_classes )[:num_classes] freq = counts / counts.sum() # 频率取倒数，clip 防止除零 weights = 1.0 / np.clip(freq, 1e-6, None) weights = weights / weights.mean() # 均值归一到 1，学习率不用重调 weights = np.clip(weights, 1.0 / max_ratio, max_ratio) return torch.tensor(weights, dtype=torch.float32)

max_ratio=20是个经验值。不裁剪的话，某个只占千分之一像素的类别会拿到上千的权重，训练前期梯度爆炸，loss 直接变nan。裁剪到 20 倍意味着最稀有的类的梯度最多是平均水平的 20 倍，既能让模型注意到它，又不至于把其他类别的学习带偏。算出来的权重存成文件，训练时读进来传给CrossEntropyLoss(weight=...)，别每次启动都重新扫一遍掩码目录，几千张图扫下来也是几十秒。

如果稀有类的样本数本身就很少（比如只有三五张图里有），光靠像素权重不够，还得配WeightedRandomSampler，按「图中是否含稀有类」给整张图加权采样。两个手段叠加时要注意，采样器已经提升了稀有类出现频率，损失权重就不该再给到 20 倍，否则等于双重加权，容易过拟合那几张图。

4. 核心模块拆解：UNet 的编码器、跳跃连接与多类别损失

4.1 DoubleConv：整个网络的积木

UNet 里重复最多的结构就是「两次 3×3 卷积 + BN + ReLU」，把它抽成一个模块，编码器和解码器都能复用。

import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): """(conv3x3 -> BN -> ReLU) x 2""" def __init__(self, in_ch, out_ch, mid_ch=None): super().__init__() mid_ch = mid_ch or out_ch self.block = nn.Sequential( # bias=False：后面紧跟 BN，偏置会被减掉，纯属浪费参数 nn.Conv2d(in_ch, mid_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(mid_ch), nn.ReLU(inplace=True), nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x)

padding=1配kernel_size=3保证空间尺寸不变，这是 UNet 能靠 shape 直接 concat 的前提。bias=False是个小优化，卷积后接 BN 时偏置项会被 BN 的均值减法抵消，留着只是多占显存。mid_ch参数是给双线性上采样版本用的，那种情况下通道数变化路径和转置卷积不一样，留个口子比事后改结构方便。

inplace=True的 ReLU 省显存，但在某些需要保留中间激活做可视化的场景下会出问题，如果后面要接 Grad-CAM 之类的分析工具，把它改成False。

4.2 上采样与跳跃连接：尺寸对不齐是最高频的报错

torch.cat要求两个张量的空间尺寸完全一致。输入尺寸不是 16 的倍数时，四次下采样再四次上采样，尺寸会差 1 到 3 个像素，直接抛RuntimeError: Sizes of tensors must match。

class Up(nn.Module): """上采样 + 与编码器特征拼接 + 两次卷积""" def __init__(self, in_ch, skip_ch, out_ch, bilinear=False): super().__init__() if bilinear: # 双线性插值不引入参数，显存紧张时优先选这个 self.up = nn.Upsample( scale_factor=2, mode="bilinear", align_corners=False ) conv_in = in_ch else: # 转置卷积可学习，小数据集上通常比双线性略好 self.up = nn.ConvTranspose2d( in_ch, in_ch // 2, kernel_size=2, stride=2 ) conv_in = in_ch // 2 self.conv = DoubleConv(conv_in + skip_ch, out_ch) def forward(self, x, skip): x = self.up(x) # 尺寸兜底：奇数尺寸下采样再上采样会差 1 像素 if x.shape[-2:] != skip.shape[-2:]: x = F.interpolate( x, size=skip.shape[-2:], mode="bilinear", align_corners=False, ) x = torch.cat([skip, x], dim=1) # 通道维拼接 return self.conv(x)

F.interpolate那三行是这套代码里最值钱的防御。它把「输入必须是 16 的倍数」这个隐性约束变成了软约束，任意尺寸的图丢进来都能跑完。代价是尺寸不整除时上采样结果略有偏差，实际影响很小，相比于训练一开始就崩，这个代价完全值得。

bilinear=True和False的选择：双线性版本参数少、显存省，但表达力弱一点；转置卷积可学习，代价是容易在输出上产生棋盘格伪影。数据量上千张、类别边界清晰时，两者差距不明显，显存优先选双线性。

4.3 完整 UNet 与输出头：为什么不接 softmax

class UNet(nn.Module): def __init__(self, in_channels=3, num_classes=5, base_channels=64, bilinear=False): super().__init__() c = base_channels factor = 2 if bilinear else 1 self.inc = DoubleConv(in_channels, c) self.down1 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(c, c * 2)) self.down2 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(c * 2, c * 4)) self.down3 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(c * 4, c * 8)) self.down4 = nn.Sequential( nn.MaxPool2d(2), DoubleConv(c * 8, c * 16 // factor) ) self.up1 = Up(c * 16, c * 8, c * 4, bilinear) self.up2 = Up(c * 4, c * 2, c * 2, bilinear) self.up3 = Up(c * 2, c, c, bilinear) self.up4 = Up(c, c, c, bilinear) # 1x1 卷积把通道数压到类别数，输出 logits self.outc = nn.Conv2d(c, num_classes, kernel_size=1) def forward(self, x): x1 = self.inc(x) # H/1 x2 = self.down1(x1) # H/2 x3 = self.down2(x2) # H/4 x4 = self.down3(x3) # H/8 x5 = self.down4(x4) # H/16 x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) return self.outc(x) # [B, num_classes, H, W]，不接 softmax

输出层不接 softmax 是个必须强调的点。nn.CrossEntropyLoss内部已经做了log_softmax，如果模型输出再过一遍 softmax，等于把概率压缩了两次，梯度会变得非常小，训练慢到怀疑人生。推理时想要概率图，在argmax之前手动torch.softmax(logits, dim=1)。

base_channels决定了整个网络的宽度。64 是原版设定，32 是显存紧张时的降级方案，128 一般没必要，自建数据集撑不起这么大的容量，反而更容易过拟合。

4.4 损失函数：CE 管像素，Dice 管重叠

多类别分割里，交叉熵和 Dice 是互补的。CE 对每个像素独立算损失，类别不均衡时会被多数类主导；Dice 直接优化预测区域和真实区域的重叠度，对前景稀少的情况更敏感。两个按 1:1 加权组合，收敛速度和最终 mIoU 通常都比单用 CE 好。

class DiceLoss(nn.Module): """多类别 soft Dice，作用在 softmax 概率上，忽略 ignore_index。""" def __init__(self, ignore_index=255, smooth=1.0): super().__init__() self.ignore_index = ignore_index self.smooth = smooth def forward(self, logits, target): num_classes = logits.shape[1] prob = torch.softmax(logits, dim=1) valid = target != self.ignore_index target_clean = target.clone() target_clean[~valid] = 0 one_hot = F.one_hot( target_clean, num_classes ).permute(0, 3, 1, 2).float() # 被忽略的位置在 prob 和 one_hot 上都清零，不参与分子分母 mask = valid.unsqueeze(1).float() prob = prob * mask one_hot = one_hot * mask dims = (0, 2, 3) inter = (prob * one_hot).sum(dims) union = prob.sum(dims) + one_hot.sum(dims) dice = (2 * inter + self.smooth) / (union + self.smooth) return 1 - dice.mean() def build_criterion(class_weights=None, ignore_index=255): ce = nn.CrossEntropyLoss( weight=class_weights, ignore_index=ignore_index ) dice = DiceLoss(ignore_index=ignore_index) def criterion(logits, target): return 0.5 * ce(logits, target) + 0.5 * dice(logits, target) return criterion

ignore_index在 Dice 里的处理比 CE 麻烦。CE 只需要传个参数，框架内部会跳过对应位置；Dice 得手动把无效像素在预测概率和 one-hot 标签上都清零，否则这些位置会污染分子和分母。smooth=1.0是拉普拉斯平滑，防止某个类别在 batch 内一个像素都没有时出现 0/0。

如果发现验证集 mIoU 抖动很大，把 Dice 的权重降到 0.3 试试，Dice 对小 batch 比较敏感，batch_size=4 时方差明显。

4.5 训练循环与验证指标：混淆矩阵是唯一可信的口径

训练循环本身不复杂，混合精度那块要写对，否则会出现scaler未定义的AttributeError。

def train_one_epoch(model, loader, optimizer, criterion, scaler, device, ignore_index=255): model.train() total_loss = 0.0 for images, masks in loader: images = images.to(device, non_blocking=True) masks = masks.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) # AMP 只在 CUDA 上启用，CPU 下 autocast 会报错 use_amp = scaler is not None and device.type == "cuda" with torch.autocast( device_type=device.type, dtype=torch.float16, enabled=use_amp ): logits = model(images) loss = criterion(logits, masks) if use_amp: scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) return total_loss / len(loader.dataset)

optimizer.zero_grad(set_to_none=True)比默认的置零更省一次显存写操作，在 PyTorch 新版本里是推荐写法。AMP 那段的device_type传device.type而不是硬编码"cuda"，是为了同一份代码在 CPU 上也能跑通，虽然慢但至少不会启动就报错。

验证指标一定用混淆矩阵算，别用「逐 batch 算 IoU 再平均」，那个值会被小目标多的 batch 严重拉偏。

@torch.no_grad() def update_confusion(conf, logits, target, num_classes, ignore_index=255): pred = logits.argmax(dim=1) valid = target != ignore_index pred, target = pred[valid], target[valid] idx = target * num_classes + pred conf += torch.bincount( idx, minlength=num_classes ** 2 ).reshape(num_classes, num_classes) return conf def iou_from_confusion(conf): conf = conf.float() inter = torch.diag(conf) union = conf.sum(0) + conf.sum(1) - inter iou = inter / union.clamp(min=1) miou = torch.nanmean(iou) return miou, iou

混淆矩阵的行是真实类别、列是预测类别，对角线是预测对的像素数。iou会对每个类别单独算一个值，打印出来能立刻看出是哪个类拖后腿——通常是最小的那个类，也可能是相邻两类互相混淆。nanmean处理的是验证集里某个类完全没出现的情况，用普通mean会得到nan，整个指标就没法看了。

5. 避坑与排查：多类别分割跑不通的五个典型现场

5.1 数据与掩码类：训练一开始就报错或者类别全错

现象一：启动就抛IndexError: Target 255 is out of bounds。

原因是掩码里的 255 被当成有效类别传给了CrossEntropyLoss，而num_classes=5，索引范围只有 0 到 4。ignore_index必须在损失函数里显式传进去，默认值是 -100，不是 255。解决方式就是在build_criterion里把ignore_index=255一路传到底，同时Dataset里做一次越界裁剪，双保险。

现象二：训练 loss 正常降，但看着验证集的预测全是背景。

大概率是掩码读错了。用Image.open(...).convert("L")读调色板 PNG 会得到错误的灰度值，这些值和真实类别 id 完全对不上，模型学到的是一个假的映射。排查方法很简单：挑一张掩码，np.unique(np.array(Image.open(path)))看一眼输出，正常的索引图应该只输出0到num_classes-1加上可能的 255；如果输出一堆散乱的值，就是这个坑。另外 JPEG 有损压缩也会把 255 变成 254，掩码必须存 PNG 或 BMP。

现象三：某些图训练到一半报文件不存在。

图像和掩码的文件名扩展名不一致，比如图是001.jpg，掩码是001.png。Dataset里用os.path.splitext按主干名匹配，别按完整文件名匹配。这个问题在小数据集上不明显，等数据量上来、标注分批交付之后才会集中爆发。

5.2 训练与指标类：loss 好看但模型不可用

现象四：loss 从 1.2 降到 0.1，mIoU 一直在 0.2 上下不动。

典型的类别不均衡。背景占九成，模型学会全预测背景，CE 会给出一个很低的 loss，但 mIoU 里前景类的 IoU 全是 0。排查方式是打印 per-class IoU，如果只有背景那一项接近 1，其他都是 0，基本可以确认。解决分两步：给 CE 传类别权重，并把 Dice 损失的权重视情况提到 0.5 以上。别只调学习率，调不动。

现象五：验证 loss 每轮都在降，但隔几轮看 mIoU 反而在掉。

验证集样本数太少，指标方差大。几百张验证图、batch_size=4，一轮下来只有几十个 batch，mIoU 波动 ±0.05 很正常。别盯着单轮指标调参，把每轮的混淆矩阵累加到一个总的 conf 里，训练结束统一算一次；或者用滑动平均看趋势。还有一个可能：验证集的归一化参数和训练集不一致，这种情况 loss 也会降，但模型输出会系统性偏移。

5.3 推理与部署类：训练好不等于跑得对

现象六：单图推理出来的掩码尺寸和原图对不上。

训练时做了Resize到 512×512，推理时直接argmax输出的就是 512×512，忘了插值回原图尺寸。修法是F.interpolate到原图 H、W，mode用"bilinear"拿到概率图再argmax，别拿 logits 直接 resize，那样会破坏类别间的相对大小关系。

现象七：模型在本地推理正常，换到服务端结果全乱。

归一化参数在训练端和服务端写了两份，其中一份的均值方差顺序写反了（BGR 当成 RGB），或者忘了除 255。这类问题最阴，因为不报错。我的习惯是把mean、std、resize尺寸打包成一个preprocess.json，训练和推理都读同一个文件，从物理上消除两边不一致的可能。

6. 二次开发：换主干、加 TTA 与导出 ONNX 的落地路径

改这套代码最常做的三件事：换数据集、换网络、加推理增强。

换数据集只需要动三处：datasets目录下的文件、YAML 里的num_classes和ignore_index、以及类别权重文件。权重一定要重算，不同数据集的分布差异很大，沿用旧权重比不传权重还糟。换完先跑 5 轮看 loss 是否正常下降，如果第一轮 loss 就超过log(num_classes)很多，说明数据管线有问题，别急着调超参。

换主干的话，UNet 的inc和down1到down4是天然的替换点。把编码器换成torchvision里的 ResNet 或 EfficientNet，取中间四层的输出作为 skip，解码器不用动。有一种常见做法是保留torchvision的预训练权重，训练时给编码器一个更小的学习率，解码器用正常学习率，这种分层学习率在小数据集上通常能快 20 轮左右收敛。

# 分层学习率：编码器小、解码器大 encoder_params = list(model.inc.parameters()) + \ list(model.down1.parameters()) + \ list(model.down2.parameters()) + \ list(model.down3.parameters()) + \ list(model.down4.parameters()) decoder_params = [p for n, p in model.named_parameters() if not any(n.startswith(k) for k in ("inc", "down1", "down2", "down3", "down4"))] optimizer = torch.optim.AdamW([ {"params": encoder_params, "lr": cfg["train"]["lr"] * 0.1}, {"params": decoder_params, "lr": cfg["train"]["lr"]}, ], weight_decay=cfg["train"]["weight_decay"])

推理增强（TTA）是性价比最高的免费涨点手段，翻转加多尺度，通常能换来 1 到 2 个点的 mIoU，代价只是推理时间变几倍。做法是把原图、水平翻转、垂直翻转分别过一遍模型，概率图翻回来之后取平均再argmax。

导出 ONNX 时有个细节必须注意：模型输出是 logits，导出之后如果推理端又套了一层 softmax，结果会和 PyTorch 侧对不上。导出前用固定尺寸的假输入跑一遍torch.onnx.export，再用onnxruntime加载，把同一张图分别送进 PyTorch 和 ORT，逐元素比对差值，atol=1e-4以内才算通过。这个验证步骤看起来多余，但少做一次，后面排查线上问题可能要花一整天。

本文还有配套的精品资源，点击获取