PyTorch UNet多类别语义分割实战：数据管线、损失与mIoU评估

简介面向需要在PyTorch中落地多类别语义分割的开发者与学习者这份工程围绕Unet的编码器—解码器结构、跳跃连接和多通道输出展开可用于医学影像、遥感图像等像素级分类场景。资源包共46个文件以19个py源码和24个pyc缓存文件为主另含2个txt与1个json配置或说明文件压缩包约69KB涵盖模型定义、数据加载与增强、损失函数、评价指标、训练流程、学习率调度及结果保存等模块目录组织便于按功能查阅。已有15287人学习说明其在入门实践与代码参考方面具有一定认可度。读者可借此理解多类别输出通道设置、交叉熵损失使用、IoU与混淆矩阵评估、训练曲线可视化等关键环节并以此为骨架迁移到自有数据集补充注意力机制、多尺度训练或更强编码器等改进思路节省从零搭建与排错的时间。1. 从一批多类别掩膜说起PyTorch 版 UNet 语义分割能解决什么手里有一批自己标注的图每张对应一张掩膜掩膜里像素值是 0、1、2、3……分别代表背景、病灶、器官、器械之类的类别。想训一个能逐像素分类的模型又不想被框架和工程细节拖住——这正是 PyTorch 下实现 UNet 多类别语义分割最典型的落点。这个工程做的事情很具体读你自己的图像和掩膜按类别编码搭一个标准 UNet用交叉熵把多类别分割训起来最后输出 mIoU 和彩色预测图。它适合两类人一类是刚拿到标注数据、需要一个能直接改路径就跑的基线另一类是已经跑过二分类分割想扩到多类别、要处理忽略标签和类别不平衡的开发者。它不是通用框架代码量不大但每个环节都能替换正好当二次开发的底座。2. 工程骨架与数据管线从原始标注图到可迭代的 Dataset这个工程不大我把目录压到能一眼看完的程度。核心思路是数据、模型、训练、推理四块解耦任何一块想换都不牵动其他部分。多类别分割翻车最多的地方不在网络而在数据管线——标签错一位、插值方式选错训练曲线能给你演一出loss 下降但 mIoU 不动的玄学。2.1 目录约定与标签编码先约定目录。图像和掩膜分开放文件名主干保持一致靠后缀区分。这是最常见也最不容易出错的布局seg_project/ ├── data/ │ ├── images/ # 0001.jpg 0002.jpg ... │ └── masks/ # 0001.png 0002.png ...单通道像素值类别 id ├── splits/ │ ├── train.txt # 每行一个文件名主干 │ └── val.txt ├── src/ │ ├── dataset.py │ ├── unet.py │ ├── losses.py │ ├── metrics.py │ └── train.py └── predict.py掩膜必须是单通道索引图不是 RGB 彩色图。很多人从标注工具导出的是彩色掩膜红绿蓝三色各代表一类直接读进来会得到三通道张量送进交叉熵就会报维度错误或者更隐蔽地能跑但全错。所以第一步永远是做一次类别重映射把彩色或跳号的 id 压成 0 到 N-1 的连续整数并输出一份统计表看看每类像素占比。import numpy as np # 原始标注里的像素值可能是 0/1/2/5/9 这种跳号也可能是 (128,0,0) 这类彩色 RAW_TO_ID {0: 0, 1: 1, 2: 2, 5: 3, 9: 4} IGNORE 255 def remap_mask(mask: np.ndarray) - np.ndarray: 把原始掩膜重映射成连续类别 id未登记的值一律记为忽略标签。 out np.full(mask.shape, IGNORE, dtypenp.uint8) for raw, cid in RAW_TO_ID.items(): out[mask raw] cid return out这段的逻辑是把标注值和训练类别彻底解耦。RAW_TO_ID是唯一事实来源改类别只改这张表网络输出通道数跟着它的长度走。IGNORE255是给交叉熵用的边界不确定、标注员拿不准的区域统一刷成 255训练时不计梯度。参数上要注意num_classes必须等于len(RAW_TO_ID)而不是原始标注里的最大值加一否则混淆矩阵会越界或者在忽略标签那一格上算出莫名其妙的 IoU。2.2 Dataset 与同步增强Dataset 的职责只有三件读图、读掩膜、做同步增强。同步两个字是重点——图像可以双线性插值、可以色彩抖动掩膜只能最近邻、只能跟着做几何变换。我一般直接上 albumentations它对 image 和 mask 的插值方式是分开配置的省得自己写一套容易漏。import cv2 import torch from torch.utils.data import Dataset import albumentations as A from albumentations.pytorch import ToTensorV2 def build_transform(size512, trainTrue): if train: return A.Compose([ A.RandomResizedCrop(size, size, scale(0.5, 1.0)), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.2), A.RandomRotate90(p0.5), A.RandomBrightnessContrast(p0.3), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) # 验证集只做 resize 归一化保持可复现 return A.Compose([ A.Resize(size, size), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) class SegDataset(Dataset): def __init__(self, root, split_file, transform): self.root root self.transform transform with open(split_file) as f: self.names [ln.strip() for ln in f if ln.strip()] def __len__(self): return len(self.names) def __getitem__(self, idx): name self.names[idx] img cv2.imread(f{self.root}/images/{name}.jpg) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # cv2 是 BGR必须转 msk cv2.imread(f{self.root}/masks/{name}.png, cv2.IMREAD_GRAYSCALE) # 强制单通道 msk remap_mask(msk) aug self.transform(imageimg, maskmsk) img, msk aug[image], aug[mask].long() # ToTensorV2 后 mask 是 HxW 的 long交叉熵需要它保持这个形状 return img, msk逐段说清楚。cv2.imread默认按 BGR 读不转 RGB 的话用 ImageNet 预训练骨干做初始化时颜色通道就是错的这种错误不会报异常只会让收敛慢一截。掩膜用IMREAD_GRAYSCALE强制读成单通道避免彩色标注图进来变成 HxWx3。aug[mask].long()这一步是必须的交叉熵的 target 要求 int64uint8 会直接抛类型错误。RandomResizedCrop的scale(0.5, 1.0)是我推荐的起点。多类别数据集的痛点是稀有类裁剪太狠的时候某些小目标类可能整张图一个像素都没有等效于给了一个空样本。所以训练前建议跑一遍统计把每类像素占比和出现频率打出来scale 下限按最小目标尺寸调不要照抄二分类分割的配置。验证集用固定Resize而不是中心裁剪是因为 mIoU 需要和原图尺度对应才好对比裁剪会让边界区域被切掉指标虚高。2.3 DataLoader 的 batch 与显存配比DataLoader 本身没什么花样但参数配错很常见。num_workers在 Linux 上一般给 4 到 8Windows 上如果 Dataset 里有 lambda 或者本地类会卡在 worker 启动上那就退回 0 先用单进程跑通。pin_memoryTrue配合non_blockingTrue的 H2D 拷贝在混合精度训练下能吃到一点收益。真正影响能不能开跑的是 batch 和分辨率。512×512、6 类、base channel 32 的 UNet单卡 12G 显存大概能放 8 到 12 张具体取决于有没有开 AMP。我一般的顺序是先用 batch2、分辨率 256 跑 20 个 step 确认链路通再往上加。如果一上来就把 batch 开到 32、分辨率 1024报 OOM 之后你会去怀疑模型写错了其实是配比问题纯属自己给自己找麻烦。数据管线到这里就算通了。判断标准很简单随便取一个 batch把图像反归一化、掩膜按调色板上色并排画出来肉眼确认掩膜和原图对得上。这步花十分钟能省掉后面几小时怀疑人生的时间。3. UNet 网络拆解编码器、解码器与跳跃连接网络部分是这份代码里最标准的一块但标准不代表可以乱改。多类别分割的通道配置、上采样方式和输出头设计直接决定训练稳定性和边界质量。下面按模块拆。3.1 双卷积块与通道配置UNet 的基本单元是两次 3×3 卷积加 BN 加 ReLU。biasFalse是因为后面紧跟 BN偏置会被减掉纯属浪费参数。每一级通道翻倍最底层到 512 或 1024 就看显存。import torch import torch.nn as nn class DoubleConv(nn.Module): (Conv3x3 - BN - ReLU) x 2 def __init__(self, in_ch, out_ch): super().__init__() self.block nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), ) def forward(self, x): return self.block(x)padding1保证特征图尺寸不变这样跳跃连接两侧的 H、W 才对得上可以直接cat。如果你把卷积换成padding0或者用了 stride2 的卷积下采样拼接前就得手动补齐这是改结构时最容易忘的一步。base32是我对自建数据集推荐的默认值第一层 32 通道逐级 64、128、256、512参数量在千万级以下单卡能训收敛也快。数据集有几千张以上、目标纹理复杂时再把 base 提到 64。3.2 下采样与上采样的实现选择下采样用MaxPool2d(2)上采样有两条路ConvTranspose2d和双线性插值 3×3 卷积。前者可学习但容易在边缘产生棋盘格伪影后者更稳、参数量更少我在多类别任务上基本默认用后者。class Up(nn.Module): def __init__(self, in_ch, skip_ch, out_ch, bilinearTrue): super().__init__() if bilinear: self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) conv_in in_ch skip_ch else: self.up nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride2) conv_in in_ch // 2 skip_ch self.conv DoubleConv(conv_in, out_ch) def forward(self, x, skip): x self.up(x) # 输入尺寸不是 32 的整数倍时上采样后可能与 skip 差 1 个像素 if x.shape[-2:] ! skip.shape[-2:]: x nn.functional.interpolate( x, sizeskip.shape[-2:], modebilinear, align_cornersTrue) x torch.cat([skip, x], dim1) # 通道维拼接 return self.conv(x)这里的if分支是血泪经验。只要输入分辨率不能被 16 或 32 整除五次下采样再上采样就会出现 H 相差 1 像素torch.cat直接报尺寸不匹配。写死A.Resize(512, 512)之后这个问题不会出现但一旦推理时喂原图就会炸所以在Up里加这层兜底比在推理脚本里到处补丁要好。拼接顺序是[skip, x]通道顺序本身不影响结果但训练和推理必须一致否则加载权重后精度会莫名其妙掉一截。3.3 多类别输出头与初始化输出头就是一个 1×1 卷积把 64 通道压到num_classes。注意这里不要加 sigmoid 或 softmaxPyTorch 的CrossEntropyLoss内部已经做了 log_softmax重复做一次等价于把梯度压扁训练会变得极慢。class UNet(nn.Module): def __init__(self, in_ch3, num_classes6, base32, bilinearTrue): super().__init__() c [base, base * 2, base * 4, base * 8, base * 16] self.inc DoubleConv(in_ch, c[0]) self.down1 nn.MaxPool2d(2) self.conv1 DoubleConv(c[0], c[1]) self.down2 nn.MaxPool2d(2) self.conv2 DoubleConv(c[1], c[2]) self.down3 nn.MaxPool2d(2) self.conv3 DoubleConv(c[2], c[3]) self.down4 nn.MaxPool2d(2) self.conv4 DoubleConv(c[3], c[4]) self.up1 Up(c[4], c[3], c[3], bilinear) self.up2 Up(c[3], c[2], c[2], bilinear) self.up3 Up(c[2], c[1], c[1], bilinear) self.up4 Up(c[1], c[0], c[0], bilinear) self.head nn.Conv2d(c[0], num_classes, 1) self._init_weights() def _init_weights(self): for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, modefan_out, nonlinearityrelu) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.BatchNorm2d): nn.init.ones_(m.weight) nn.init.zeros_(m.bias) def forward(self, x): x1 self.inc(x) x2 self.conv1(self.down1(x1)) x3 self.conv2(self.down2(x2)) x4 self.conv3(self.down3(x3)) x5 self.conv4(self.down4(x4)) x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.head(x) # 输出 B x C x H x W未归一化 logits前向路径就是编码四次下采样、解码四次上采样每次拼接对应层级的特征。Up的第一个参数是来自更深层的通道数skip_ch是同级编码器的通道数输出通道统一回到 skip 的宽度这样后面能一层层对称地接下去。初始化用 kaiming 配合 ReLUBN 权重置 1、偏置置 0是卷积网络的常规起手式。别小看这一步去掉自定义初始化、全用 PyTorch 默认值深层网络前几个 epoch 的 loss 会明显更抖。输出是 logits推理时自己argmax(dim1)就行。num_classes一定要和RAW_TO_ID的长度严格对齐我见过把背景类漏掉、输出通道写成 5 但标签里有 0 到 5 共 6 个值的情况CrossEntropyLoss会直接抛 Target out of bounds这类报错反而是好事最怕的是不报错但静默错位。4. 训练与评估闭环损失、mIoU 与 checkpoint模型搭好只是开始多类别分割的效果差异八成出在损失设计和评估口径上。这一章把训练循环、指标计算和推理可视化串成一条完整链路。4.1 多类别损失与忽略标签主损失用CrossEntropyLossignore_index255把边界和不确定区域排除掉。类别极度不平衡时比如背景占 90%再加一路 Dice 或者给类别权重两者选其一即可不建议全堆上去。import torch.nn as nn import torch.nn.functional as F class ComboLoss(nn.Module): def __init__(self, num_classes, ignore_index255, dice_weight0.5): super().__init__() self.ce nn.CrossEntropyLoss(ignore_indexignore_index) self.num_classes num_classes self.ignore_index ignore_index self.dice_weight dice_weight def forward(self, logits, target): ce_loss self.ce(logits, target) # Dice 部分要先把忽略标签从 target 里剔掉否则 255 会被当成类别 valid target ! self.ignore_index logits_v logits.permute(0, 2, 3, 1)[valid] # N x C target_v target[valid] # N probs F.softmax(logits_v, dim1) # N x C onehot F.one_hot(target_v, self.num_classes).float() inter (probs * onehot).sum(0) union probs.sum(0) onehot.sum(0) dice ((2 * inter 1e-6) / (union 1e-6)).mean() return ce_loss self.dice_weight * (1 - dice)交叉熵负责逐像素的梯度Dice 负责把稀有类的整体重叠度拉上来dice_weight0.5是个安全的中间值小目标类仍然学不动就往上加到 1.0但别超过 1否则训练前期会很不稳定。注意permute的位置logits 原本是N x C x H x W要在空间维上做 softmax必须先换成N x H x W x C再展平成N x C顺序搞反得到的 Dice 是错的而且不报错。那个1e-6是防除零别省。如果某一类样本极少比如几千张图里只有几十张带它光调损失不够回到 Dataset 层面用WeightedRandomSampler按含该类样本的比例给采样权重效果比在损失里硬压权重更直接。4.2 训练循环与学习率调度训练循环要包含的几件事AMP、梯度裁剪、按 epoch 记录指标、只在 mIoU 提升时存 checkpoint。少一样都能跑但少了后面调起来会很难受。import torch from torch.amp import autocast, GradScaler def train_one_epoch(model, loader, optimizer, criterion, scaler, device): model.train() total 0.0 for img, msk in loader: img, msk img.to(device, non_blockingTrue), msk.to(device) optimizer.zero_grad(set_to_noneTrue) with autocast(device_typecuda, dtypetorch.float16): logits model(img) loss criterion(logits, msk) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) scaler.step(optimizer) scaler.update() total loss.item() * img.size(0) return total / len(loader.dataset)set_to_noneTrue比置零省一点显存和带宽现在基本是默认写法。AMP 的autocast只包前向反向交给GradScaler顺序别弄反。梯度裁剪max_norm1.0在多类别任务上能压住前期偶发的梯度尖峰尤其是用了 Dice 的时候。学习率调度我一般用余弦退火初始3e-4配合 AdamWweight_decay1e-4optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_maxargs.epochs, eta_min1e-6)每个 epoch 结束调一次scheduler.step()注意 CosineAnnealingLR 是按调用次数推进的别在 batch 里调否则学习率掉得飞快模型还没学到东西就停了。checkpoint 我存三样东西model.state_dict()、optimizer.state_dict()、当前最优 mIoU。只存权重不带优化器状态中断后 resume 会丢掉动量续训前几个 epoch 的指标会明显掉一截。这个坑我踩过不止一次。4.3 mIoU 与混淆矩阵评估口径必须固定。最省事也最不容易算错的方式是累一个num_classes x num_classes的混淆矩阵最后统一算 IoU。import numpy as np import torch torch.no_grad() def evaluate(model, loader, num_classes, device, ignore_index255): model.eval() conf np.zeros((num_classes, num_classes), dtypenp.int64) for img, msk in loader: img img.to(device) pred model(img).argmax(dim1).cpu().numpy() gt msk.numpy() valid gt ! ignore_index # 把 (gt, pred) 编码成单个整数再一次性统计 idx gt[valid].astype(np.int64) * num_classes pred[valid] conf np.bincount(idx, minlengthnum_classes ** 2 ).reshape(num_classes, num_classes) inter np.diag(conf) union conf.sum(1) conf.sum(0) - inter iou inter / np.maximum(union, 1) miou np.nanmean(iou) return miou, iouconf[i][j]表示真实类别 i 被预测成 j 的次数。行和是真实分布列和是预测分布inter是对角线union是并集。np.maximum(union, 1)防除零np.nanmean让验证集里根本没出现的类别不参与平均——这一步很重要如果直接用mean某个类别在整个验证集里出现 0 次IoU 会算成 0 或 nan把整体指标拉得没法看。多类别任务里建议把每类的 IoU 单独 print 出来只看 mIoU 会掩盖问题背景 0.98、四个类 0.6、一个小类 0.05平均下来 0.6 出头看着还行实际上那个小类等于没学会。4.4 推理与可视化推理脚本要跟验证集用同一套归一化参数这是最常见的翻车点——训练用 ImageNet 均值方差推理忘了写预测图出来一片噪声。import cv2 import numpy as np PALETTE np.array([ [0, 0, 0], [255, 0, 0], [0, 255, 0], [0, 0, 255], [255, 255, 0], [255, 0, 255], ], dtypenp.uint8) torch.no_grad() def predict(model, image_bgr, size512, devicecuda): model.eval() img cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) img cv2.resize(img, (size, size)) x img.astype(np.float32) / 255.0 x (x - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225]) x torch.from_numpy(x).permute(2, 0, 1).unsqueeze(0).float().to(device) pred model(x).argmax(dim1)[0].cpu().numpy().astype(np.uint8) color PALETTE[pred] color cv2.resize(color, (image_bgr.shape[1], image_bgr.shape[0]), interpolationcv2.INTER_NEAREST) # 上色图最近邻放大 return colorPALETTE的长度必须等于num_classes少一行就会IndexError。最后放大回原图尺寸时用INTER_NEAREST双线性会把类别边界糊成渐变看起来像模型边界很软其实是可视化插值的锅。把原图、真值上色图、预测上色图三张并排存下来是每次跑完实验的第一道人工检查比盯数字直观得多。5. 避坑与排查标签错位、不收敛与显存溢出这一章是前面几步的实际翻车记录按现象 → 原因 → 解决写遇到对得上的直接照方抓药。第一条loss 正常下降mIoU 始终在 0.1 附近。现象是训练集 loss 从 1.8 降到 0.3但验证 mIoU 死活上不去预测图几乎全是背景色。原因基本锁定在标签错位或类别塌缩要么掩膜 resize 用了双线性插值类别 id 被插出 1.5、2.7 这种非整数值取整后大面积错位要么背景占比太高模型学会了全预测背景这个局部最优。解决分两步先把所有Resize的插值改成INTER_NEAREST再统计每类像素占比背景超过 90% 就上 ComboLoss 的 Dice 部分或者在 Dataset 里按类别做加权采样。改完重新从零训别在旧权重上续。第二条报错 Target out of bounds或者混淆矩阵索引越界。现象是CrossEntropyLoss抛Target x is out of bounds或者评估时np.bincount报错。原因是掩膜里的类别 id 不连续比如实际值是 0/1/2/5/9而num_classes按 5 来设5 和 9 直接越界另一种是原始掩膜带 255 以外的杂值比如边缘抗锯齿产生的 254。解决办法是先跑一遍全量掩膜的像素值统计np.unique打出来看清楚再用RAW_TO_ID白名单映射未登记值一律刷成 255别用mask[mask N] 255这种粗暴写法它会把合法类别也一起干掉。第三条batch 调到 2 以下就报 BN 相关的错或者指标剧烈抖动。现象是小 batch 训练时 loss 突然变成 nan或者验证指标上下跳十几个点。原因是BatchNorm2d在 batch 很小时统计量极不稳定batch1 时方差几乎为 0。解决有三个方向换GroupNorm(num_groups8, num_channelsch)对 batch 不敏感或者冻结 BN 的 running statsmodel.eval()后只训参数适合微调场景再或者用梯度累积把等效 batch 撑上去。我一般优先换 GroupNorm改起来就一行代价是收敛稍慢一点。第四条验证阶段显存越跑越涨最后 OOM。现象是训练正常跑evaluate几十个 batch 后显存爆掉。原因十有八九是忘了torch.no_grad()或者忘了model.eval()前者会把整个计算图挂在显存里不释放后者会让 BN 继续更新 running stats验证指标一轮一个样。检查方式是看验证函数有没有装饰器、循环里有没有torch.no_grad()包住。另外conf用 numpy 累加、预测结果及时.cpu()不要留在 GPU 上。第五条训练集 Dice 很高换一批同分布的新图就崩。现象是验证集指标 0.85实际拿新数据一跑边界全是锯齿和误检。原因是增强策略在验证集上泄露了或者验证集和训练集来自同一批采集、分布过于接近。检查一下build_transform(trainFalse)是不是误用了训练增强再确认 train/val 的文件名主干没有交叉。我一般还会额外留一个完全独立的小测试集哪怕只有 20 张专门做这种换数据就崩的兜底检查。6. 二次开发换骨干、加模块与导出部署的验证路径代码能跑通之后真正花时间的是改造和验证。这一章讲三个最常改的方向以及我自己固定会走的一道体检流程。换骨干是最划算的改造。原生 UNet 的编码器是从零训的数据量上千张时换成带预训练权重的 ResNet 或 EfficientNet 编码器mIoU 通常能涨几个点。做法是把编码器的四个 stage 输出接出来当 skip通道分别是 64/128/256/512Up的skip_ch跟着改就行。要注意预训练模型的第一层是 stride2 的 7×7 卷积直接替换self.inc会让输出分辨率比原来小一半后面所有拼接尺寸都得重算稳妥的做法是保留第一层 stride2 的 stem在它的输出上开始接conv1并在最后的上采样多补一次。加模块要克制。CBAM、SE 这类注意力放在跳跃连接拼接之前作用是对齐编码器特征代码量很小class SEBlock(nn.Module): def __init__(self, ch, r8): super().__init__() self.fc nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(ch, ch // r, 1), nn.ReLU(inplaceTrue), nn.Conv2d(ch // r, ch, 1), nn.Sigmoid(), ) def forward(self, x): return x * self.fc(x) # 在 Up.forward 里skip self.se(skip) 之后再 cat这类模块一次只加一处加完必须做消融对比否则你不知道涨的点是来自模块还是随机种子。我见过一口气加注意力、加空洞卷积、换损失最后指标涨了 2 个点但根本说不清是谁的功劳后续也没法复用经验。导出部署前先做一致性验证。ONNX 导出用opset_version11dynamic_axes把 batch 和 H/W 都设成动态然后用 onnxruntime 跑同一张图和 PyTorch 的 logits 逐元素比对最大绝对误差控制在 1e-3 以内再往下走torch.onnx.export( model, dummy, unet.onnx, input_names[input], output_names[logits], opset_version11, dynamic_axes{input: {0: b, 2: h, 3: w}, logits: {0: b, 2: h, 3: w}}, )最容易忽略的是Up里那个尺寸兜底分支导出时如果走的是nn.functional.interpolate动态 size 分支某些推理后端对动态插值支持不好会退化成固定尺寸。所以导出前建议先固定分辨率跑一遍确认静态图没问题再开动态轴。还有一个习惯我每次接新数据集都会先做十张图过拟合取 10 张图关掉增强学习率调到 1e-3训 200 个 step看 loss 能不能压到接近 0、mIoU 能不能到 0.95 以上。压不下去就说明数据管线或标签有问题这时候再去调网络结构纯属浪费时间。从那以后我每次动 Dataset 或者改标签映射都强制先走一遍这个十张图过拟合比看任何日志都直接。希望帮到你。本文还有配套的精品资源点击获取