简介这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者聚焦在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 是类别 A2 是类别 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_size256, augmentTrue): 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], interpolationTF.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 起步显存够再上 512augment只做了水平翻转多类别场景下旋转和裁剪要谨慎容易把稀有类别裁没。2.3 类别不平衡时先统计像素分布多类别数据集往往背景占 90% 以上直接训练会让模型学会全预测背景。在写 Dataset 之后、开训之前先跑一遍统计from collections import Counter dataset MultiClassSegDataset(dataset, augmentFalse) counter Counter() for _, mask in dataset: counter.update(mask.numpy().flatten().tolist()) total sum(counter.values()) for cls, cnt in sorted(counter.items()): print(fclass {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_channels64, num_classes4): super().__init__() self.conv nn.Conv2d(in_channels, num_classes, kernel_size1) def forward(self, x): return self.conv(x) # 输出 [B, num_classes, H, W]不接激活逻辑说明kernel_size1只做通道映射不改变空间尺寸。num_classes要和你掩码里的最大类别索引加一严格一致比如掩码里有 0、1、2、3 四类就设 4。常见错误是设成 3 却喂了标签 3训练时直接越界。参数上in_channels取决于你 UNet 最后一层上采样后的通道数常见实现是 64如果你改了基础通道数这里要同步改。3.2 带类别权重的损失函数多类别分割默认用nn.CrossEntropyLoss它内部已经包含 Softmax所以模型输出不要再加 Softmax。类别不平衡时传入weightimport 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(weightweights)逻辑说明weight越大该类在损失里占比越高。归一化是为了让总损失量级和原来接近否则学习率要重新调。如果小类别仍然学不出来可以进一步用 Dice Loss 或 Focal Loss 组合但建议先用带权交叉熵跑通基线再换复杂损失不然出问题分不清是数据还是损失。3.3 训练循环里必须打印每类 IoU只看总 loss 下降没有意义多类别分割要看每一类的 IoU。下面是一个简化的验证片段def compute_iou(pred, target, num_classes): ious [] pred pred.argmax(dim1) # [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(dim1)把每个像素变成类别索引再逐类算交并比。如果某一类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_classes4).cuda() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-5) scheduler optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience5 )逻辑说明weight_decay加一点正则防止过拟合。ReduceLROnPlateau在验证 loss 不降时把学习率减半patience5表示连续 5 个 epoch 没改善才动。参数上factor不要低于 0.1否则学习率掉太快模型还没收敛就停了。4.2 batch size 与显存的关系256×256 输入、UNet 基础通道 64 的情况下batch size 8 大约占 6GB 显存batch size 16 大约 10GB。如果显存不够优先降 batch size 而不是降分辨率因为分辨率对分割细节影响更大。常见做法是用梯度累积模拟大 batchaccum_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_classesresize 掩码强制用最近邻如果原始掩码是调色板 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_classes4): model.eval() img, mask dataset[index] with torch.no_grad(): out model(img.unsqueeze(0).cuda()) pred out.argmax(dim1).squeeze(0).cpu().numpy() # 固定颜色表保证每次可视化一致 palette np.array([ [0, 0, 0], # 背景黑 [255, 0, 0], # 类1红 [0, 255, 0], # 类2绿 [0, 0, 255], # 类3蓝 ], dtypenp.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], axis1) Image.fromarray(combined).save(save_path)逻辑说明argmax(dim1)得到每像素类别palette把类别索引映射成固定颜色方便对比。combined横向拼接原图、真值、预测保存成一张图。参数上num_classes和颜色表行数要一致颜色表按类别顺序定义不要用随机颜色否则每次跑出来颜色不同没法对比。推理阶段还有几个容易忽略的点。第一输入归一化必须和训练时完全一致训练用了 ImageNet 均值和方差推理也要用同一套否则预测会整体偏移。第二如果训练时做了 resize推理也要 resize 到同一尺寸预测完再放大回原图时用最近邻不要用双线性。第三多类别模型输出不要加 Softmax 再 argmax直接 argmax 等价且更省算力。验证时我习惯按类别分别统计 IoU 和像素准确率并且把每类 IoU 低于 0.3 的样本单独挑出来看。常见情况是某一类在验证集里只出现几张小图IoU 波动极大这时候不要急着改模型先确认验证集里该类样本量是否足够。如果某一类始终学不出来优先检查标注是否一致比如同一类在不同图像里标了不同索引这种问题模型再调也救不回来。从那以后我每次跑多类别分割都强制先跑一遍np.unique(mask)和类别分布统计再开训。这个习惯帮我省掉了至少一半的无效训练。希望帮到你。本文还有配套的精品资源点击获取