企业官网建设流程全解析

简介：这是一份面向图像分割初学者与深度学习实践者的 Python+U-Net 实践资料包，聚焦医学影像、卫星遥感等场景中从图像到像素级掩模的分割任务。压缩包共 21 个文件，约 5.6MB，以 jpg、png 示例图和分割前后对比图为主要素材，辅以 py 预测脚本、md 说明文档、txt 依赖清单及 gif 动图，覆盖图像预处理、模型预测与结果可视化等环节。目前已有 10164 人学习下载，具备较好的参考热度。核心预测脚本展示图像切片与平滑融合思路，配合说明文档、依赖清单和示例图，读者可据此理解 U-Net 收缩—扩展结构、跳跃连接及 Dice/交叉熵损失在分割任务中的作用，也能参考目录组织方式搭建自己的训练与验证流程，适合作为课程设计、项目实践和算法入门的对照材料。

1. 工业质检线上的那点事：为什么我盯着这个 UNet 图像分割项目看了半小时

拿一张金属件表面照片，人眼能圈出划痕，但产线上的程序要的是"每个像素标成缺陷还是背景"，这就落到 Python 图像分割的主场。UNet 结构简单、对小样本友好，做像素级任务时最省心：编码器一路下采样抽语义，解码器靠 skip connection 把浅层边缘信息接回来，薄划痕不容易丢。这份资源给的是 Python + UNet 的整套工程代码，数据组织、增强、训练、推理评估都在，不是只丢一个网络定义的半成品。它适合两类人：刚接触语义分割、需要一份能跑通全流程的骨架照着改的；手里已有业务图数据、想直接替换数据集就出结果的。下面按我拆项目的顺序，从目录一路讲到改哪里。

2. 项目骨架与技术选型：UNet 的编码器-解码器怎么落到代码里

先看骨架再谈模型，是我拆任何分割项目的固定顺序。原因很直白：图像分割的坑八成不在网络结构上，而在数据通路、掩码格式、归一化一致性这些地方，目录结构恰好把这些东西暴露得最清楚。这套工程是典型的脚本式组织，不套框架、不上多余封装，train.py和predict.py两个入口直接对应"训练"和"出图"两个动作，中间靠models/、datasets/、utils/三个目录分责，读起来没什么认知负担。

2.1 目录结构与文件职责

目录长这样，我按职责列一遍，方便你对号入座：

unet-segmentation/ ├── data/ │ ├── train/ │ │ ├── images/ # 训练原图，jpg/png 均可 │ │ └── masks/ # 与 images 同名，单通道 0/255 掩码 │ ├── val/ │ │ ├── images/ │ │ └── masks/ │ └── test/ │ └── images/ # 只放待推理图，不放 mask ├── models/ │ ├── __init__.py │ └── unet.py # 网络定义，DoubleConv/Down/Up 三段 ├── datasets/ │ ├── __init__.py │ └── segmentation_dataset.py ├── utils/ │ ├── losses.py # BCE + Dice │ ├── metrics.py # Dice / IoU │ └── visualize.py # 掩码叠加原图 ├── configs/ │ └── train.yaml # 路径、超参、图像尺寸集中配置 ├── train.py ├── predict.py └── requirements.txt

这个布局的价值在于"换数据集只动data/和configs/train.yaml"。我见过不少同学把路径硬编码在train.py里，换个数据集要全局搜索替换，改漏一处就报 FileNotFoundError，翻车翻得很冤枉。

依赖方面，常见做法是 PyTorch + torchvision + albumentations + opencv-python + PyYAML + tqdm，写进requirements.txt即可：

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 pip install albumentations opencv-python PyYAML tqdm

CUDA 版本按你机器实际的驱动选，torch.cuda.is_available()返回 False 时先别急着调模型，那是环境问题不是代码问题。

2.2 编码器与解码器：UNet 网络逐层拆解

网络部分只有三个组件，理解了这三个，UNet 就没有黑匣子了。第一个是DoubleConv，两次 3x3 卷积加 BN 加 ReLU，是每一层的基本砖块：

import torch import torch.nn as nn class DoubleConv(nn.Module): """(conv3x3 -> BN -> ReLU) x2，UNet 的基础卷积块""" def __init__(self, in_ch, out_ch, mid_ch=None): super().__init__() mid_ch = mid_ch or out_ch self.block = nn.Sequential( nn.Conv2d(in_ch, mid_ch, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(mid_ch), nn.ReLU(inplace=True), nn.Conv2d(mid_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保证卷积后尺寸不变，bias=False是因为紧跟 BN，偏置会被归一化抵消，省这点参数对显存没坏处。mid_ch参数是给后面换 backbone 时留的口子，标准 UNet 用不上，保持默认即可。

第二个是下采样和上采样两个模块，下采样用最大池化丢一半分辨率，上采样用转置卷积恢复，然后把编码器同名层的特征在通道维拼接：

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) # 拼接后通道 = in_ch//2 + 对应编码器层通道 self.conv = DoubleConv(in_ch, out_ch) def forward(self, x, skip): x = self.up(x) # 尺寸对齐，防止输入非 16 倍数导致的 1px 误差 if x.shape[-2:] != skip.shape[-2:]: x = nn.functional.interpolate(x, size=skip.shape[-2:], mode="nearest") return self.conv(torch.cat([skip, x], dim=1))

torch.cat那行的通道数要算对：以标准 64/128/256/512/1024 通道配置为例，最深层上采样后是 512，拼接编码器倒数第二层的 512，合计 1024，正好是DoubleConv(1024, 512)的输入。这里数错是新手最常见的报错来源，RuntimeError: given groups=1, weight of size [512, 1024, 3, 3]基本都是这个原因。

完整的 UNet 组装起来就是编码 4 次下采样、解码 4 次上采样，最后 1x1 卷积把通道压到类别数：

class UNet(nn.Module): def __init__(self, in_ch=1, num_classes=1, base=64): super().__init__() c = [base, base * 2, base * 4, base * 8, base * 16] self.inc = DoubleConv(in_ch, c[0]) self.down1 = Down(c[0], c[1]) self.down2 = Down(c[1], c[2]) self.down3 = Down(c[2], c[3]) self.down4 = Down(c[3], c[4]) self.up1 = Up(c[4], c[3]) self.up2 = Up(c[3], c[2]) self.up3 = Up(c[2], c[1]) self.up4 = Up(c[1], c[0]) self.out = nn.Conv2d(c[0], num_classes, kernel_size=1) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = 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.out(x) # 输出 logits，不加 sigmoid

注意最后一行返回的是 logits 而不是概率，sigmoid 放到损失函数和推理阈值里做。这个约定很重要，nn.BCEWithLogitsLoss内部已经把 sigmoid 和 BCE 合在一起算，数值上更稳；如果你在forward里手动加sigmoid，外面再用BCEWithLogitsLoss，就等于套了两层，训练会明显变慢甚至不收敛。

2.3 输入尺寸、通道数与参数量核算

标准 UNet 的参数量大约 3100 万，输入单通道 256x256 时，单卡 8G 显存跑 batch 8 是没问题的；如果改成 RGB 三通道，in_ch改成 3，DoubleConv第一层的参数量会多出约 1.7 万，可以忽略。真正影响显存的是分辨率和 batch，512x512 相比 256x256 特征图面积翻四倍，显存大概要 3 倍以上。

输入尺寸有个硬约束：必须能被 16 整除。四次MaxPool2d(2)意味着特征图会被连续除 16，254 这样的尺寸在第 4 次下采样时会出现 15.875 取整成 15，上采样再翻倍就回到 30 而不是 32，skip connection 拼接时直接报尺寸不匹配。我在Up.forward里加了interpolate兜底，但那是后悔药，正确做法是数据管线里就统一Resize到 256 或 512，别让网络去处理这种不确定性。

通道数的选择上，base=64是原论文配置；如果数据集只有几百张、类别单一，把base降到 32 参数量直接砍到四分之一，训练更快，精度掉得通常不多，这是我一般在样本量小于 1000 时先试的做法。反之样本上万、目标细碎，再往 96 或 128 加，但要注意显存。

3. 数据管线与训练主循环：从原始图到能收敛的 loss 曲线

网络定义清楚了，接下来是决定成败的部分。同样的 UNet 结构，数据管线写歪了，Dice 卡在 0.2 上不去；数据管线写对了，Dice 能到 0.85 以上。这一章把Dataset、增强、损失和训练循环串成一条线，中间每步都给出可抄的代码。

3.1 数据目录约定与掩码的预处理

目录约定是images/和masks/同名文件，训练图001.jpg对应掩码001.png。掩码必须是单通道灰度图，前景 255、背景 0，不能是 RGB 三通道的白色——cv2.imread读三通道图返回(H,W,3)，后面ToTensorV2会变成(3,H,W)，和输出的(1,H,W)对不上，loss 直接形状报错。读掩码时统一加cv2.IMREAD_GRAYSCALE就能规避：

img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) mask = (mask > 127).astype("float32") # 0/255 -> 0.0/1.0

阈值取 127 是为了兼容标注工具导出的 0/1、0/255 两种格式，存在 JPEG 压缩噪声时也能稳住。如果你的数据是多类别分割，这时候要改成np.uint8的类别索引，损失函数也换成CrossEntropyLoss，不能再用 BCE，这一点在二次开发时会重点讲。

增强管线用 albumentations，训练和验证分开写，验证集只做 Resize 和归一化，不能加任何随机变换，否则每次验证的指标都在抖，你根本分不清是模型在学还是增强在捣乱：

import albumentations as A from albumentations.pytorch import ToTensorV2 def build_transforms(image_size, is_train=True): if is_train: return A.Compose([ A.RandomResizedCrop(height=image_size, width=image_size, scale=(0.7, 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,), std=(0.229,)), # 灰度图单通道 ToTensorV2(), ]) return A.Compose([ A.Resize(height=image_size, width=image_size), A.Normalize(mean=(0.485,), std=(0.229,)), ToTensorV2(), ])

RandomResizedCrop的scale下限我一般设 0.7 而不是 0.6，因为细划痕这种目标本来就小，裁得太狠会把目标整个裁掉，模型学到的都是纯背景。Normalize的均值方差是 ImageNet 那套按三通道算的，单通道时取第一个通道的值即可，虽然不严格，但在这个场景下够用；要严谨就自己统计一遍训练集的均值和标准差，阈值别拍脑袋。

3.2 Dataset 与 DataLoader 的落地写法

Dataset的职责就三件：读文件、做增强、返回张量。写法上有个细节值得说：__getitem__里不要做任何 resize 之外的几何运算，所有增强交给 albumentations，因为它会同步处理 image 和 mask，你自己手写cv2.resize容易忘了对 mask 改用最近邻插值，导致边缘出现 128 这种中间值，二值化后掩码边界被吃掉一圈。

import os import cv2 import numpy as np import torch from torch.utils.data import Dataset class SegDataset(Dataset): def __init__(self, root, split="train", transform=None): self.img_dir = os.path.join(root, split, "images") self.mask_dir = os.path.join(root, split, "masks") self.names = sorted(os.listdir(self.img_dir)) self.transform = transform def __len__(self): return len(self.names) def __getitem__(self, idx): name = self.names[idx] img = cv2.imread(os.path.join(self.img_dir, name), cv2.IMREAD_GRAYSCALE) mask = cv2.imread(os.path.join(self.mask_dir, name), cv2.IMREAD_GRAYSCALE) if img is None or mask is None: raise FileNotFoundError(f"读图失败: {name}") mask = (mask > 127).astype("float32") augmented = self.transform(image=img, mask=mask) image = augmented["image"] # (1,H,W) float32 mask = augmented["mask"].float() # (H,W) float32 if mask.ndim == 2: mask = mask.unsqueeze(0) # -> (1,H,W)，和 logits 对齐 return image, mask

DataLoader 的num_workers在 Linux 下给 4 到 8 比较舒服，Windows 下必须包在if __name__ == "__main__":里，否则会疯狂重复启动进程，这是我在 Windows 上第一次跑时踩的坑，屏幕刷满了打印信息，一度以为是死循环。

3.3 损失函数与训练主循环

二值分割最常见的问题是前景占比极低，一张 512x512 图里缺陷可能只占 1%，纯 BCE 会让模型倾向于全预测背景，loss 降得很好看，Dice 却接近 0。解决办法是 BCE 和 Dice 加权组合：

import torch def dice_loss(logits, targets, eps=1e-6): probs = torch.sigmoid(logits) num = 2 * (probs * targets).sum(dim=(1, 2, 3)) den = probs.sum(dim=(1, 2, 3)) + targets.sum(dim=(1, 2, 3)) + eps return 1 - (num / den).mean() def combined_loss(logits, targets, bce_weight=0.5): bce = torch.nn.functional.binary_cross_entropy_with_logits(logits, targets) return bce_weight * bce + (1 - bce_weight) * dice_loss(logits, targets)

Dice 那项按 batch 内每个样本单独算再平均，比整批求和更能照顾小目标；eps防止某张图全背景时除零。bce_weight我一般从 0.5 起步，如果验证集 Dice 长期在 0.4 以下、预测结果大片空白，把它降到 0.3 让 Dice 项占主导，通常两周内的收敛速度会明显改善。

训练循环本身不复杂，关键是验证阶段的三件事不能漏：model.eval()、torch.no_grad()、指标在阈值化之后算。

for epoch in range(cfg["epochs"]): model.train() running = 0.0 for images, masks in train_loader: images = images.to(device, non_blocking=True) masks = masks.to(device, non_blocking=True) optimizer.zero_grad(set_to_none=True) logits = model(images) loss = combined_loss(logits, masks) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0) optimizer.step() running += loss.item() * images.size(0) train_loss = running / len(train_loader.dataset) model.eval() dice_total, n = 0.0, 0 with torch.no_grad(): for images, masks in val_loader: images = images.to(device) masks = masks.to(device) preds = (torch.sigmoid(model(images)) > 0.5).float() dice_total += dice_score(preds, masks) * images.size(0) n += images.size(0) val_dice = dice_total / max(n, 1) print(f"epoch {epoch} train_loss={train_loss:.4f} val_dice={val_dice:.4f}")

clip_grad_norm_那行是我吃过亏才加上的：Dice loss 在预测全零的早期梯度很陡，偶尔一次异常样本能把梯度炸到几百，训练直接发散，loss 变成 nan。加了梯度裁剪之后这种情况基本再见。学习率用1e-3起步配ReduceLROnPlateau，连续 5 个 epoch 验证 Dice 不涨就乘 0.5，比手动调省事。

4. 推理与指标验证：把 .pth 权重变成一张能看的掩码

训练完拿到best.pth，很多人卡在"怎么用"。推理脚本要解决三件事：预处理必须和训练完全一致、输出要还原到原图尺寸、指标要算得可信。尤其第一条，训练时用了Normalize(mean=0.485, std=0.229)，推理时忘了或者用了别的均值，模型表现会掉一大截，这是分割项目里最隐蔽的翻车点。

4.1 单图推理的完整流程

推理函数保持和Dataset一致的预处理，输出概率图后按 0.5 阈值二值化：

import cv2 import numpy as np import torch @torch.no_grad() def predict_image(model, img_path, device, image_size=256, threshold=0.5): model.eval() raw = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) h, w = raw.shape[:2] # 1. 预处理：Resize -> 归一化 -> 张量，顺序不能变 img = cv2.resize(raw, (image_size, image_size), interpolation=cv2.INTER_LINEAR) img = img.astype(np.float32) / 255.0 img = (img - 0.485) / 0.229 tensor = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(device) # 2. 前向 logits = model(tensor) prob = torch.sigmoid(logits)[0, 0].cpu().numpy() # 3. 还原到原图尺寸，用最近邻避免引入中间值 prob = cv2.resize(prob, (w, h), interpolation=cv2.INTER_NEAREST) mask = (prob > threshold).astype(np.uint8) * 255 return raw, prob, mask

参数里image_size必须和训练时保持一致，这是硬约束；threshold在样本不均衡时可以调，缺陷召回率不满意就降到 0.35，误检多了就抬到 0.6，这个值没有标准答案，按业务更怕漏检还是更怕误检来定。INTER_NEAREST那行是有意为之，概率图放大用双线性会让边界模糊，二值化后边缘多出毛刺。

4.2 大图滑窗推理

原图分辨率远大于 256 时，直接 resize 会丢掉小目标。工业图常见 2000x2000 以上，我的做法是滑窗切块，加 25% 重叠，最后对重叠区取平均：

@torch.no_grad() def sliding_window_predict(model, img_path, device, win=512, stride=384): model.eval() raw = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) h, w = raw.shape[:2] acc = np.zeros((h, w), dtype=np.float32) cnt = np.zeros((h, w), dtype=np.float32) for y in range(0, h, stride): for x in range(0, w, stride): y2, x2 = min(y + win, h), min(x + win, w) y1, x1 = max(0, y2 - win), max(0, x2 - win) # 贴边时回退窗口 patch = raw[y1:y2, x1:x2].astype(np.float32) / 255.0 patch = (patch - 0.485) / 0.229 t = torch.from_numpy(patch)[None, None].to(device) prob = torch.sigmoid(model(t))[0, 0].cpu().numpy() acc[y1:y2, x1:x2] += prob cnt[y1:y2, x1:x2] += 1.0 return acc / np.maximum(cnt, 1e-6)

stride小于win才有重叠，384 对 512 大约是 25% 重叠，再小推理时间线性上涨收益不明显。贴边那两行max(0, y2 - win)是保证每个窗口都是完整的win x win，否则网络见到比训练尺寸小的输入，BN 统计会偏移，边缘效果变差。

4.3 Dice / IoU 指标与叠加可视化

评估指标只算 Dice 容易偏乐观，配上 IoU 一起看更稳。两个指标都在二值化之后算，且要处理全背景样本：

def dice_score(pred, target, eps=1e-6): inter = (pred * target).sum() return (2 * inter + eps) / (pred.sum() + target.sum() + eps) def iou_score(pred, target, eps=1e-6): inter = (pred * target).sum() union = pred.sum() + target.sum() - inter return (inter + eps) / (union + eps)

可视化部分把掩码叠成红色盖在原图上，存盘前先转回 BGR，不然颜色会反：

def save_overlay(raw_gray, mask, out_path, alpha=0.4): base = cv2.cvtColor(raw_gray, cv2.COLOR_GRAY2BGR) color = np.zeros_like(base) color[mask > 0] = (0, 0, 255) # BGR 红色 overlay = cv2.addWeighted(base, 1 - alpha, color, alpha, 0) cv2.imwrite(out_path, overlay)

提示：评估集要和训练集严格隔离，同一批次拍的图分别进 train 和 val，验证 Dice 会虚高十几个点，这个结论拿去汇报会误导决策。

5. 常见问题排查：训练不收敛、显存爆掉与掩码全黑的现场记录

这一章是我从别人的代码和自己踩过的坑里筛出来的五条，每条都按现象、原因、解决写，照着对号入座能省掉大半天。

5.1 掩码全黑，验证 Dice 一直是 0

现象：训练 loss 能降，但验证集 Dice 从第一个 epoch 到最后一个都是 0.0000，可视化出来全黑。

原因：掩码读取时没加IMREAD_GRAYSCALE，或者标注导出的 PNG 是带透明通道的 RGBA，cv2.imread读出来是四通道，二值化后全部落在错误维度上；也有一种情况是掩码前景是 1 而不是 255，mask > 127把所有像素判成 0。

解决：读掩码后立刻打印mask.shape和np.unique(mask)，确认为(H,W)且取值只有 0 和 255 两个值；不是的话在Dataset里统一做一次归一化映射，把任意前景值压到 1。

5.2 训练 loss 一路下降，验证 Dice 卡在 0.3 不动

现象：训练集 Dice 能到 0.9，验证集死活上不去，两条曲线分叉得厉害。

原因：验证集的增强写错了，误用了训练那份带RandomResizedCrop的Compose；或者数据划分时把同一张原图的不同裁剪分别放进了 train 和 val，造成信息泄漏的反向版本——验证集分布和训练集差太多。

解决：验证变换只保留 Resize + Normalize；按原图 ID 做划分，同一来源的切片只能进同一侧；再确认model.eval()和torch.no_grad()都加上了，漏掉任何一个都会让 BN 用 batch 统计，结果随机性很大。

5.3 显存爆掉，batch size 一调就 OOM

现象：CUDA out of memory在第 2 个 epoch 或者换更大分辨率时出现，nvidia-smi显示显存被占满但利用率不高。

原因：UNet 在 512x512 下单样本特征图占用就不小，base=64时第一层就是 64 通道；另外循环里持有未释放的中间变量，或者num_workers太多导致主进程也在涨内存。

解决：先降分辨率到 256 验证能否跑通；再开混合精度torch.cuda.amp.autocast()，通常能省 30% 以上显存；最后调base=32或把batch_size降到 4 配合梯度累积，累积 4 步等效 batch 16，效果基本对齐。

5.4 输入尺寸不是 16 的倍数，报 shape 不匹配

现象：RuntimeError: Sizes of tensors must match except in dimension 1，报错位置在torch.cat。

原因：四次最大池化要求 H、W 都能被 16 整除，原始图直接送进去，遇到 250x250 这种尺寸，下采样取整后上采样回不来，skip connection 两边尺寸差 1 到 2 个像素。

解决：数据管线里统一A.Resize(256, 256)，别指望interpolate兜底；如果业务必须保留原始长宽比，就先 pad 到 16 的倍数，推理后再裁回去。

5.5 训练集效果好，换一张新图就崩

现象：验证集 Dice 0.88，产线上换一批光照条件不同的图，输出几乎全黑或者满屏噪点。

原因：归一化参数不一致，训练用 0.485/0.229，推理用了别的值；或者推理时没做/255，输入落在 0 到 255 区间，网络直接饱和。

解决：把归一化参数写进configs/train.yaml，训练和推理都从同一个配置读，杜绝两处硬编码；上线前抽 20 张新图跑一遍，人工看叠加图，比看指标更能发现问题。

6. 二次开发：换数据集、换 backbone 与导出 ONNX 的三个改造点

拿到这套代码之后，真正花时间的不是跑通，而是改成自己的业务。我按改动量从小到大排了三个方向，每个都给出验证方法，改完能立刻判断有没有改坏。

6.1 换数据集：只动三处

第一处是data/下的目录，保持images/和masks/同名；第二处是configs/train.yaml里的data_root、image_size、batch_size；第三处是models/unet.py里的in_ch，灰度图保持 1，RGB 图改成 3。多类别分割时num_classes改成类别数，损失换成nn.CrossEntropyLoss，掩码在Dataset里保持np.uint8的 0/1/2 索引，不要再乘 255。

验证方法很直接：先只跑 2 个 epoch，看train_loss是否在下降、验证 Dice 是否大于 0.1。如果第一个 epoch 结束 Dice 还是 0，八成是掩码格式或归一化的问题，回到第 5 章对照排查，别急着调超参。

6.2 换 backbone 或加注意力

想换编码器，把DoubleConv堆叠替换成预训练 ResNet 的下采样层，接口约定只有一条：每一层输出通道数要和Up模块里torch.cat的期望对上。最省事的做法是保留Down/Up的签名，只替换Down内部实现，把MaxPool2d + DoubleConv换成resnet.layer1这类模块，输入输出通道在构造时手动传。

加注意力模块也是同样思路，插在Up的torch.cat之后、DoubleConv之前，对拼接后的特征做通道加权。改完先跑一次torchsummary或直接print(model)，确认参数量没有异常暴涨（比如从 31M 变成 300M，多半是通道数算错了），再开训。

6.3 导出 ONNX 并做数值对齐验证

部署侧最常问的是能不能不依赖 PyTorch 跑。导出本身一行的事，难点在动态尺寸：

import torch model.eval() dummy = torch.randn(1, 1, 256, 256).to(device) torch.onnx.export( model, dummy, "unet.onnx", input_names=["input"], output_names=["logits"], dynamic_axes={"input": {0: "batch", 2: "h", 3: "w"}, "logits": {0: "batch", 2: "h", 3: "w"}}, opset_version=12, )

导出后必须做数值对齐：同一张图分别用 PyTorch 和 onnxruntime 跑，比较两张概率图的最大绝对误差，正常情况应该在 1e-4 以内。超过 1e-3 就要查是不是interpolate的mode在 ONNX 里被降级了，或者 BN 的track_running_stats状态不一致。

我第一次给分割模型做导出时省了对齐这步，线上跑了两周才发现边缘像素差了 2 个灰度级，定位回去花了一整天。从那以后我每次动网络结构，导完 ONNX 都强制走一遍"同一张图、两个后端、逐像素比对"的流程，误差表打印出来贴在提交记录里。这套代码的价值不在于它多完美，而在于每个环节都留了可替换的口子，你顺着目录改下去就行。希望帮到你。

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