简介面向希望在 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, palettePALETTE): rgb np.array(Image.open(mask_path).convert(RGB), dtypenp.uint8) label np.full(rgb.shape[:2], 255, dtypenp.uint8) # 255 ignore for idx, color in enumerate(palette): # 三个通道同时相等才算命中避免把近似色误判 hit np.all(rgb np.array(color, dtypenp.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_index255): counts np.zeros(num_classes, dtypenp.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], minlengthnum_classes) total counts.sum() for i, c in enumerate(counts): print(fclass {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 和 maskimport albumentations as A from albumentations.pytorch import ToTensorV2 train_tf A.Compose([ A.RandomResizedCrop(256, 256, scale(0.5, 1.0)), A.HorizontalFlip(p0.5), A.RandomBrightnessContrast(p0.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(imageimage, maskmask) # mask 转 longCrossEntropy 只接受 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_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, kernel_size3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)padding1配kernel_size3是保持 H、W 不变的标准组合UNet 靠这个让跳跃连接两端的尺寸天然对齐。biasFalse是因为后面紧跟 BNBN 的 beta 参数已经承担了偏置作用再加 bias 属于冗余参数显存和收敛速度都吃亏。ReLU(inplaceTrue)省显存但如果后面要对同一份激活做多次读取比如某些注意力分支inplace 会改掉原值导致结果错误换成nn.ReLU()更安全。out_ch的选择上原始 UNet 是 64-128-256-512-1024自己数据集通常只有几百上千张图直接照搬 1024 那层会严重过拟合显存也顶不住。我一般在base_ch32或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_size2, stride2) 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, sizeskip.shape[-2:], modebilinear, align_cornersFalse ) x torch.cat([skip, x], dim1) # 沿通道维拼接不是相加 return self.conv(x)这段里的尺寸对齐判断是必须的。输入 256×256 时一路 pool 下去都是偶数不会出问题一旦换成 250×250 或者做了随机裁剪没对齐第四次下采样后会变成 15×15上采样回来是 30×30而 skip 是 31×31torch.cat直接报维度错误。加一层interpolate兜底比在 dataset 里强行 pad 到 16 的倍数更省事。用torch.cat而不是相加是 UNet 和 FCN 的核心差别拼接保留了编码器的高分辨率细节解码器自己学怎么融合而不是强迫两边特征在同一空间里对齐。代价是通道数翻倍显存占用明显上升这也是为什么base_ch不能设太大的直接原因。上采样用ConvTranspose2d还是interpolate conv是个常见分歧。转置卷积可学习、表达能力强但 stride2、kernel2 时容易出现棋盘格伪影表现为分割图上有规律的网格纹理。我一般训练阶段用转置卷积推理时如果看到网格就换成F.interpolate(modebilinear) DoubleConv参数量少一点边缘更平滑。3.3 输出层与类别数从单通道 sigmoid 改成 N 通道 softmax网络主体搭好后最后接一个 1×1 卷积把通道数压到类别数class UNet(nn.Module): def __init__(self, in_ch3, num_classes4, base_ch64): 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_size1) 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少算一个会导致训练时索引越界或者某一类永远学不到。输出层不要加 softmaxnn.CrossEntropyLoss内部已经做了 log_softmax NLL外面再加一层等于算了两遍梯度会被压平收敛明显变慢推理时再手动argmax或softmax。二分类改多分类时原来用BCEWithLogitsLoss 单通道输出的写法要整体换掉不能只改通道数。另外BatchNorm2d在小 batch 下统计量不稳训练时 batch 设到 4 以下BN 的 running mean 会剧烈抖动验证指标跟着跳。显存不够时优先降分辨率或减base_ch尽量别把 batch 压到 2。实在压不下去把 BN 换成GroupNorm(num_groups8, num_channelsout_ch)小 batch 下稳定得多代价是收敛稍慢。4. 训练、验证与推理从 DataLoader 到彩色分割图的完整链路4.1 损失函数与 mIoU多类别分割盯这两个数单用交叉熵在类别极不平衡时会被背景主导模型倾向于把所有像素判成背景。组合一个 Dice 项能明显改善前景类的召回class CombinedLoss(nn.Module): def __init__(self, num_classes, ignore_index255, ce_w0.5, dice_w0.5): super().__init__() self.ce nn.CrossEntropyLoss(ignore_indexignore_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, dim1) 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 * diceignore_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_index255): valid target ! ignore_index pred, target pred[valid], target[valid] idx target * num_classes pred conf torch.bincount(idx, minlengthnum_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(min1) return iou.mean().item(), iou # 返回总 mIoU 和每类 IoUconf.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_noneTrue) with autocast(device_typecuda, dtypetorch.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_noneTrue)比默认的置零更快梯度张量直接释放而不是填 0。AMP 里scaler.scale(loss).backward()和scaler.step必须成对出现中间插一句普通loss.backward()会直接报错。autocast只包前向和 loss反向由 scaler 接管别把backward也包进去。优化器一般用AdamW(lr1e-3, weight_decay1e-4)起步比 SGD 收敛快小数据集上更友好。学习率调度用余弦退火scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxepochs, eta_min1e-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, dtypetorch.int64) for image, mask in loader: image, mask image.to(device), mask.to(device) logits model(image) pred logits.argmax(dim1) 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, }, fruns/{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, dim1) pred prob.argmax(dim1).squeeze(0).cpu().numpy().astype(np.uint8) return predunsqueeze(0)补 batch 维argmax(dim1)沿类别维取最大值输出的就是 0 到 N-1 的索引图。上色用同一张调色板反过来查PALETTE_ARR np.array(PALETTE, dtypenp.uint8) def colorize(label): return PALETTE_ARR[label] # (H, W, 3) def overlay(image_rgb, color_mask, alpha0.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 outPALETTE_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 计算强制转到 fp32prob torch.softmax(logits.float(), dim1) one_hot one_hot.float()同时把eps从1e-6提到1e-4。如果 NaN 出现在训练初期优先把学习率从1e-3降到3e-4再开启梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.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-size2g。另外pin_memoryTrue只在 GPU 训练时开CPU 训练开着反而拖慢。5.5 加载保存的权重后模型输出和训练时不一样现象是torch.load之后load_state_dict报一堆 missing keys 和 unexpected keys或者加载成功但 mIoU 掉一大截。原因通常是推理时新建的网络num_classes和训练时不一致或者base_ch忘了改。load_state_dict(strictFalse)能强行加载但那些没匹配上的层还是随机初始化结果自然不对。解决办法就是 4.3 节的写法保存时把num_classes、base_ch一起存进字典推理时读出来再建网络ckpt torch.load(runs/xxx/best.pth, map_locationcpu) model UNet(num_classesckpt[num_classes], base_chckpt[base_ch]) model.load_state_dict(ckpt[model])map_locationcpu别忘了在没 GPU 的机器上直接 load 会报 CUDA 设备错误。6. 二次开发改造换数据集、加类别、换主干的三条路线这套代码的扩展点其实就三个改之前先想清楚动哪一层改造目标要改的文件关键注意点换自己的数据集dataset.py、palette.py颜色表顺序必须和 mask 一致ignore_index统一用 255增加/减少类别palette.py、训练命令的num_classes输出层通道数自动跟着变旧权重不能直接复用换主干网络unet.py保持forward返回(B, C, H, W)的 logits跳跃连接的通道数要重新对齐换数据集是最常见的场景。我的做法是先把新数据的 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, strictFalse)这样编码器和解码器沿用旧权重输出层重新初始化通常比从头训收敛快一倍以上。新加的两类如果样本很少前几个 epoch 可以冻结编码器只训解码器等 loss 稳住再放开。换主干网络是提升精度的主要手段。UNet 的编码器可以换成 ResNet34 或 EfficientNet 的预训练权重把self.inc到self.d4替换掉解码器保持不变。这时候要注意两点一是预训练主干的输出通道数ResNet34 是 512和你解码器的输入通道要手动对齐二是预训练权重的归一化参数是 ImageNet 的 mean/stddataset 里的Normalize必须改成同一套否则预训练特征会被打乱反而不如从零训。验证改造是否有效我一般固定三件事同一个验证集、同一个随机种子、同样的 epoch 数。只改一个变量跑对比mIoU 涨了才算数。有次我同时换了损失权重和数据增强mIoU 涨了 5 个点回头单独跑发现是增强起的作用损失权重那部分其实是负收益——从那以后我每次做对比实验都强制只动一个开关跑完记在runs/目录的notes.md里省得两周后自己都不记得改了什么。这套流程看着笨但比事后猜要快得多。希望帮到你。本文还有配套的精品资源点击获取