简介：面向需要在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, dtype=np.uint8) for raw, cid in RAW_TO_ID.items(): out[mask == raw] = cid return out

这段的逻辑是把"标注值"和"训练类别"彻底解耦。RAW_TO_ID是唯一事实来源，改类别只改这张表，网络输出通道数跟着它的长度走。IGNORE=255是给交叉熵用的：边界不确定、标注员拿不准的区域统一刷成 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(size=512, train=True): if train: return A.Compose([ A.RandomResizedCrop(size, size, scale=(0.5, 1.0)), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.2), A.RandomRotate90(p=0.5), A.RandomBrightnessContrast(p=0.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(image=img, mask=msk) 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 要求 int64，uint8 会直接抛类型错误。

RandomResizedCrop的scale=(0.5, 1.0)是我推荐的起点。多类别数据集的痛点是稀有类，裁剪太狠的时候，某些小目标类可能整张图一个像素都没有，等效于给了一个空样本。所以训练前建议跑一遍统计，把每类像素占比和出现频率打出来，scale 下限按最小目标尺寸调，不要照抄二分类分割的配置。

验证集用固定Resize而不是中心裁剪，是因为 mIoU 需要和原图尺度对应才好对比；裁剪会让边界区域被切掉，指标虚高。

2.3 DataLoader 的 batch 与显存配比

DataLoader 本身没什么花样，但参数配错很常见。num_workers在 Linux 上一般给 4 到 8，Windows 上如果 Dataset 里有 lambda 或者本地类，会卡在 worker 启动上，那就退回 0 先用单进程跑通。pin_memory=True配合non_blocking=True的 H2D 拷贝，在混合精度训练下能吃到一点收益。

真正影响能不能开跑的是 batch 和分辨率。512×512、6 类、base channel 32 的 UNet，单卡 12G 显存大概能放 8 到 12 张，具体取决于有没有开 AMP。我一般的顺序是：先用 batch=2、分辨率 256 跑 20 个 step 确认链路通，再往上加。如果一上来就把 batch 开到 32、分辨率 1024，报 OOM 之后你会去怀疑模型写错了，其实是配比问题，纯属自己给自己找麻烦。

数据管线到这里就算通了。判断标准很简单：随便取一个 batch，把图像反归一化、掩膜按调色板上色，并排画出来，肉眼确认掩膜和原图对得上。这步花十分钟，能省掉后面几小时怀疑人生的时间。

3. UNet 网络拆解：编码器、解码器与跳跃连接

网络部分是这份代码里最"标准"的一块，但标准不代表可以乱改。多类别分割的通道配置、上采样方式和输出头设计，直接决定训练稳定性和边界质量。下面按模块拆。

3.1 双卷积块与通道配置

UNet 的基本单元是两次 3×3 卷积加 BN 加 ReLU。bias=False是因为后面紧跟 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, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), ) def forward(self, x): return self.block(x)

padding=1保证特征图尺寸不变，这样跳跃连接两侧的 H、W 才对得上，可以直接cat。如果你把卷积换成padding=0或者用了 stride=2 的卷积下采样，拼接前就得手动补齐，这是改结构时最容易忘的一步。base=32是我对自建数据集推荐的默认值，第一层 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, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True) conv_in = in_ch + skip_ch else: self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, 2, stride=2) 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, size=skip.shape[-2:], mode="bilinear", align_corners=True) x = torch.cat([skip, x], dim=1) # 通道维拼接 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 或 softmax，PyTorch 的CrossEntropyLoss内部已经做了 log_softmax，重复做一次等价于把梯度压扁，训练会变得极慢。

class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=6, base=32, bilinear=True): 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, mode="fan_out", nonlinearity="relu") 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 配合 ReLU，BN 权重置 1、偏置置 0，是卷积网络的常规起手式。别小看这一步，去掉自定义初始化、全用 PyTorch 默认值，深层网络前几个 epoch 的 loss 会明显更抖。

输出是 logits，推理时自己argmax(dim=1)就行。num_classes一定要和RAW_TO_ID的长度严格对齐，我见过把背景类漏掉、输出通道写成 5 但标签里有 0 到 5 共 6 个值的情况，CrossEntropyLoss会直接抛 "Target out of bounds"，这类报错反而是好事，最怕的是不报错但静默错位。

4. 训练与评估闭环：损失、mIoU 与 checkpoint

模型搭好只是开始，多类别分割的效果差异，八成出在损失设计和评估口径上。这一章把训练循环、指标计算和推理可视化串成一条完整链路。

4.1 多类别损失与忽略标签

主损失用CrossEntropyLoss，ignore_index=255把边界和不确定区域排除掉。类别极度不平衡时（比如背景占 90%），再加一路 Dice 或者给类别权重，两者选其一即可，不建议全堆上去。

import torch.nn as nn import torch.nn.functional as F class ComboLoss(nn.Module): def __init__(self, num_classes, ignore_index=255, dice_weight=0.5): super().__init__() self.ce = nn.CrossEntropyLoss(ignore_index=ignore_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, dim=1) # 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_weight=0.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_blocking=True), msk.to(device) optimizer.zero_grad(set_to_none=True) with autocast(device_type="cuda", dtype=torch.float16): logits = model(img) loss = criterion(logits, msk) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) scaler.step(optimizer) scaler.update() total += loss.item() * img.size(0) return total / len(loader.dataset)

set_to_none=True比置零省一点显存和带宽，现在基本是默认写法。AMP 的autocast只包前向，反向交给GradScaler，顺序别弄反。梯度裁剪max_norm=1.0在多类别任务上能压住前期偶发的梯度尖峰，尤其是用了 Dice 的时候。

学习率调度我一般用余弦退火，初始3e-4，配合 AdamW，weight_decay=1e-4：

optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs, eta_min=1e-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_index=255): model.eval() conf = np.zeros((num_classes, num_classes), dtype=np.int64) for img, msk in loader: img = img.to(device) pred = model(img).argmax(dim=1).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, minlength=num_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, iou

conf[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], ], dtype=np.uint8) @torch.no_grad() def predict(model, image_bgr, size=512, device="cuda"): 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(dim=1)[0].cpu().numpy().astype(np.uint8) color = PALETTE[pred] color = cv2.resize(color, (image_bgr.shape[1], image_bgr.shape[0]), interpolation=cv2.INTER_NEAREST) # 上色图最近邻放大 return color

PALETTE的长度必须等于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 很小时统计量极不稳定，batch=1 时方差几乎为 0。解决有三个方向：换GroupNorm(num_groups=8, num_channels=ch)，对 batch 不敏感；或者冻结 BN 的 running stats（model.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(train=False)是不是误用了训练增强，再确认 train/val 的文件名主干没有交叉。我一般还会额外留一个完全独立的小测试集，哪怕只有 20 张，专门做这种"换数据就崩"的兜底检查。

6. 二次开发：换骨干、加模块与导出部署的验证路径

代码能跑通之后，真正花时间的是改造和验证。这一章讲三个最常改的方向，以及我自己固定会走的一道体检流程。

换骨干是最划算的改造。原生 UNet 的编码器是从零训的，数据量上千张时换成带预训练权重的 ResNet 或 EfficientNet 编码器，mIoU 通常能涨几个点。做法是把编码器的四个 stage 输出接出来当 skip，通道分别是 64/128/256/512，Up的skip_ch跟着改就行。要注意预训练模型的第一层是 stride=2 的 7×7 卷积，直接替换self.inc会让输出分辨率比原来小一半，后面所有拼接尺寸都得重算，稳妥的做法是保留第一层 stride=2 的 stem，在它的输出上开始接conv1，并在最后的上采样多补一次。

加模块要克制。CBAM、SE 这类注意力放在跳跃连接拼接之前，作用是对齐编码器特征，代码量很小：

class SEBlock(nn.Module): def __init__(self, ch, r=8): super().__init__() self.fc = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(ch, ch // r, 1), nn.ReLU(inplace=True), 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_version=11，dynamic_axes把 batch 和 H/W 都设成动态，然后用 onnxruntime 跑同一张图，和 PyTorch 的 logits 逐元素比对，最大绝对误差控制在 1e-3 以内再往下走：

torch.onnx.export( model, dummy, "unet.onnx", input_names=["input"], output_names=["logits"], opset_version=11, 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 或者改标签映射，都强制先走一遍这个十张图过拟合，比看任何日志都直接。希望帮到你。

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