企业官网建设流程全解析

简介：这份资源面向具备一定深度学习基础的开发者与图像分析方向的学习者，聚焦在PyTorch框架下用Unet完成多类别语义分割任务，可应用于医学影像、遥感图像等场景。压缩包共46个文件，约69KB，以19个py源码文件为核心，涵盖模型定义、数据加载、损失函数、学习率调度、指标计算与训练保存等模块，另有24个pyc缓存文件、2个txt说明与1个json配置，结构紧凑便于直接运行与二次修改。目前已有15264人学习下载，热度较高。资源围绕Unet的编码器—解码器结构、跳跃连接、多通道输出层设计，以及交叉熵损失、Adam优化器、数据增强、IoU与混淆矩阵评估等关键环节给出可复用代码，读者可据此快速搭建自己的多类别分割流程，并在此基础上尝试更换编码器、引入注意力机制或调整学习率策略以提升性能。

1. 从一张三通道标注图说起：Unet 多类别语义分割到底在做什么

你手里有一批自己拍的路面、遥感、医学或者工业质检图像，用 LabelMe 或 ArcGIS 标出了「裂缝、坑槽、标线」三类目标，导出成一张三通道 PNG。现在想用 Pytorch 跑一个 Unet，让模型对每个像素输出这三类的概率——这就是多类别语义分割要解决的事。它和二分类分割最大的区别在于：输出通道数从 1 变成 N，损失函数从 BCE 换成 CrossEntropy，标注图的像素值不再是 0/255 而是 0/1/2 的类别索引。很多人第一次跑 Unet 训练自己的数据集时，模型 loss 一直卡在 0.6 不降，八成是标注图没转成单通道索引图，或者忽略了一个叫 ignore_index 的参数。这篇笔记按「数据怎么准备 → 模型怎么改 → 训练怎么调 → 坑在哪」的顺序，把 Pytorch 下 Unet 多类别分割的完整链路拆开讲，新手能照着复现，熟手能看到边界条件。

2. 多类别数据集怎么变成 Unet 能吃的张量

2.1 标注格式的三种来源与统一策略

自己做的数据集，标注来源无非三类：LabelMe 的 JSON、ArcGIS 导出的三通道彩色 mask、以及用脚本批量生成的灰度索引图。Unet 的多类别训练只认最后一种——单通道、像素值等于类别索引、背景为 0。彩色 mask 看着直观，但直接喂进去会出大问题：模型会把 RGB 三个通道当成三个独立特征，而实际上它们只是同一个类别标签的编码。

常见做法是写一个转换脚本，把彩色 mask 按颜色映射表转成索引图。映射表要自己维护，比如(0,0,0)→0、(255,0,0)→1、(0,255,0)→2。转换时用 numpy 的布尔索引比逐像素循环快几十倍，一张 512×512 的图毫秒级完成。

import numpy as np from PIL import Image # 颜色到类别索引的映射，按自己的标注规范改 COLOR_MAP = { (0, 0, 0): 0, # 背景 (255, 0, 0): 1, # 类别1 (0, 255, 0): 2, # 类别2 (0, 0, 255): 3, # 类别3 } def color_mask_to_index(mask_path, save_path): img = np.array(Image.open(mask_path).convert('RGB')) index = np.zeros(img.shape[:2], dtype=np.uint8) for color, idx in COLOR_MAP.items(): # 三个通道同时匹配才算命中，避免颜色串扰 match = np.all(img == np.array(color), axis=-1) index[match] = idx Image.fromarray(index).save(save_path) return index

这段代码的关键在np.all(..., axis=-1)：必须三个通道同时相等才判定为该类别。如果只用单通道比较，红色 (255,0,0) 和黄色 (255,255,0) 会在第一个通道上撞车，导致类别错乱。转换完建议抽 10 张图用np.unique检查像素值分布，确认没有出现映射表之外的杂色。

2.2 Dataset 与 DataLoader：同步增强是翻车重灾区

图像和 mask 必须做完全相同的几何变换，这是多类别分割里最容易翻车的地方。用 torchvision 的transforms分别处理 image 和 mask，随机裁剪和翻转的参数不一致，训练几轮后 mask 和原图就错位了，loss 会诡异地停在某个值不动。

稳妥的写法是自己封装一个 Dataset，在__getitem__里用同一个随机种子控制几何变换，颜色变换只作用于原图。

import torch from torch.utils.data import Dataset import random import numpy as np from PIL import Image class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, img_list, size=512): self.img_dir = img_dir self.mask_dir = mask_dir self.img_list = img_list self.size = size def __len__(self): return len(self.img_list) def __getitem__(self, idx): name = self.img_list[idx] img = Image.open(f'{self.img_dir}/{name}.jpg').convert('RGB') mask = Image.open(f'{self.mask_dir}/{name}.png') # 同步缩放，mask 用最近邻避免插值出小数类别 img = img.resize((self.size, self.size), Image.BILINEAR) mask = mask.resize((self.size, self.size), Image.NEAREST) img = np.array(img, dtype=np.float32) / 255.0 mask = np.array(mask, dtype=np.int64) # 同步随机翻转，用同一个随机数决定 if random.random() > 0.5: img = np.fliplr(img).copy() mask = np.fliplr(mask).copy() # 只对原图做颜色抖动，mask 不参与 if random.random() > 0.5: img = img * random.uniform(0.8, 1.2) img = np.clip(img, 0, 1) img = torch.from_numpy(img).permute(2, 0, 1).float() mask = torch.from_numpy(mask).long() return img, mask

参数上有三个点要盯住：mask 缩放必须用NEAREST，用双线性会插出 1.5 这种非整数类别，CrossEntropy 直接报错；mask 的 dtype 必须是int64，uint8在部分 Pytorch 版本里会触发类型不匹配；颜色抖动后要clip回 [0,1]，否则归一化输入超出范围会让 BN 层统计量漂移。DataLoader 的num_workers在 Windows 上建议设 0 或 2，设大了容易在 WSL 里卡死，这是 pytorch 环境搭建 wsl 场景下的高频问题。

2.3 类别不平衡：先看分布再决定要不要加权

多类别数据集里背景往往占 80% 以上，直接训练模型会倾向于全预测背景，mIoU 看着有 0.8 其实全是背景刷出来的。动手前先用脚本统计每个类别的像素占比：

import numpy as np from PIL import Image import glob counts = np.zeros(4) # 假设4类 for p in glob.glob('masks/*.png'): m = np.array(Image.open(p)) for c in range(4): counts[c] += (m == c).sum() freq = counts / counts.sum() print('类别频率:', freq) # 权重取频率倒数，归一化后传给 CrossEntropyLoss weights = 1.0 / (freq + 1e-6) weights = weights / weights.sum() * len(weights) print('建议权重:', weights)

如果最小类频率低于 1%，加权是必要的；如果在 5% 以上，先不加权跑一轮看混淆矩阵，很多时候问题不在权重而在标注质量。权重别直接取倒数，要归一化到均值为 1 附近，否则 loss 量级变化太大会让学习率失效。

3. Unet 输出层与损失函数：多类别的两个必改点

3.1 输出通道数、激活函数与上采样方式

标准 Unet 的下采样是 4 次，通道从 64 翻到 1024，上采样回到原分辨率。做多类别时，最后 1×1 卷积的输出通道要改成num_classes，且不要接 Softmax。Pytorch 的CrossEntropyLoss内部已经包含 LogSoftmax，外面再接一层 Softmax 会导致梯度被压两次，训练极慢甚至不收敛。这是新手最常踩的坑，没有之一。

上采样方式上，原版 Unet 用转置卷积，容易出现棋盘格伪影。我一般把ConvTranspose2d换成nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)加一个 3×3 卷积，伪影明显减少，参数量还略降。align_corners设 True 是为了让缩放后的像素对齐，设 False 在边缘会有半像素偏移，小目标分割上能看出差别。

import torch.nn as nn class Up(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv = 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, skip): x = self.up(x) # 尺寸对齐，防止奇数分辨率下 skip 和 up 差一个像素 diff_y = skip.size(2) - x.size(2) diff_x = skip.size(3) - x.size(3) x = nn.functional.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2]) x = torch.cat([skip, x], dim=1) return self.conv(x)

尺寸对齐那几行别省。输入 512×512 时四次下采样刚好整除，但换成 500×500 就会出现 skip 和上采样特征差 1 像素，torch.cat直接报维度错误。用F.pad补齐比裁剪更安全，裁剪会丢边缘信息。

3.2 CrossEntropyLoss 的 ignore_index 与权重传参

CrossEntropyLoss有两个参数在多类别分割里必须显式设置。ignore_index用来屏蔽标注边界或未标注区域，通常设 255，标注时把不确定的像素涂成 255，训练时自动跳过。weight传类别权重张量，要和类别数等长，且放在 GPU 上。

import torch import torch.nn as nn num_classes = 4 weights = torch.tensor([0.1, 1.2, 1.5, 2.0]).cuda() # 按统计结果填 criterion = nn.CrossEntropyLoss(weight=weights, ignore_index=255) # 前向：logits 形状 [B, num_classes, H, W]，target 形状 [B, H, W] logits = model(imgs) # 不要接 softmax loss = criterion(logits, masks)

注意 target 的形状是[B, H, W]而不是 one-hot 的[B, C, H, W]。如果你从别的框架转过来习惯传 one-hot，会报 "Expected target size" 之类的错。另外ignore_index=255要求 mask 里未标注区域确实是 255，如果标注工具用的是 0，那背景和忽略区就混了，得在 Dataset 里手动改。

3.3 评价指标：mIoU 和混淆矩阵怎么算才不骗自己

训练时只看 loss 会被类别不平衡骗。我习惯在每个 epoch 结束后累计混淆矩阵，再算 mIoU 和每类 IoU。混淆矩阵用bincount实现最快：

def update_confusion(conf_mat, pred, target, num_classes, ignore=255): # pred: [B,H,W] 的 argmax 结果, target: [B,H,W] mask = target != ignore pred = pred[mask] target = target[mask] idx = target * num_classes + pred conf_mat += torch.bincount(idx, minlength=num_classes**2).reshape(num_classes, num_classes) return conf_mat def compute_miou(conf_mat): # 每类 IoU = TP / (TP + FP + FN) tp = conf_mat.diag() fp = conf_mat.sum(0) - tp fn = conf_mat.sum(1) - tp iou = tp / (tp + fp + fn + 1e-6) return iou, iou.mean()

bincount的minlength必须设成num_classes**2，否则某些类别没出现时混淆矩阵维度会缩水，reshape 报错。算 mIoU 时背景类要不要算进去，看你的任务——如果背景是有效类别就算，如果只是「非目标」建议单独看前景类的 IoU，背景刷出来的高分没有意义。

4. 训练循环与调参：从 loss 不降到 mIoU 上 0.7

4.1 优化器、学习率与 batch size 的搭配

Unet 分割任务上，Adam 比 SGD 收敛快，初始学习率 1e-3 是安全起点，配合ReduceLROnPlateau在 mIoU 停滞时降半。batch size 受显存限制，512×512 输入下 8GB 显存大概能跑 batch 4，用梯度累积凑等效 batch 16 也行。

import torch.optim as optim from torch.optim.lr_scheduler import ReduceLROnPlateau optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4) scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5) for epoch in range(100): model.train() for imgs, masks in train_loader: imgs, masks = imgs.cuda(), masks.cuda() optimizer.zero_grad() logits = model(imgs) loss = criterion(logits, masks) loss.backward() # 梯度裁剪，Unet 深了容易梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() miou = validate(model, val_loader) scheduler.step(miou) # mode='max' 时传指标 print(f'epoch {epoch}, mIoU {miou:.4f}, lr {optimizer.param_groups[0]["lr"]:.6f}')

weight_decay设 1e-4 是轻量正则，数据量小于 2000 张时可以加到 1e-3。梯度裁剪的max_norm=1.0对 Unet 够用，如果 loss 出现 NaN，先降到 0.5 试试。ReduceLROnPlateau的mode要和监控指标方向一致，监控 mIoU 用max，监控 loss 用min，设反了学习率会往错误方向调。

4.2 数据增强的边界：哪些增强会破坏 mask

几何增强（翻转、旋转、缩放、弹性形变）对分割都安全，只要 image 和 mask 同步。颜色增强（亮度、对比度、HSV 抖动）只作用于 image。有两类增强要慎用：Cutout 和 MixUp。Cutout 在分类里好用，但分割里挖掉的区域 mask 怎么填是个问题，填 255 忽略会浪费样本，填背景会引入噪声。MixUp 把两张图线性叠加，mask 没法线性叠加，多类别下直接失效。

我一般用翻转 + 随机旋转 90 度的倍数 + 亮度抖动三件套，简单可靠。弹性形变对医学和遥感有用，但scipy.ndimage.map_coordinates对 mask 要用order=0，用默认的三次插值会插出小数类别。

4.3 从 loss 曲线判断问题出在哪

loss 不降分几种情况。前 10 个 epoch loss 在 0.6 附近震荡，多半是学习率太大或者 mask 没转成索引图；loss 缓慢下降但 mIoU 不动，检查是不是背景类主导，看混淆矩阵；loss 突然变 NaN，查梯度裁剪和输入归一化，常见的是某张图全黑导致 BN 方差为 0。训练集 loss 降验证集不降，加数据增强或者减模型深度，Unet 在 2000 张以下的数据集上很容易过拟合。

验证时记得model.eval()加torch.no_grad()，忘了 eval 会让 BN 用 batch 统计量，验证结果波动大。这两行是血泪经验，我见过不止一个项目因为漏了 eval 导致验证 mIoU 比实际低 10 个点。

5. 避坑与排查：多类别 Unet 训练里最常见的 5 个翻车现场

5.1 现象：loss 卡在 0.69 附近不动，预测全是背景

原因：标注图是 RGB 三通道，Dataset 里没转索引，CrossEntropyLoss拿到的 target 形状是[B, H, W, 3]或者被错误地当成 3 类，模型学不到有效信号，退化成全预测背景，而背景占比高时 loss 恰好停在-ln(背景频率)附近。

解决：在 Dataset 的__getitem__里打印一次 mask 的shape和np.unique，确认是[H, W]且像素值在[0, num_classes)范围内。是彩色图就先跑 2.1 的转换脚本。

5.2 现象：训练中途报 "Expected target size [B, H, W], got [B, C, H, W]"

原因：target 传成了 one-hot 编码。有些教程为了兼容 Dice Loss 会先转 one-hot，但CrossEntropyLoss要的是类别索引。

解决：要么在 loss 前argmax回索引，要么改用nn.BCEWithLogitsLoss配 one-hot。多类别下推荐前者，CrossEntropy 数值更稳。

5.3 现象：mIoU 在 0.3 到 0.4 之间反复横跳，验证结果不稳定

原因：验证时忘了model.eval()，BN 层用当前 batch 的均值和方差，batch 小的时候统计量噪声大，指标自然抖。

解决：验证循环开头加model.eval()，结束加model.train()，整个验证包在torch.no_grad()里。另外验证集 batch size 可以设大一点，BN 统计更稳。

5.4 现象：小目标类别 IoU 始终为 0，混淆矩阵里该类一行全空

原因：小类像素占比太低，加权不够，或者下采样 4 次后小目标在特征图上只剩一两个像素，skip connection 也救不回来。

解决：先确认标注里小类确实存在且像素数不为零；然后提高该类权重到频率倒数的 2 到 3 倍；如果还不行，减少一次下采样（Unet 改成 3 层），或者用带空洞卷积的变体保持分辨率。7900xtx 这类大显存卡可以开更大的输入尺寸，对小目标帮助明显。

5.5 现象：转 ONNX 后推理结果和 Pytorch 不一致

原因：nn.Upsample的align_corners在 ONNX 里默认行为和 Pytorch 不同，或者动态轴没设对导致尺寸对不上。

解决：导出时显式指定dynamic_axes，align_corners在 Pytorch 和 ONNX 两边都设 True。导出后用onnxruntime跑同一张图，逐像素比对，差异应该在 1e-4 以内。pytorch 转 onnx 这一步别等到部署才做，训练完就验一遍，早发现早改。

6. 把 mIoU 再抬 5 个点的三个具体技巧

第一个技巧是在验证集上做 TTA（测试时增强）。推理时把图翻转、旋转 90 度各跑一次，logits 平均后再 argmax。多类别分割上 TTA 通常能涨 2 到 3 个点，代价是推理时间翻 4 倍。实现上不用改模型，包一层函数就行：

def tta_predict(model, img): # img: [1, 3, H, W] preds = [] for k in range(4): # 0/90/180/270 度旋转 x = torch.rot90(img, k, dims=[2, 3]) with torch.no_grad(): out = model(x) out = torch.rot90(out, -k, dims=[2, 3]) preds.append(out) # 水平翻转再来一轮 x = torch.flip(img, dims=[3]) with torch.no_grad(): out = model(x) preds.append(torch.flip(out, dims=[3])) return torch.stack(preds).mean(0)

注意旋转用torch.rot90而不是transforms，后者会引入插值误差。平均的是 logits 不是 softmax 后的概率，logits 平均再 softmax 数值更稳。

第二个技巧是用 Lovasz Loss 或 Dice Loss 和 CrossEntropy 组合。CrossEntropy 是逐像素的，对 IoU 这种区域指标是间接优化。Lovasz 直接优化 IoU 的凸 surrogate，小类上提升明显。组合方式一般是loss = 0.5 * CE + 0.5 * Lovasz，权重按验证集 mIoU 微调。Lovasz 对 ignore_index 的支持要自己处理，把忽略像素的梯度置零。

第三个技巧是冻结编码器做 warmup。如果用了预训练的 ResNet 或 EfficientNet 做编码器，前 5 个 epoch 只训解码器，学习率 1e-3，之后再解冻全部用 1e-4 微调。这样能避免随机初始化的解码器在早期把预训练编码器的特征带偏。我自己的习惯是 warmup 阶段监控解码器 loss，降到 0.3 以下再解冻，比固定 epoch 数更靠谱。

最后说个验证方法：别只看 mIoU 一个数。把验证集里 IoU 最低的 10 张图挑出来可视化，看是标注错了还是模型真不行。我做过的一个路面裂缝项目，mIoU 卡在 0.62，可视化后发现是标注把阴影也标成了裂缝，重新标了 50 张后直接到 0.71。数据质量的天花板比模型结构高得多，调参之前先把数据过一遍，这是我踩了无数坑之后养成的习惯。希望帮到你。

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