简介：面向希望在 PyTorch 中完成多类别语义分割实践的开发者与学习者，内容围绕 U-Net 的编码器与解码器构建、跳跃连接、多通道输出、交叉熵损失以及训练评估流程展开，适合医学影像、遥感图像等像素级分类场景的入门学习与代码复现。压缩包共 46 个文件，以 19 个 Python 源码和 24 个缓存文件为主，另含 2 个文本说明与 1 个 JSON 配置，整体约 69KB；目录涵盖数据加载、自定义增强、模型定义、损失函数、评价指标、学习率调度和训练结果保存等模块，结构清晰，便于按训练链路逐层查阅。已有 15284 人学习下载，可作为多类别 U-Net 项目的参考骨架，帮助读者理解数据集组织、权重计算、训练曲线记录与分割结果可视化等实现细节，也能为排查通道数设置、标签映射和评估指标计算等问题提供思路。

1. 拿自己的多类别数据集训一个 UNet：这份工程到底解决了什么

大多数人第一次接触语义分割，都是拿现成数据集跑一遍别人的脚本，指标出来了，换成自己标注的几十上百张图就立刻卡住：mask 是三通道 RGB 对不上类别编号、背景类被算进 loss、训十几个 epoch 预测出来全是同一种颜色。这份工程解决的就是这一段——它不推导 UNet 论文，而是把 PyTorch 下 UNet 做多类别语义分割的完整链路摊开：自己的数据集怎么组织、调色板 mask 怎么转成类别索引、类别数怎么从 1 改成 N、loss 和 mIoU 怎么算、训完怎么把 logits 还原成能看的彩色图。适合手里已经有一批标注图、想把分割跑通并做二次开发的开发者，也适合学完 PyTorch 基础想找个完整工程练手的人。多类别数据集是这套代码的核心变量，改类别数、改调色板、改损失权重都围绕它展开。

2. 工程骨架与数据管线：多类别 mask 从磁盘到张量的全过程

2.1 目录结构：数据、模型、训练、推理四层分开

我一般把这类项目拆成四层，好处是换数据集只动datasets/，换网络只动src/unet.py，训练脚本几乎不用碰。下面是我习惯的目录形态：

unet-multiclass/ ├── datasets/ │ ├── images/ │ │ ├── train/ # 原图，jpg/png 都行 │ │ └── val/ │ └── masks/ │ ├── train/ # 与 images/train 同名，后缀可以不同 │ └── val/ ├── src/ │ ├── dataset.py # Dataset + 增强 │ ├── unet.py # 网络定义 │ ├── losses.py # CE / Dice / 组合损失 │ ├── metrics.py # 混淆矩阵与 mIoU │ ├── train.py # 训练入口 │ ├── predict.py # 推理入口 │ └── palette.py # 类别颜色表 ├── runs/ # 权重与日志输出 └── requirements.txt

图像和 mask 放在两个平行目录、用同名文件对应，是这套结构里最省心的约定：dataset.py只需要把路径里的images替换成masks就能找到标签，不用维护额外的 csv。验证集一定是单独目录而不是训练时随机切，原因在 2.3 节会说。runs/目录按时间戳建子目录，每次训练一个文件夹，里面放best.pth、last.pth和一个train.log，回头对比实验时不用靠记忆。

requirements.txt里通常会锁这几样：torch、torchvision、numpy、Pillow、albumentations、opencv-python、tqdm、tensorboard。版本不用写死到小版本，但 torch 和 torchvision 要成对，torchvision 版本高于 torch 会导致 import 报符号缺失，这种玄学问题在装环境阶段最容易浪费时间。

2.2 多类别 mask 的读取：别把 RGB 掩码直接喂给 CrossEntropy

多类别分割翻车最多的一步就在这。标注工具导出的掩码几乎都是调色板 PNG 或 RGB 图，每个类别是一种颜色，而nn.CrossEntropyLoss要的是取值范围[0, num_classes-1]的单通道整型张量。直接把三通道 mask 塞进去，要么维度报错，要么被当成三个独立通道的回归问题，loss 看着在降，预测却完全不对。

常见做法是建一张颜色表，逐像素比对颜色转成类别索引：

import numpy as np from PIL import Image # 调色板第 i 个元素就是类别 i 的颜色，(0,0,0) 一般留给背景 PALETTE = [ (0, 0, 0), # 0 background (255, 0, 0), # 1 class_a (0, 255, 0), # 2 class_b (0, 0, 255), # 3 class_c ] def rgb_mask_to_label(mask_path, palette=PALETTE): rgb = np.array(Image.open(mask_path).convert("RGB"), dtype=np.uint8) label = np.full(rgb.shape[:2], 255, dtype=np.uint8) # 255 = ignore for idx, color in enumerate(palette): # 三个通道同时相等才算命中，避免把近似色误判 hit = np.all(rgb == np.array(color, dtype=np.uint8), axis=-1) label[hit] = idx return label

这段代码有三个参数值得说清。palette的顺序就是类别编号的顺序，palette[0]必须是背景，因为大多数标注工具默认先给背景分配颜色。255是ignore_index，凡是颜色表里找不到的像素都标成 255，训练时会被 loss 忽略；如果数据里存在标注错误或边缘过渡色，这个兜底比强行归到某一类安全得多。convert("RGB")不能省，某些调色板 PNG 直接读出来是P模式，强转 numpy 会得到索引值而不是颜色，转成 RGB 后再比对才稳定。

如果标注工具导出的是索引图（P模式），还有更快的路子：直接np.array(Image.open(path))拿到索引，再按调色板映射表查一次。但索引图和颜色表强绑定，换工具就失效，所以我一般还是走 RGB 比对，慢一点但不会错。

2.3 训练/验证划分与类别不平衡：先统计再动手

多类别数据集最典型的特征是不平衡。背景占 70%、目标类占 3% 是常态，模型只要把所有像素预测成背景，准确率就能到 70%，看着很漂亮，实际完全不可用。所以在写 Dataset 之前先做一次像素占比统计，这个习惯能省掉后面大量返工：

import numpy as np from pathlib import Path def class_pixel_ratio(mask_dir, num_classes, ignore_index=255): counts = np.zeros(num_classes, dtype=np.int64) for p in Path(mask_dir).glob("*.png"): label = rgb_mask_to_label(str(p)) valid = label != ignore_index counts += np.bincount(label[valid], minlength=num_classes) total = counts.sum() for i, c in enumerate(counts): print(f"class {i}: {c / total:.4%}") return counts

统计结果直接决定后面两件事：一是ignore_index之外的类别要不要加权，二是验证集怎么切。加权可以在CrossEntropyLoss(weight=...)里传入按频率倒数归一化的权重，也可以交给 Dice 项去平衡，两种都行，但别同时猛调，容易把梯度拉爆。

划分上的血泪经验是：验证集必须按场景或采集批次切，不能随机抽图。同一段视频连续抽帧、同一个场景连拍几十张，随机切分会让训练集和验证集共享几乎相同的画面，mIoU 虚高十几个点，上线就原形毕露。如果数据本身是连续采集的，按文件序号每 10 张取 1 张进验证集，比random_split靠谱。

2.4 数据增强：图像和 mask 必须同步变换

分割的增强和分类不一样，图像翻转、裁剪、缩放，mask 都要跟着做同样的几何变换，而且 mask 只能最近邻插值。用torchvision.transforms分别处理很容易踩坑，我一般用 albumentations，它在一次调用里同时处理 image 和 mask：

import albumentations as A from albumentations.pytorch import ToTensorV2 train_tf = A.Compose([ A.RandomResizedCrop(256, 256, scale=(0.5, 1.0)), A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.3), # 只作用于 image，不影响 mask A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ]) val_tf = A.Compose([ A.Resize(256, 256), A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), ToTensorV2(), ])

A.Compose的几何类变换会自动作用到 mask 上，插值方式由库里控制成最近邻，不会凭空造出类别 1.5 这种值。颜色类变换只作用于 image，这是它比手写 transform 省事的地方。RandomResizedCrop的scale=(0.5, 1.0)对小目标数据集要慎重，裁太狠会把仅有的几个目标裁没，我通常先设(0.8, 1.0)训一版看 mIoU 再决定要不要放开。验证集只做 Resize 和 Normalize，不加任何随机项——验证结果要可复现，随机增强会让每个 epoch 的 mIoU 抖动，没法判断模型到底有没有变好。

Dataset 里的__getitem__就是把这套流程串起来：

class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, transform): self.img_paths = sorted(Path(img_dir).glob("*.jpg")) self.mask_dir = Path(mask_dir) self.transform = transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img_path = self.img_paths[idx] image = np.array(Image.open(img_path).convert("RGB")) mask = rgb_mask_to_label(str(self.mask_dir / (img_path.stem + ".png"))) out = self.transform(image=image, mask=mask) # mask 转 long，CrossEntropy 只接受 int64 return out["image"], out["mask"].long()

sorted不能省，否则不同进程的 DataLoader 拿到顺序不一致，调试时对不上号。mask.long()也必须加，albumentations 的ToTensorV2给 mask 的默认类型是 uint8，直接进 loss 会报 dtype 不匹配。

3. UNet 网络拆解：编码器、跳跃连接与解码器的逐层实现

3.1 双卷积块：UNet 的最小单元

UNet 的每个尺度都是「卷积-归一化-激活」重复两次，把这个单元抽出来，整个网络就是它的堆叠：

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, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_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是保持 H、W 不变的标准组合，UNet 靠这个让跳跃连接两端的尺寸天然对齐。bias=False是因为后面紧跟 BN，BN 的 beta 参数已经承担了偏置作用，再加 bias 属于冗余参数，显存和收敛速度都吃亏。ReLU(inplace=True)省显存，但如果后面要对同一份激活做多次读取（比如某些注意力分支），inplace 会改掉原值导致结果错误，换成nn.ReLU()更安全。

out_ch的选择上，原始 UNet 是 64-128-256-512-1024，自己数据集通常只有几百上千张图，直接照搬 1024 那层会严重过拟合，显存也顶不住。我一般在base_ch=32或64之间选，深层通道数封顶 512，参数量从 31M 降到 7M 左右，小数据集上 mIoU 反而更高。

3.2 下采样与上采样拼接：尺寸对齐是最容易翻车的地方

编码器每层先 maxpool 再双卷积，解码器每层先上采样再和对应尺度的特征拼接：

class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_ch, out_ch)) def forward(self, x): return self.block(x) class Up(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() # 转置卷积把通道数减半、尺寸翻倍 self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2) self.conv = DoubleConv(in_ch, out_ch) # 拼接后通道翻倍，所以是 in_ch def forward(self, x, skip): x = self.up(x) # 输入边长不是 16 的整数倍时，上采样后可能和 skip 差 1 像素 if x.shape[-2:] != skip.shape[-2:]: x = nn.functional.interpolate( x, size=skip.shape[-2:], mode="bilinear", align_corners=False ) x = torch.cat([skip, x], dim=1) # 沿通道维拼接，不是相加 return self.conv(x)

这段里的尺寸对齐判断是必须的。输入 256×256 时一路 pool 下去都是偶数，不会出问题；一旦换成 250×250 或者做了随机裁剪没对齐，第四次下采样后会变成 15×15，上采样回来是 30×30，而 skip 是 31×31，torch.cat直接报维度错误。加一层interpolate兜底，比在 dataset 里强行 pad 到 16 的倍数更省事。

用torch.cat而不是相加，是 UNet 和 FCN 的核心差别：拼接保留了编码器的高分辨率细节，解码器自己学怎么融合，而不是强迫两边特征在同一空间里对齐。代价是通道数翻倍，显存占用明显上升，这也是为什么base_ch不能设太大的直接原因。

上采样用ConvTranspose2d还是interpolate + conv，是个常见分歧。转置卷积可学习、表达能力强，但 stride=2、kernel=2 时容易出现棋盘格伪影，表现为分割图上有规律的网格纹理。我一般训练阶段用转置卷积，推理时如果看到网格，就换成F.interpolate(mode="bilinear") + DoubleConv，参数量少一点，边缘更平滑。

3.3 输出层与类别数：从单通道 sigmoid 改成 N 通道 softmax

网络主体搭好后，最后接一个 1×1 卷积把通道数压到类别数：

class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=4, base_ch=64): super().__init__() self.inc = DoubleConv(in_ch, base_ch) self.d1 = Down(base_ch, base_ch * 2) self.d2 = Down(base_ch * 2, base_ch * 4) self.d3 = Down(base_ch * 4, base_ch * 8) self.d4 = Down(base_ch * 8, base_ch * 16) self.u1 = Up(base_ch * 16, base_ch * 8) self.u2 = Up(base_ch * 8, base_ch * 4) self.u3 = Up(base_ch * 4, base_ch * 2) self.u4 = Up(base_ch * 2, base_ch) self.out = nn.Conv2d(base_ch, num_classes, kernel_size=1) def forward(self, x): x1 = self.inc(x) x2 = self.d1(x1) x3 = self.d2(x2) x4 = self.d3(x3) x5 = self.d4(x4) x = self.u1(x5, x4) x = self.u2(x, x3) x = self.u3(x, x2) x = self.u4(x, x1) return self.out(x) # 返回 logits，不在网络里做 softmax

三个地方是改造重点。num_classes要包含背景，四类目标加背景就是 5，少算一个会导致训练时索引越界或者某一类永远学不到。输出层不要加 softmax，nn.CrossEntropyLoss内部已经做了 log_softmax + NLL，外面再加一层等于算了两遍，梯度会被压平，收敛明显变慢；推理时再手动argmax或softmax。二分类改多分类时，原来用BCEWithLogitsLoss+ 单通道输出的写法要整体换掉，不能只改通道数。

另外，BatchNorm2d在小 batch 下统计量不稳，训练时 batch 设到 4 以下，BN 的 running mean 会剧烈抖动，验证指标跟着跳。显存不够时优先降分辨率或减base_ch，尽量别把 batch 压到 2。实在压不下去，把 BN 换成GroupNorm(num_groups=8, num_channels=out_ch)，小 batch 下稳定得多，代价是收敛稍慢。

4. 训练、验证与推理：从 DataLoader 到彩色分割图的完整链路

4.1 损失函数与 mIoU：多类别分割盯这两个数

单用交叉熵，在类别极不平衡时会被背景主导，模型倾向于把所有像素判成背景。组合一个 Dice 项能明显改善前景类的召回：

class CombinedLoss(nn.Module): def __init__(self, num_classes, ignore_index=255, ce_w=0.5, dice_w=0.5): super().__init__() self.ce = nn.CrossEntropyLoss(ignore_index=ignore_index) self.num_classes = num_classes self.ignore_index = ignore_index self.ce_w, self.dice_w = ce_w, dice_w def forward(self, logits, target): ce = self.ce(logits, target) # Dice 只在有效像素上算，ignore 区域不参与 prob = torch.softmax(logits, dim=1) valid = target != self.ignore_index target_safe = target.clone() target_safe[~valid] = 0 one_hot = nn.functional.one_hot( target_safe, self.num_classes ).permute(0, 3, 1, 2).float() one_hot = one_hot * valid.unsqueeze(1).float() dims = (0, 2, 3) inter = (prob * one_hot).sum(dims) union = prob.sum(dims) + one_hot.sum(dims) dice = 1 - ((2 * inter + 1e-6) / (union + 1e-6)).mean() return self.ce_w * ce + self.dice_w * dice

ignore_index在 Dice 项里要手动屏蔽，CrossEntropyLoss会自己忽略，但 one-hot 会把 255 当成非法索引直接抛错，所以先clone再把无效位置置 0，然后乘 mask 把贡献清零。1e-6是数值稳定项，防止某类在当前 batch 里一个像素都没有导致除零。ce_w和dice_w各 0.5 是起点，如果背景类占 90% 以上，把 dice 权重提到 0.7，前景的 IoU 通常能涨几个点。

评估指标用混淆矩阵算 mIoU，比逐 batch 平均靠谱：

def update_confusion(conf, pred, target, num_classes, ignore_index=255): 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 miou_from_confusion(conf): inter = torch.diag(conf).float() union = conf.sum(0).float() + conf.sum(1).float() - inter iou = inter / union.clamp(min=1) return iou.mean().item(), iou # 返回总 mIoU 和每类 IoU

conf.sum(0)是预测为该类的像素总数，conf.sum(1)是真实为该类的像素总数，减去交集就是并集。返回每类 IoU 很有必要，mIoU 涨了但某一类掉到 0.1，说明模型在拿大类换小类，光看总数发现不了。

4.2 训练循环：AMP、优化器、学习率调度

训练脚本的主干不长，容易出错的是几个开关的位置：

import torch from torch.utils.data import DataLoader from torch.amp import autocast, GradScaler def train_one_epoch(model, loader, optimizer, criterion, scaler, device): model.train() total_loss = 0.0 for image, mask in loader: image, mask = image.to(device), mask.to(device) optimizer.zero_grad(set_to_none=True) with autocast(device_type="cuda", dtype=torch.float16): logits = model(image) loss = criterion(logits, mask) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() total_loss += loss.item() return total_loss / len(loader)

optimizer.zero_grad(set_to_none=True)比默认的置零更快，梯度张量直接释放而不是填 0。AMP 里scaler.scale(loss).backward()和scaler.step必须成对出现，中间插一句普通loss.backward()会直接报错。autocast只包前向和 loss，反向由 scaler 接管，别把backward也包进去。

优化器一般用AdamW(lr=1e-3, weight_decay=1e-4)起步，比 SGD 收敛快，小数据集上更友好。学习率调度用余弦退火：

scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=epochs, eta_min=1e-6 )

调度器在每个 epoch 结束后调一次scheduler.step()，不要每次迭代都调，否则学习率掉得飞快，十几个 epoch 后就几乎不更新了。

设备选择上，device = torch.device("cuda" if torch.cuda.is_available() else "cpu")，然后model.to(device)。DataLoader 的num_workers在 Linux 下设 4 到 8 都行，Windows 下必须放在if __name__ == "__main__":里，否则多进程会递归启动直接把内存吃光，这个坑几乎人人都踩过一次。

4.3 验证与权重保存：只存 mIoU 最好的那一份

验证阶段要关掉梯度、切 eval 模式，否则 BN 会用当前 batch 的统计量，结果不可复现：

@torch.no_grad() def validate(model, loader, num_classes, device): model.eval() conf = torch.zeros(num_classes, num_classes, dtype=torch.int64) for image, mask in loader: image, mask = image.to(device), mask.to(device) logits = model(image) pred = logits.argmax(dim=1) conf = update_confusion(conf.cpu(), pred.cpu(), mask.cpu(), num_classes) return miou_from_confusion(conf)

model.eval()和torch.no_grad()一个都不能少，前者切 BN 和 Dropout 的行为，后者省显存。混淆矩阵累积在 CPU 上算，避免验证过程中显存越占越多。

保存逻辑很直接：

miou, per_class = validate(model, val_loader, num_classes, device) if miou > best_miou: best_miou = miou torch.save({ "model": model.state_dict(), "miou": miou, "num_classes": num_classes, "base_ch": base_ch, }, f"runs/{run_id}/best.pth")

存字典而不是裸的state_dict，把num_classes、base_ch一起塞进去。后面推理时直接读这两个值建网络，不用回去翻训练命令，也不会出现「网络结构对不上、权重加载一半报 key 不匹配」的情况。这个习惯在同时跑好几个实验时特别值钱。

4.4 推理与上色：把 logits 变成能贴回原图的结果

推理就是训练验证那套去掉标签：

@torch.no_grad() def predict(model, image_tensor, device, num_classes): model.eval() logits = model(image_tensor.unsqueeze(0).to(device)) prob = torch.softmax(logits, dim=1) pred = prob.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.uint8) return pred

unsqueeze(0)补 batch 维，argmax(dim=1)沿类别维取最大值，输出的就是 0 到 N-1 的索引图。上色用同一张调色板反过来查：

PALETTE_ARR = np.array(PALETTE, dtype=np.uint8) def colorize(label): return PALETTE_ARR[label] # (H, W, 3) def overlay(image_rgb, color_mask, alpha=0.5): # 背景类不上色，避免整张图被黑色压暗 fg = color_mask.sum(-1) > 0 out = image_rgb.copy() out[fg] = (image_rgb[fg] * (1 - alpha) + color_mask[fg] * alpha).astype(np.uint8) return out

PALETTE_ARR[label]是 numpy 的高级索引，一步完成 (H,W) 到 (H,W,3) 的映射，比循环快几个数量级。overlay里先把背景排除掉，只对前景做半透明混合，否则整张图会被黑色调暗，看不清原图细节。

如果推理图和训练图分辨率不一致，别直接 resize 到 256 再预测，还原回来边缘会糊。常见做法是滑窗推理：按 256 的窗口带重叠切块，每块预测后在重叠区取概率平均，最后拼回原尺寸。写起来多几十行，但大图上的边缘质量差得很明显。

5. 排查现场：多类别分割最容易翻车的五处

5.1 训练 loss 正常下降，预测却全是同一种颜色

现象是 loss 从 1.2 稳降到 0.3，验证 mIoU 也有 0.6 左右，但推理出来的图整片是背景色。原因通常有两个：一是 mask 转换时颜色表顺序和训练时不一致，类别编号错位，模型学到的其实是另一套映射；二是上色时用了PALETTE[label]，但label是 float 类型，numpy 索引报隐式转换，落到了错误位置。解决办法是把训练和推理的PALETTE抽到同一个palette.py里 import，推理前打印np.unique(pred)确认取值范围确实在[0, num_classes-1]之间，再检查label.dtype是不是 uint8。

还有一种更隐蔽的情况：验证集 mIoU 正常，但推理脚本忘了model.eval()，BN 用单张图的统计量，输出全乱。这种问题不看代码基本查不出来，我现在习惯在predict函数第一行就写model.eval()。

5.2 loss 变成 NaN，而且是从第几个 batch 突然开始的

现象是前几十个 step 正常，某个 batch 之后 loss 直接 NaN，之后再也回不来。原因多半是学习率过高或 Dice 项除零。Dice 里如果某个类别在当前 batch 中既没有预测也没有真实像素，union为 0，即使加了1e-6，在 fp16 下也可能下溢成 0。解决办法是把 Dice 计算强制转到 fp32：

prob = torch.softmax(logits.float(), dim=1) one_hot = one_hot.float()

同时把eps从1e-6提到1e-4。如果 NaN 出现在训练初期，优先把学习率从1e-3降到3e-4，再开启梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)。AMP 下还要确认GradScaler用得对，跳过scaler.step直接optimizer.step()会让未缩放的梯度进网络，几轮就爆。

5.3 mIoU 忽高忽低，验证集有时比训练集还高

这不算 bug，但会让实验判断完全失准。原因通常是验证集太小（只有十几张图），加上验证时忘了model.eval()，或者 BN 的 running stats 还没稳定。另一个常见原因是验证集和训练集来自同一段连续采集，2.3 节说过的场景泄漏。

解决办法：验证集至少占到总数据的 15%，且按场景切分；验证阶段严格eval()+no_grad()；如果数据量实在太小，用 5 折交叉验证取平均，单次结果不做结论。另外别每个 epoch 都存权重，只存 mIoU 最优的那份，否则最后拿到的是过拟合最严重的模型。

5.4 DataLoader 报错或者直接把内存吃满

现象是训练启动几秒后进程被杀，或者报RuntimeError: DataLoader worker (pid xxx) is killed by signal。在 Windows 上是num_workers > 0且没写if __name__ == "__main__":保护，多进程递归启动。在 Linux 上多半是共享内存不够，容器环境默认/dev/shm只有 64MB，几个 worker 一跑就满。

解决办法：Windows 下把训练入口包进if __name__ == "__main__":；Linux 容器里把num_workers降到 2 或 0，或者给容器加--shm-size=2g。另外pin_memory=True只在 GPU 训练时开，CPU 训练开着反而拖慢。

5.5 加载保存的权重后，模型输出和训练时不一样

现象是torch.load之后load_state_dict报一堆 missing keys 和 unexpected keys，或者加载成功但 mIoU 掉一大截。原因通常是推理时新建的网络num_classes和训练时不一致，或者base_ch忘了改。load_state_dict(strict=False)能强行加载，但那些没匹配上的层还是随机初始化，结果自然不对。

解决办法就是 4.3 节的写法：保存时把num_classes、base_ch一起存进字典，推理时读出来再建网络：

ckpt = torch.load("runs/xxx/best.pth", map_location="cpu") model = UNet(num_classes=ckpt["num_classes"], base_ch=ckpt["base_ch"]) model.load_state_dict(ckpt["model"])

map_location="cpu"别忘了，在没 GPU 的机器上直接 load 会报 CUDA 设备错误。

6. 二次开发改造：换数据集、加类别、换主干的三条路线

这套代码的扩展点其实就三个，改之前先想清楚动哪一层：

换数据集是最常见的场景。我的做法是先把新数据的 mask 全部过一遍rgb_mask_to_label，统计有多少像素落到了 255，超过 5% 就说明颜色表没配全，得回去补。然后跑一次 2.3 节的占比统计，看类别分布，再决定要不要给CrossEntropyLoss传weight。

加类别要特别小心旧权重。比如从 4 类扩到 6 类，self.out的权重形状从(4, 64, 1, 1)变成(6, 64, 1, 1)，load_state_dict必然报 size mismatch。可行的做法是加载时过滤掉输出层：

state = torch.load("runs/old/best.pth")["model"] state = {k: v for k, v in state.items() if not k.startswith("out.")} model.load_state_dict(state, strict=False)

这样编码器和解码器沿用旧权重，输出层重新初始化，通常比从头训收敛快一倍以上。新加的两类如果样本很少，前几个 epoch 可以冻结编码器只训解码器，等 loss 稳住再放开。

换主干网络是提升精度的主要手段。UNet 的编码器可以换成 ResNet34 或 EfficientNet 的预训练权重，把self.inc到self.d4替换掉，解码器保持不变。这时候要注意两点：一是预训练主干的输出通道数（ResNet34 是 512）和你解码器的输入通道要手动对齐；二是预训练权重的归一化参数是 ImageNet 的 mean/std，dataset 里的Normalize必须改成同一套，否则预训练特征会被打乱，反而不如从零训。

验证改造是否有效，我一般固定三件事：同一个验证集、同一个随机种子、同样的 epoch 数。只改一个变量跑对比，mIoU 涨了才算数。有次我同时换了损失权重和数据增强，mIoU 涨了 5 个点，回头单独跑发现是增强起的作用，损失权重那部分其实是负收益——从那以后我每次做对比实验都强制只动一个开关，跑完记在runs/目录的notes.md里，省得两周后自己都不记得改了什么。这套流程看着笨，但比事后猜要快得多。希望帮到你。

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