简介：这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者，聚焦在PyTorch框架下用Unet完成多类别语义分割任务，可应用于医学影像、遥感图像等场景。压缩包共46个文件，以19个py脚本和24个pyc缓存文件为主，另含2个txt与1个json配置，整体约69KB，涵盖数据加载、自定义变换、模型定义、损失函数、学习率调度、指标计算与训练保存等模块，目录结构清晰便于按功能查阅。目前已有15274人学习下载，热度较高。读者可据此搭建完整的训练与评估流程，理解编码器-解码器结构、跳跃连接与多通道输出设计，并借助交叉熵损失、IoU等指标完成分割效果验证与可视化对比，快速迁移到自己的多类别数据集上。
1. 从一张标注图到可训练模型：这套 PyTorch UNet 多类别分割方案到底能干什么
如果你手里有一批自己标注的多类别语义分割数据，比如遥感影像里的水体、建筑、植被，或者工业质检里的划痕、凹坑、脏污，第一反应往往是找一套能直接跑通的 UNet 代码。但真正动手时才发现，网上大部分示例要么只支持二分类，要么把数据加载写死成 VOC 格式，要么训练完预测出来全是同一类。这套 PyTorch 下实现 UNet 对自己多类别数据集做语义分割的资源，解决的就是这个断层：它不绑定任何公开数据集，而是把重点放在“如何把自己的多类别标注喂进去、如何让输出通道数对上类别数、如何验证每一类真的被分出来”。适合已经会写基础 PyTorch 训练循环、但卡在数据组织和多类别输出配置上的从业者。下面按数据准备、模型改造、训练调参、踩坑排查、进阶验证的顺序拆开讲，每一步都落到能直接抄的代码和参数上。
2. 多类别数据集的目录结构与 Dataset 写法：别让标签从第一步就错位
2.1 图像与掩码的对应关系怎么组织
多类别分割翻车最多的地方不是模型，而是数据索引对不上。常见做法是建两个平行文件夹，图像和掩码同名不同后缀，或者用同一个文件名分别放在images/和masks/下。我一般会强制要求文件名完全一致，只靠目录区分，这样 Dataset 里用sorted(os.listdir)就能保证顺序一致，避免训练时图像和标签错位这种玄学问题。
目录结构建议如下：
dataset/ images/ 0001.png 0002.png masks/ 0001.png 0002.png掩码必须是单通道灰度图，像素值直接就是类别索引，比如 0 是背景，1 是类别 A，2 是类别 B。不要用 RGB 彩色掩码，也不要用调色板 PNG 存了颜色却忘了转索引，否则后面算损失时会出现类别值超出范围。
2.2 自定义 Dataset 的完整实现
下面这段代码可以直接替换掉你项目里的 Dataset 类，核心是保证图像和掩码同步做增强，并且把掩码转成long类型。
import os import torch from torch.utils.data import Dataset from PIL import Image import numpy as np import torchvision.transforms.functional as TF import random class MultiClassSegDataset(Dataset): def __init__(self, root_dir, img_size=256, augment=True): self.img_dir = os.path.join(root_dir, "images") self.mask_dir = os.path.join(root_dir, "masks") # 只保留图像和掩码都存在的样本，避免训练中途报错 self.names = sorted([ n for n in os.listdir(self.img_dir) if os.path.exists(os.path.join(self.mask_dir, n)) ]) self.img_size = img_size self.augment = augment 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 = TF.resize(img, [self.img_size, self.img_size]) mask = TF.resize(mask, [self.img_size, self.img_size], interpolation=TF.InterpolationMode.NEAREST) if self.augment and random.random() > 0.5: img = TF.hflip(img) mask = TF.hflip(mask) img = TF.to_tensor(img) # 归一化到 [0,1] img = TF.normalize(img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) mask = torch.from_numpy(np.array(mask)).long() # 关键：long 类型 return img, mask逻辑说明：convert("L")把掩码强制转成单通道，resize时掩码用最近邻插值，这是多类别分割里最容易忽略的一步，用双线性会生成 1.5 这种不存在的类别值，训练时直接报Target out of bounds。mask转long是因为CrossEntropyLoss要求目标为int64。参数上img_size建议从 256 起步，显存够再上 512；augment只做了水平翻转，多类别场景下旋转和裁剪要谨慎，容易把稀有类别裁没。
2.3 类别不平衡时先统计像素分布
多类别数据集往往背景占 90% 以上，直接训练会让模型学会全预测背景。在写 Dataset 之后、开训之前，先跑一遍统计：
from collections import Counter dataset = MultiClassSegDataset("dataset", augment=False) counter = Counter() for _, mask in dataset: counter.update(mask.numpy().flatten().tolist()) total = sum(counter.values()) for cls, cnt in sorted(counter.items()): print(f"class {cls}: {cnt/total:.4f}")如果发现某一类占比低于 1%，后面损失函数就要考虑带权重的CrossEntropyLoss，权重可以按类别频率的倒数来设，但不要极端到把背景权重压成 0，否则背景会大量误检。
3. UNet 输出层改造与多类别损失配置：通道数、激活和权重怎么对齐
3.1 输出通道数必须等于类别数
原始 UNet 最后一般是Conv2d(64, 1, 1)做二分类。多类别要把输出通道改成num_classes，并且不要加 Sigmoid，因为多类别用的是 Softmax 交叉熵。下面是一个可直接嵌入的 UNet 尾部改造：
import torch.nn as nn class UNetOut(nn.Module): def __init__(self, in_channels=64, num_classes=4): super().__init__() self.conv = nn.Conv2d(in_channels, num_classes, kernel_size=1) def forward(self, x): return self.conv(x) # 输出 [B, num_classes, H, W]，不接激活逻辑说明：kernel_size=1只做通道映射，不改变空间尺寸。num_classes要和你掩码里的最大类别索引加一严格一致，比如掩码里有 0、1、2、3 四类，就设 4。常见错误是设成 3 却喂了标签 3，训练时直接越界。参数上in_channels取决于你 UNet 最后一层上采样后的通道数，常见实现是 64，如果你改了基础通道数，这里要同步改。
3.2 带类别权重的损失函数
多类别分割默认用nn.CrossEntropyLoss，它内部已经包含 Softmax，所以模型输出不要再加 Softmax。类别不平衡时传入weight：
import torch # 假设统计出四类频率，背景高，小类别低 class_counts = torch.tensor([0.70, 0.10, 0.15, 0.05]) weights = 1.0 / class_counts weights = weights / weights.sum() # 归一化，避免整体损失尺度爆炸 criterion = nn.CrossEntropyLoss(weight=weights)逻辑说明：weight越大，该类在损失里占比越高。归一化是为了让总损失量级和原来接近，否则学习率要重新调。如果小类别仍然学不出来，可以进一步用 Dice Loss 或 Focal Loss 组合，但建议先用带权交叉熵跑通基线，再换复杂损失，不然出问题分不清是数据还是损失。
3.3 训练循环里必须打印每类 IoU
只看总 loss 下降没有意义，多类别分割要看每一类的 IoU。下面是一个简化的验证片段：
def compute_iou(pred, target, num_classes): ious = [] pred = pred.argmax(dim=1) # [B,H,W] for cls in range(num_classes): pred_cls = (pred == cls) target_cls = (target == cls) inter = (pred_cls & target_cls).sum().item() union = (pred_cls | target_cls).sum().item() ious.append(inter / union if union > 0 else float('nan')) return ious逻辑说明：argmax(dim=1)把每个像素变成类别索引，再逐类算交并比。如果某一类union为 0，说明验证集里没有该类，返回nan而不是 0，避免误判模型完全没学到。参数上num_classes和模型输出通道一致。训练时每几个 epoch 跑一次验证，把每类 IoU 打印出来，比只看 loss 可靠得多。
4. 训练参数、学习率与显存控制：把 256 分辨率跑稳的实操配置
4.1 优化器与学习率选择
UNet 分割任务常用 Adam 或 SGD。Adam 起步快，适合先跑通；SGD 配合动量最终精度可能更好，但调参麻烦。我一般先用 Adam，学习率 1e-3，如果 loss 震荡就降到 1e-4。下面是一个可直接用的训练配置：
import torch.optim as optim model = UNet(num_classes=4).cuda() optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5) scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=5 )逻辑说明：weight_decay加一点正则，防止过拟合。ReduceLROnPlateau在验证 loss 不降时把学习率减半，patience=5表示连续 5 个 epoch 没改善才动。参数上factor不要低于 0.1，否则学习率掉太快，模型还没收敛就停了。
4.2 batch size 与显存的关系
256×256 输入、UNet 基础通道 64 的情况下，batch size 8 大约占 6GB 显存，batch size 16 大约 10GB。如果显存不够，优先降 batch size 而不是降分辨率，因为分辨率对分割细节影响更大。常见做法是用梯度累积模拟大 batch：
accum_steps = 4 for i, (img, mask) in enumerate(loader): img, mask = img.cuda(), mask.cuda() out = model(img) loss = criterion(out, mask) / accum_steps loss.backward() if (i + 1) % accum_steps == 0: optimizer.step() optimizer.zero_grad()逻辑说明：损失除以accum_steps是为了让梯度尺度和真实大 batch 一致。注意optimizer.zero_grad()要放在累积之后，不能每个 batch 都清。参数上accum_steps根据显存和目标 batch 调整，一般 2 到 8 之间。
4.3 数据增强的边界
多类别分割的增强比分类更敏感。水平翻转、垂直翻转、小角度旋转通常安全；随机裁剪要保证裁剪后每类还有像素；颜色抖动只对图像做，不能对掩码做。常见做法是把增强写在 Dataset 里同步处理，而不是用现成的分类增强库分别处理图像和掩码，否则很容易错位。
5. 多类别 UNet 训练避坑：从标签越界到全背景预测的排查清单
5.1 报错 Target out of bounds
现象：训练一开始就抛Target out of bounds或Assertion failed。原因：掩码像素值超过了num_classes - 1，常见于用彩色掩码转灰度后值变成 128、255，或者 resize 用了双线性插值生成了中间值。解决：先跑一遍np.unique(mask)看所有像素值，确认最大值小于num_classes；resize 掩码强制用最近邻；如果原始掩码是调色板 PNG，用convert("L")之前先确认调色板索引就是类别索引。
5.2 预测结果全是背景
现象：训练 loss 降到很低，但预测图几乎全黑，每类 IoU 里背景接近 1，其他类接近 0。原因：类别极度不平衡，背景占比过高，模型学到全预测背景就能拿到很低 loss。解决：给CrossEntropyLoss加类别权重，或者改用 Dice Loss；同时在采样时对含小类别的图像做过采样。验证时不要只看准确率，要看每类 IoU。
5.3 图像和掩码错位
现象：训练 loss 能降，但预测出来的形状和原图对不上，像是平移或翻转了。原因：图像和掩码用了不同的增强随机种子，或者文件名排序不一致导致配对错误。解决：增强必须在同一个__getitem__里同步做，用同一个随机数控制；文件名用sorted保证顺序；可以在 Dataset 里加一句断言，打印当前样本名，确认图像和掩码来自同一文件。
5.4 验证集 loss 不降但训练 loss 一直降
现象：训练 loss 持续下降，验证 loss 先降后升。原因：过拟合，数据量太少或者模型太大。解决：加数据增强、加weight_decay、减小模型通道数，或者早停。多类别分割里如果某一类样本特别少，过拟合会更快，建议先检查每类在训练集和验证集的分布是否一致。
5.5 显存溢出但 batch size 已经很小
现象：batch size 降到 2 还是 OOM。原因：UNet 在 512 分辨率下中间特征图很大，或者没有用torch.no_grad()做验证。解决：验证阶段包with torch.no_grad():，减少缓存；用混合精度训练torch.cuda.amp；如果还不够，把输入分辨率降到 256 或 192。不要盲目加torch.cuda.empty_cache()，它解决不了根本的显存占用。
6. 进阶验证与推理部署：把单张预测做成可复现的检查流程
训练完模型只是第一步，真正要确认多类别分割可用，得把推理结果可视化并逐类检查。我一般会写一个独立的推理脚本，加载权重后对验证集随机抽几张，把原图、真值掩码、预测掩码并排保存，用固定颜色映射，这样一眼就能看出哪一类被漏了或混了。
import torch import numpy as np from PIL import Image def predict_and_save(model, dataset, index, save_path, num_classes=4): model.eval() img, mask = dataset[index] with torch.no_grad(): out = model(img.unsqueeze(0).cuda()) pred = out.argmax(dim=1).squeeze(0).cpu().numpy() # 固定颜色表，保证每次可视化一致 palette = np.array([ [0, 0, 0], # 背景黑 [255, 0, 0], # 类1红 [0, 255, 0], # 类2绿 [0, 0, 255], # 类3蓝 ], dtype=np.uint8) pred_color = palette[pred] mask_color = palette[mask.numpy()] img_np = (img.permute(1, 2, 0).numpy() * 255).astype(np.uint8) combined = np.concatenate([img_np, mask_color, pred_color], axis=1) Image.fromarray(combined).save(save_path)逻辑说明：argmax(dim=1)得到每像素类别，palette把类别索引映射成固定颜色，方便对比。combined横向拼接原图、真值、预测，保存成一张图。参数上num_classes和颜色表行数要一致，颜色表按类别顺序定义，不要用随机颜色，否则每次跑出来颜色不同，没法对比。
推理阶段还有几个容易忽略的点。第一，输入归一化必须和训练时完全一致，训练用了 ImageNet 均值和方差，推理也要用同一套，否则预测会整体偏移。第二，如果训练时做了 resize，推理也要 resize 到同一尺寸，预测完再放大回原图时用最近邻，不要用双线性。第三，多类别模型输出不要加 Softmax 再 argmax，直接 argmax 等价且更省算力。
验证时我习惯按类别分别统计 IoU 和像素准确率，并且把每类 IoU 低于 0.3 的样本单独挑出来看。常见情况是某一类在验证集里只出现几张小图，IoU 波动极大，这时候不要急着改模型，先确认验证集里该类样本量是否足够。如果某一类始终学不出来，优先检查标注是否一致，比如同一类在不同图像里标了不同索引，这种问题模型再调也救不回来。
从那以后我每次跑多类别分割，都强制先跑一遍np.unique(mask)和类别分布统计，再开训。这个习惯帮我省掉了至少一半的无效训练。希望帮到你。
本文还有配套的精品资源，点击获取
