简介：面向具备一定 PyTorch 基础、希望完成多类别语义分割实战的开发者与学习者，这份资料围绕 U-Net 在自建数据集上的训练与评估展开，适用于医学影像、遥感图像等像素级分类场景。压缩包共 46 个文件，以 19 个 py 源码为主体，另有 24 个 pyc 缓存、2 个 txt 与 1 个 json 辅助记录，整体约 69KB。代码按数据加载与增强、模型定义、损失函数、评价指标、学习率调度、训练日志与结果保存等模块组织，便于对照理解多类别输出通道设置、交叉熵训练和 IoU 等指标评估流程。资源还保留了训练脚本与演示入口，可帮助读者梳理从数据读取、前向计算到指标统计的完整链路，并据此改造成自己的数据集实验。已有 15284 人学习或下载，适合用作课程实验、项目复现与二次开发的参考底稿。

1. U-Net 多类别分割：一份能直接替换自己数据集的 PyTorch 工程

拿到一个语义分割项目，最让人头疼的往往不是模型写不出来，而是数据进来就崩：标签是 RGB 彩色图、类别编号不连续、背景类没算进去、训练 loss 一直卡在 0.0x。这份资源做的事很朴素——用 PyTorch 从零把 U-Net 搭起来，跑通一套多类别（不只是前景/背景二分类）的自定义数据集训练与推理流程，并且把"换数据"这件事的成本压到最低。它适合两类人：一是刚接触 pytorch Unet 语义分割、需要一份能跑通全链路参考实现的新手；二是手里已有自己标注数据、想找个干净 baseline 直接改的熟手。整套代码不含花哨的 trick，胜在结构清楚、每一环都能单独拆开改。

2. 工程骨架与数据流：从文件夹到张量的完整链路

先把项目拆开看结构。语义分割项目最容易被写成一坨，训练脚本里塞数据预处理、塞模型定义、塞指标计算，改一行牵全身。这份工程走的是常见的分层组织路线，目录职责相对清晰：

unet-multiclass-seg/ ├── datasets/ │ ├── images/ # 原图，与标签同名 │ │ ├── 0001.png │ │ └── 0002.png │ └── masks/ # 标签图，单通道或 RGB 索引色 │ ├── 0001.png │ └── 0002.png ├── models/ │ └── unet.py # U-Net 网络定义 ├── datasets.py # Dataset / DataLoader 构建 ├── train.py # 训练主循环 ├── predict.py # 单图 / 批量推理 ├── metrics.py # 混淆矩阵、mIoU、Dice └── configs.py # 类别数、尺寸、超参集中管理

2.1 目录约定与标签映射

多类别分割与二分类最大的差别就在标签。二分类里前景是 1、背景是 0，随手mask > 0就完事；多类别不行，标签图往往是 RGB 三通道的伪彩色索引图，比如 (0,0,0) 是背景、(128,0,0) 是类别一、(0,128,0) 是类别二。直接丢进网络会得到三通道输出，维度对不上，CrossEntropy 会直接报错。

常见做法是在datasets.py里先定义一张颜色到类号的映射表，读图时做一次转换：

import numpy as np from PIL import Image # 颜色 -> 类别编号，顺序必须与训练时一致 PALETTE = [ (0, 0, 0), # 0 背景 (128, 0, 0), # 1 类别A (0, 128, 0), # 2 类别B (0, 0, 128), # 3 类别C ] NUM_CLASSES = len(PALETTE) def mask_to_index(mask_path): rgb = np.array(Image.open(mask_path).convert("RGB")) index = np.zeros(rgb.shape[:2], dtype=np.uint8) for cls_id, color in enumerate(PALETTE): # 逐类比对，命中即赋类号；也可用 (rgb == color).all(axis=-1) match = (rgb[:, :, 0] == color[0]) & \ (rgb[:, :, 1] == color[1]) & \ (rgb[:, :, 2] == color[2]) index[match] = cls_id return index

这段逻辑说明几个关键点：NUM_CLASSES由调色板长度决定，而不是写死；index用uint8足够撑到 255 类，省内存；逐通道比对比np.all(rgb == color, axis=-1)慢一点点，但避免了三维广播临时数组，对大图更友好。参数上必须盯住的是PALETTE的顺序——训练时类 1 是红，推理时绘图却把红画成类 2，结果看起来全对，实际全是错的，这种玄学 bug 排查起来非常费神。

提示：如果你的标注工具导出的是单通道灰度索引图（像素值就是类号），可以跳过上面整段，直接Image.open(p).convert("L")，但务必确认标签里没有 255 这种"忽略区域"值，否则要单独映射成 255 并在 loss 里设ignore_index=255。

2.2 Dataset 与增强：同步变换是硬约束

图像分割的数据增强和分类完全不同：分类只要翻转图片，分割必须让图和标签同步做同一个几何变换。这一步做错，模型训练 loss 会缓慢下降到某个值之后再也不动，属于典型的血泪经验。

import torch from torch.utils.data import Dataset from PIL import Image import numpy as np class SegDataset(Dataset): def __init__(self, img_paths, mask_paths, size=(256, 256), train=True): self.img_paths = img_paths self.mask_paths = mask_paths self.size = size self.train = train def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img = Image.open(self.img_paths[idx]).convert("RGB") mask = Image.fromarray(mask_to_index(self.mask_paths[idx])) # 几何变换必须双输入同步，用同一次随机参数 if self.train and np.random.rand() > 0.5: img = img.transpose(Image.FLIP_LEFT_RIGHT) mask = mask.transpose(Image.FLIP_LEFT_RIGHT) img = img.resize(self.size, Image.BILINEAR) # 标签必须最近邻，双线性会插值出并不存在的类号 mask = mask.resize(self.size, Image.NEAREST) img = np.array(img, dtype=np.float32) / 255.0 img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225] img = torch.from_numpy(img).permute(2, 0, 1).float() mask = torch.from_numpy(np.array(mask)).long() return img, mask

逻辑上分四步：读图、同步几何增强、缩放、归一化转张量。参数要点有三个。第一，size必须能被 16 整除，U-Net 下采样四次，尺寸不是 16 的倍数时，上采样拼接会因为尺寸差 1 而报 shape 不匹配。第二，mask.resize只能用NEAREST，用双线性会在类别边界插出 1.5 这种非法类号，CrossEntropyLoss不会报错，但会当成合法类别去学，指标就永远上不去。第三，归一化用的是 ImageNet 均值方差，如果你的数据是医学影像或工业缺陷图，可以换成自己算的统计量，但训练和推理必须用同一套。

DataLoader 侧没什么特殊，num_workers在 Linux 下给 4~8，Windows 下建议给 0 避免多进程反复重启：

from torch.utils.data import DataLoader train_loader = DataLoader( SegDataset(train_imgs, train_masks, size=(256, 256), train=True), batch_size=8, shuffle=True, num_workers=4, pin_memory=True )

2.3 网络结构：U-Net 的编码解码与跳连

U-Net 的结构本身不复杂，双卷积块 + 四次下采样 + 四次上采样 + 同级跳连。判断一份实现是否可用，主要看跳连拼接的方式和通道数是否对齐：

import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x) class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=4, base=64): super().__init__() self.d1 = DoubleConv(in_ch, base) self.d2 = DoubleConv(base, base * 2) self.d3 = DoubleConv(base * 2, base * 4) self.d4 = DoubleConv(base * 4, base * 8) self.pool = nn.MaxPool2d(2) self.bottleneck = DoubleConv(base * 8, base * 16) self.u4 = nn.ConvTranspose2d(base * 16, base * 8, 2, stride=2) self.c4 = DoubleConv(base * 16, base * 8) # 拼接后通道翻倍 self.u3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2) self.c3 = DoubleConv(base * 8, base * 4) self.u2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2) self.c2 = DoubleConv(base * 4, base * 2) self.u1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2) self.c1 = DoubleConv(base * 2, base) self.out = nn.Conv2d(base, num_classes, 1) # 输出类别数通道 def forward(self, x): x1 = self.d1(x) x2 = self.d2(self.pool(x1)) x3 = self.d3(self.pool(x2)) x4 = self.d4(self.pool(x3)) b = self.bottleneck(self.pool(x4)) b = self.c4(torch.cat([self.u4(b), x4], dim=1)) b = self.c3(torch.cat([self.u3(b), x3], dim=1)) b = self.c2(torch.cat([self.u2(b), x2], dim=1)) b = self.c1(torch.cat([self.u1(b), x1], dim=1)) return self.out(b)

参数上，base是通道基数，默认 64 对应标准 U-Net 的 64/128/256/512/1024；显存吃紧就降到 32，四路通道会等比缩减。num_classes必须包含背景，四类目标就写 5，这是最高频的翻车点。bias=False是因为紧跟 BN，偏置会被 BN 的 beta 覆盖，省参数。跳连用torch.cat而不是相加，这是 U-Net 与 FCN 的核心区别，拼接后通道翻倍，所以后面的DoubleConv输入通道是当前层两倍，写错这里会在 forward 时报错但不会给明确指出哪一层。

输出层用 1x1 卷积给出num_classes个通道，不接 softmax。因为训练用CrossEntropyLoss，它内部自带 log_softmax，模型再 softmax 就重复了，梯度会被压平，训练会异常缓慢。推理时要自己补 softmax 或直接 argmax。

3. 训练与评估：把 loss 和 mIoU 落到每一轮

3.1 损失函数与优化器组合

多类别分割的标准组合是CrossEntropyLoss+Adam，如果要处理类别不平衡（比如背景占了 90% 像素），再叠加 Dice 或 Focal：

import torch.nn as nn import torch device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = UNet(in_ch=3, num_classes=NUM_CLASSES, base=64).to(device) # ignore_index 用于跳过标注中的边界/无效像素 ce = nn.CrossEntropyLoss(ignore_index=255) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode="max", factor=0.5, patience=5 )

ignore_index=255要和前面标签预处理对得上，标签里没有 255 就不需要这个参数，写了反而潜藏问题。学习率 1e-3 是 Adam 的通用起点，小数据集建议 5e-4。ReduceLROnPlateau用mode="max"是因为监控的 mIoU 越高越好，如果监控 val loss 就改成"min"，模式填反了会导致学习率一直往下降，模型早早停止学习。

3.2 混淆矩阵驱动的 mIoU

分割的指标不能只看像素准确率。背景占大头时，全预测背景也能拿到 95% 准确率，毫无意义。工程上一般统一用混淆矩阵算 mIoU：

import numpy as np class ConfusionMatrix: def __init__(self, num_classes): self.n = num_classes self.mat = np.zeros((self.n, self.n), dtype=np.int64) def update(self, pred, target): # pred/target: (B, H, W) 的整型张量 mask = (target >= 0) & (target < self.n) hist = np.bincount( self.n * target[mask].astype(int) + pred[mask].astype(int), minlength=self.n ** 2 ).reshape(self.n, self.n) self.mat += hist def miou(self): ious = [] for i in range(self.n): tp = self.mat[i, i] fp = self.mat[:, i].sum() - tp fn = self.mat[i, :].sum() - tp denom = tp + fp + fn if denom > 0: ious.append(tp / denom) # 该类别未出现则跳过 return float(np.mean(ious))

逻辑上先累加一个 N×N 的混淆矩阵，行是真实类别、列是预测类别，第 i 行第 i 列就是真正例。miou里对分母为 0 的类别跳过，避免小数据集某些类没出现时除零算出 nan。参数上注意minlength必须给到n**2，否则 bincount 返回的数组长度按最大值来，reshape 会直接崩。

训练循环里每轮在验证集上跑一遍，用torch.no_grad()包住，pred.argmax(dim=1)取类别：

model.eval() cm = ConfusionMatrix(NUM_CLASSES) with torch.no_grad(): for imgs, masks in val_loader: imgs, masks = imgs.to(device), masks.to(device) logits = model(imgs) pred = logits.argmax(dim=1).cpu().numpy() cm.update(pred, masks.cpu().numpy()) print(f"mIoU: {cm.miou():.4f}")

3.3 从零跑通的执行顺序

实际动手按这个顺序走，比对着脚本瞎点要快。先装依赖，torch、torchvision、pillow、numpy四件套就够，版本按本机 CUDA 选对应轮子。然后把数据按 2.1 的目录结构摆好，确保images和masks文件名一一对应。接着改configs.py或直接改UNet的num_classes，确认PALETTE长度与类别数一致。最后跑：

python train.py --data-root ./datasets --epochs 50 --batch-size 8 --lr 1e-3

跑起来先别急着看 50 轮，观察前 5 轮：loss 应该从 2.x（类别数的对数级别）稳定下降，如果第 1 轮就掉到 0.0x，多半是标签全 0 或全预测某一类，去查标签读取那段。mIoU 在 5 轮后应该有明显上升，如果一直是 0，检查PALETTE是否命中、ignore_index是否把全部像素忽略了。

推理阶段换成predict.py，注意模型必须load_state_dict后再eval()，否则 BN 仍在用 batch 统计量，单张图推理结果会飘：

model.load_state_dict(torch.load("best.pth", map_location=device)) model.eval() logits = model(img_tensor.to(device)) pred = logits.argmax(dim=1).squeeze(0).cpu().numpy() # 用逆映射把类号染回可视化颜色

4. 常见问题排查：那些让 loss 不动、shape 对不上的坑

4.1 现象：训练 loss 一直在 0.0x 附近，mIoU 几乎为 0

原因通常两处。一是标签是 RGB 图但没走mask_to_index，直接用.long()转成了三个通道张量，CrossEntropyLoss收到形状不对的 target，要么报错要么把通道数当类别，语义全乱。二是num_classes含不含背景没统一，模型输出 4 类而标签里出现了 4，越界索引被静默截断。解决方式很直接：训练前单独写三行打印标签的min()、max()、unique()，类号最大值得等于num_classes - 1，对不上就先修数据。

4.2 现象：上采样拼接时报 size 不匹配

U-Net 四次MaxPool2d(2)会把尺寸除以 16，ConvTranspose2d再乘 2。输入尺寸不是 16 的倍数（比如 250×250）时，下采样向下取整成 15×15，上采样变成 30×30，而对应的编码层是 32×32，torch.cat立刻报错。解决是训练和验证统一 resize 到 16 的倍数，或者改用先F.interpolate对齐尺寸再拼接的写法，代价是多一次插值。我一般直接在 Dataset 里锁死 256 或 512，省得后面每处都要处理。

4.3 现象：训练正常，推理结果整体偏移或全是某一类

八成是model.eval()漏了，BatchNorm 在推理时用了单张图的统计量，而这个统计量本身噪声极大，输出分布直接飘掉。就像黑匣子没关好。另一个可能是load_state_dict加载的是训练过程中保存的 epoch 权重，但num_classes改过，strict=False下部分层被静默跳过。解决：推理脚本里eval()+torch.no_grad()一个不能少，加载权重时保持strict=True，让它报错而不是装没事。

4.4 现象：显存 OOM，batch 已经降到 1 还是炸

大图是主因。256×256 的图在 base=64 时 8G 显存大约能跑 batch 4~8；512×512 会翻四倍。解决优先级：先降base到 32，再考虑混合精度torch.cuda.amp，最后才动 batch。还有个小坑是pin_memory=True配大分辨率时会额外占内存，num_workers过高时每个 worker 都会缓存一份数据，显存没炸内存先炸，排查时容易找错方向。

4.5 现象：验证集 mIoU 忽高忽低，同权重两次评估差很多

数据增强里把验证集也做了随机翻转，或者SegDataset的train=True没在验证集上关掉。分割的随机增强对标签极其敏感，验证必须确定性输入，否则同一模型每次评估都是不同数据分布。解决：验证集train=False，并在 DataLoader 里shuffle=False。

5. 二次开发：换数据集与加 trick 的稳妥改法

想把这份 baseline 变成能打的模型，改动的顺序应该由外到内，先数据后模型。

数据侧最先动的是类别定义。改PALETTE和num_classes是必须的，但别只改这两处就收工。建议在configs.py里把它们集中成一份配置，训练、推理、可视化全部从配置读，避免出现训练 5 类推理 4 类的错位。数据集变大后，把 Dataset 的路径来源改成os.walk或读取一个 csv 索引，比手工两两对应文件靠谱得多。增强可以补随机裁剪、颜色抖动，但每次加完都拿 20 张图跑一遍可视化，确认图和标签仍然贴合，这是最省时间的后悔药。

模型侧可以从三个方向加，按收益从高到低排：

换 backbone 的改动其实不大，把d1~d4替换成 torchvision 预训练主干，记住replace_stride_with_dilation或直接拿四个 stage 的输出，跳连通道数随之变化，DoubleConv的输入通道要跟着重算。这一步最容易翻车的地方在通道数对不上，我一般先在 forward 里加一行 print 打各层 shape，确认无误再删。

验证方法上，别只盯 mIoU。每次改完跑一遍混淆矩阵，看是哪一类在拉低均值，比盯着一个总分有方向得多。如果某个类别 IoU 长期接近 0，八成是样本太少或标签画得有问题，先查数据再调模型。另外建议固定一个 8~16 张图的小验证集，改一次看一次可视化叠加图，肉眼比指标更早发现"整体偏移半格"这类系统性错误。

最后一条习惯：动任何结构前先跑通一遍原始版本并记下基准 mIoU，改完再对比。我吃过太多次"以为改好了其实基线上就没跑对"的亏，从那以后每次重构都强制先跑基线、存权重、打指标，再动手。希望这套流程帮你少绕几个弯。

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