UNet图像分割实战：PyTorch从掩膜生成到滑窗推理与ONNX校验-北京尧图网络

简介面向Python图像处理与深度学习实践者的U-Net图像分割资料包适合希望从数据准备、模型构建、损失函数选择到训练预测完整走通分割流程的初中级开发者也可用于医学影像、遥感等场景的入门复现。压缩包共21个文件、约5.6MB以jpg/png示例图与效果对比图为主辅以Python核心脚本、Markdown说明文档、依赖清单及LICENSE另有gif动图直观展示分割前后差异。已有10164人学习下载具备较好的参考热度。包内代码围绕U-Net收缩与扩展路径、跳跃连接、交叉熵与Dice损失、Adam等优化器设置展开并涉及图像切片与平滑融合预测思路可帮助读者理解小目标与类不平衡条件下的分割处理。借助说明文档和示例素材读者能较快搭建训练与验证流程对照效果图检查模型输出并在此基础上替换自建数据集做迁移实验。1. 从一张图到一张掩膜UNet 图像分割先解决什么问题某质检小组每天要从产线相机里导出几千张图人工圈出划痕和脏污的位置再量面积归档。这种活用传统阈值加形态学也能做但一换光源、一换批次参数就得重调属于典型的“调参两小时上线五分钟崩”。用 Python 做图片处理时UNet 是这条路上性价比最高的一档几百张标注图、一张消费级显卡、一天之内就能跑出可用的逐像素掩膜。它给的不是检测框而是每个像素的 0/1 判定面积、周长、连通域数量都能直接从掩膜上算出来。适合样本量在几百到几千、目标是二分类或少量类别、对边界精度有要求的场景如果目标极大且形状规则先别急着上网络传统方法更快。2. 拆开 UNet编码器、解码器与跳跃连接的工程写法先把结构讲清楚再动手否则后面调不通的时候你连该怀疑哪一层都不知道。UNet 的形状像一个“U”左边一路下采样把空间尺寸压小、通道数翻倍右边一路上采样把尺寸还原中间用跳跃连接把左边的特征直接搬到右边拼接。整个过程没有全连接层输入多大就输出多大这也是它能处理任意尺寸图片的根本原因。2.1 编码器-解码器为什么能同时保住语义和边界下采样路径收缩路径每经过一次池化特征图边长减半、感受野翻倍。走到最底部时一个神经元“看到”的原图范围已经覆盖了很大一块区域它携带的是“这块地方整体像不像缺陷”这种语义信息但位置已经模糊了。上采样路径扩张路径负责把这种粗粒度判断还原回像素级位置问题在于光靠反卷积从低分辨率恢复边界一定是糊的。跳跃连接就是给解码器补细节的通道。编码器第一层的特征图分辨率最高、位置信息最准直接拼到解码器最后一层的输入上等于告诉它“语义你按上采样来边界照我这个抄”。这也是 UNet 在小样本数据集上比同量级 FCN 更容易收敛的原因——梯度有了一条短路径可以回传浅层不至于训不动。代价是显存。拼接操作让解码器每层的输入通道数翻倍第一层解码的输入通道是base*8 base*1这种量级显存占用比纯 encoder-decoder 高出三成左右。如果显存吃紧优先砍base_channels别砍层数层数一少感受野就不够了。配置项常用取值影响base_channels32 / 6432 约 7~8M 参数适合单卡 8G64 约 31M边界更细但吃显存下采样次数4 次输入尺寸需为 16 的倍数工程上一般 pad 到 32 的倍数上采样方式双线性插值卷积 / 转置卷积前者无棋盘格伪影后者可学习但易产生网格输出通道1二分类/ N多类单通道配 sigmoid/BCE多通道配 softmax/CE2.2 用 PyTorch 把双卷积块和上下采样写成可复用模块网上流传的 UNet 实现有一堆变体差别主要在归一化层和上采样方式。我一般固定用Conv-BN-ReLU双卷积双线性上采样理由是 BN 让学习率好设双线性上采样不会像转置卷积那样在掩膜上留下周期性网格——这个网格在放大看的时候特别明显后处理都救不回来。import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): 两层 3x3 卷积每层后接 BN 和 ReLU 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) class Down(nn.Module): 先最大池化降采样再走双卷积 def __init__(self, in_ch, out_ch): super().__init__() self.pool nn.MaxPool2d(2) self.conv DoubleConv(in_ch, out_ch) def forward(self, x): return self.conv(self.pool(x)) class Up(nn.Module): 先上采样再与跳跃连接拼接最后双卷积 def __init__(self, in_ch, skip_ch, out_ch): super().__init__() # 用双线性插值避免转置卷积的棋盘格伪影 self.up nn.Upsample(scale_factor2, modebilinear, align_cornersFalse) self.conv DoubleConv(in_ch skip_ch, out_ch) def forward(self, x, skip): x self.up(x) # 尺寸兜底即使前面 pad 过这里也再对齐一次 if x.shape[-2:] ! skip.shape[-2:]: x F.interpolate(x, sizeskip.shape[-2:], modebilinear, align_cornersFalse) x torch.cat([skip, x], dim1) return self.conv(x) class UNet(nn.Module): def __init__(self, in_ch3, num_classes1, base32): super().__init__() self.in_conv DoubleConv(in_ch, base) self.down1 Down(base, base * 2) self.down2 Down(base * 2, base * 4) self.down3 Down(base * 4, base * 8) self.down4 Down(base * 8, base * 16) # bottleneck self.up1 Up(base * 16, base * 8, base * 8) self.up2 Up(base * 8, base * 4, base * 4) self.up3 Up(base * 4, base * 2, base * 2) self.up4 Up(base * 2, base, base) # 1x1 卷积映射到类别数二分类时 num_classes1 self.head nn.Conv2d(base, num_classes, kernel_size1) def forward(self, x): x1 self.in_conv(x) # 1/1 x2 self.down1(x1) # 1/2 x3 self.down2(x2) # 1/4 x4 self.down3(x3) # 1/8 x5 self.down4(x4) # 1/16 x self.up1(x5, x4) x self.up2(x, x3) x self.up3(x, x2) x self.up4(x, x1) return self.head(x) # 输出 logits不做 sigmoid几个参数说清楚。base32时整个网络参数量在 7~8M 这个量级8G 显存跑batch_size8、512x512输入基本能压住base64参数量翻四倍边界确实更细腻但除非你的目标只有几个像素宽否则收益不明显。head输出的是原始 logits不要在里面加sigmoid因为后面要用BCEWithLogitsLoss这个损失内部已经融合了 sigmoid重复加会导致数值不稳和梯度变小——这是新手最常见的翻车点之一。另外注意biasFalse卷积后面紧接 BNBN 的 beta 参数已经承担了偏置的作用再加卷积偏置是冗余的。2.3 跳跃连接的尺寸对齐concat 报错的根因与修法如果你直接拿一张 1000×750 的图喂进去大概率会在torch.cat那行炸掉报Sizes of tensors must match之类。原因很直白四次 2×2 池化要求边长能被 16 整除否则每次下取整都会丢一点来回四趟之后上采样恢复出来的尺寸和编码器留下的 skip 对不上。1000 除 16 是 62.5池化取整成 62上采样回来是 992跟 1000 差了 8 个像素。常见的两种修法。一种是在Up.forward里用F.interpolate强制对齐上面代码已经这么写了能跑通但会把边缘几个像素拉伸长宽比极端时会轻微变形。另一种更稳推理前把图 pad 到 32 的倍数推理完再裁回原尺寸。我一般两种都用——数据侧做 pad网络侧留 interp 兜底双保险。import torch.nn.functional as F def pad_to_multiple(x, m32): 把张量右下角补零到 m 的倍数返回补齐后的张量和裁回用的尺寸 h, w x.shape[-2:] ph (m - h % m) % m pw (m - w % m) % m if ph 0 and pw 0: return x, (h, w) # modereplicate 比补零更接近真实边缘避免边界出现硬边 x F.pad(x, (0, pw, 0, ph), modereplicate) return x, (h, w) def unpad(x, size): return x[..., :size[0], :size[1]]这里modereplicate是个细节补零会让网络在右下角看到一块纯黑学出来的边界容易在那附近判错复制边缘像素更自然。如果你是训练时在线 pad、验证时也 pad网络见到的一直是补齐后的图一致性最好但要注意 pad 出来的区域在算 loss 时要 mask 掉否则网络会去学“补出来的那块不是目标”这种无意义信号。3. 数据管线从原始图片到能进网络的 mask 张量模型结构定下来之后真正决定最终 Dice 的是数据。我见过太多人网络改了三版没提升最后发现是 mask 读取时被 JPEG 压缩噪点污染了阈值一切边界全是毛刺。3.1 目录约定与 mask 二值化0/255 与 0/1 的致命差别先定目录别用“图片和掩膜放同一个文件夹靠文件名后缀区分”这种约定后期做划分、做增强都别扭。我一般这样组织dataset/ ├── images/ │ ├── train/ a001.png a002.png ... │ └── val/ b001.png b002.png ... └── masks/ ├── train/ a001.png a002.png ... # 与 images 同名 └── val/ b001.png b002.png ...标注工具导出的 mask 通常是 8 位灰度图目标区域 255、背景 0也有导出成 1 的。读进来必须先统一否则后面BCEWithLogitsLoss会拿到 0~255 的 target梯度直接爆掉。import cv2 import numpy as np def read_mask(path, threshold127): 读单通道掩膜并二值化为 0/1 的 float32 m cv2.imread(path, cv2.IMREAD_GRAYSCALE) if m is None: raise FileNotFoundError(path) # 大于阈值算前景 m (m threshold).astype(np.float32) return m def check_pair(img_path, mask_path): 训练前跑一遍尺寸不一致的样本直接暴露出来 img cv2.imread(img_path) msk read_mask(mask_path) assert img.shape[:2] msk.shape, f尺寸不一致: {img_path} return img.shape[:2], float(msk.mean())threshold127是经验值如果你的标注是抗锯齿边缘边缘有 100~200 的过渡灰度阈值取中间会略微收缩目标取低一点比如 50会略微膨胀。二分类分割里这半像素的差别对 Dice 影响不大但对面积统计影响明显做质检计量的话建议统一取 127 并在文档里写死。另外check_pair这步别省我遇到过标注同事导出时勾了“保持长宽比”结果 mask 被等比缩放过训练时 resize 一视同仁地看着正常实际标签已经偏了几个像素Dice 怎么都上不去查了两天才发现。3.2 图像与掩膜同步增强插值方式选错等于毁标签数据增强必须图像和掩膜同步旋转、翻转、缩放、亮度扰动都要作用在两张图上。手写随机变换容易出岔子比较省事的做法是用albumentations它的A.Compose会保证同一组随机参数同时作用于 image 和 mask。import albumentations as A from albumentations.pytorch import ToTensorV2 def build_train_tf(size512): return A.Compose([ A.RandomResizedCrop(heightsize, widthsize, scale(0.5, 1.0), ratio(0.75, 1.33), p1.0), A.HorizontalFlip(p0.5), A.VerticalFlip(p0.2), A.RandomRotate90(p0.5), # 亮度对比度扰动只作用于图像mask 参数必须传 None A.RandomBrightnessContrast(brightness_limit0.2, contrast_limit0.2, p0.5), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ]) def build_val_tf(size512): return A.Compose([ A.Resize(heightsize, widthsize), A.Normalize(mean(0.485, 0.456, 0.406), std(0.229, 0.224, 0.225)), ToTensorV2(), ])这里有个坑得单独拎出来说。A.Resize和A.RandomResizedCrop对 mask 默认用的是最近邻插值这一点它做得对但如果你自己写cv2.resize把interpolation忘了传默认是INTER_LINEAR二值 mask 缩放之后会出现 0.3、0.7 这类中间值再经过(m 0.5)二值化细长目标就断了。所以网络侧的 mask 插值一律INTER_NEAREST没有例外。Normalize用的均值方差是 ImageNet 的统计量用预训练权重时必须照抄从头训练可以用你们自己数据集统计出来的值会稍微快一点收敛但差别有限不值得为它多写一套代码。ToTensorV2会把 HWC 的 uint8 转成 CHW 的 float同时把 mask 也转成张量这一步别自己手写容易漏掉维度顺序。3.3 DataLoader 参数与类别不均衡的处理Dataset 写起来很直白麻烦的是不均衡。缺陷分割里前景像素经常只占 0.5%~3%一个 batch 里可能整张图都是背景。这时候有两个抓手一是采样层面让含前景的图多被抽到二是损失层面给正样本加权。import os import numpy as np from torch.utils.data import Dataset, DataLoader class SegDataset(Dataset): def __init__(self, img_dir, mask_dir, tf, mask_thr127): self.names sorted(os.listdir(img_dir)) self.img_dir img_dir self.mask_dir mask_dir self.tf tf self.mask_thr mask_thr 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)) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) msk read_mask(os.path.join(self.mask_dir, name), self.mask_thr) out self.tf(imageimg, maskmsk) # 返回 float32 的 [1,H,W] 掩膜与单通道 logits 对齐 return out[image], out[mask].unsqueeze(0).float() def make_loader(ds, batch_size8, trainTrue): return DataLoader( ds, batch_sizebatch_size, shuffletrain, num_workers4, # 与 CPU 核数匹配别超过 8 pin_memoryTrue, # 有 GPU 时打开减少 H2D 拷贝开销 drop_lasttrain, # 训练时丢掉最后一个不满的 batchBN 更稳 persistent_workerstrain, # 复用 worker 进程避免每 epoch 重启 prefetch_factor2, )num_workers不是越大越好。设成 CPU 核数的 1~2 倍之后继续加收益递减反而可能因为进程间拷贝拖慢。pin_memoryTrue只在你有 GPU 的时候有意义纯 CPU 训练开着是浪费。drop_lastTrue在训练集很小的时候要慎用——如果训练集只有 20 张、batch 设 16那每个 epoch 只剩 1 个 batchBN 统计量极其不稳这种情况把 batch 降到 8 或者关掉drop_last。不均衡还可以用加权采样给每张含前景的图一个更高的采样概率from torch.utils.data import WeightedRandomSampler def build_sampler(ds, pos_ratio0.7): 让含前景的样本被抽到的概率更高 weights [] for name in ds.names: m read_mask(os.path.join(ds.mask_dir, name)) weights.append(1.0 if m.mean() 0 else pos_ratio) return WeightedRandomSampler(weights, num_sampleslen(weights), replacementTrue)用了WeightedRandomSampler之后DataLoader 里的shuffle必须设成False两者互斥同时开会直接报错。4. 训练循环损失函数、指标口径与学习率调度网络和数据都就位剩下的就是训练。这一节的东西看着琐碎但每一项都影响最终能不能收敛到 0.85 以上。4.1 BCE、Dice 与组合损失谁在前面承担主要任务二分类分割最常用的三个损失BCEWithLogitsLoss、Dice Loss、以及两者的加权组合。BCE 是逐像素独立计算的梯度稳定但在前景占比 1% 的数据上把所有像素都判成背景就能拿到 99% 的准确率模型很容易躺平。Dice 是区域级别的度量直接优化预测和标签的重叠度对不均衡天然友好缺点是训练早期梯度波动大容易震荡。我的默认配置是组合BCE 承担主要工作Dice 作为修正项。权重一般 1:1前景占比低于 0.5% 时把 Dice 权重提到 2。import torch import torch.nn as nn import torch.nn.functional as F class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, logits, target): prob torch.sigmoid(logits) prob prob.reshape(prob.size(0), -1) target target.reshape(target.size(0), -1) inter (prob * target).sum(dim1) union prob.sum(dim1) target.sum(dim1) dice (2 * inter self.smooth) / (union self.smooth) return 1 - dice.mean() class ComboLoss(nn.Module): def __init__(self, pos_weight5.0, dice_w1.0): super().__init__() # pos_weight 直接放大正样本的 BCE 梯度 self.bce nn.BCEWithLogitsLoss( pos_weighttorch.tensor([pos_weight])) self.dice DiceLoss() self.dice_w dice_w def forward(self, logits, target): return self.bce(logits, target) self.dice_w * self.dice(logits, target)pos_weight的含义是正样本损失的放大倍数取负样本数/正样本数是理论值实际我会从 3 到 10 之间试。设太大比如 50会让模型疯狂预测前景掩膜上出现大片误检。smooth1.0是为了防止某张图完全没有前景时除以零也让梯度在早期更平滑别改成 1e-7 这种小值那样空 mask 的图会贡献巨大梯度。注意DiceLoss里对 logits 先做sigmoid再算这个顺序不能反。arch 输出的如果是概率BCEWithLogitsLoss会二次 sigmoidloss 会卡在 0.69 附近不动——这是我最常看到的“训练不收敛”求助场景。4.2 一个能直接跑的训练与验证循环下面这段是完整骨架包含了 AMP 混合精度、梯度累积、best 权重保存和早停计数。import torch from torch.amp import autocast, GradScaler def train_one_epoch(model, loader, optimizer, criterion, scaler, device, accum_steps1): model.train() total 0.0 for i, (img, msk) in enumerate(loader): img, msk img.to(device, non_blockingTrue), msk.to(device, non_blockingTrue) with autocast(device_typecuda, dtypetorch.float16): logits model(img) loss criterion(logits, msk) / accum_steps scaler.scale(loss).backward() if (i 1) % accum_steps 0: # 先反缩放梯度再裁剪否则裁的是被放大过的值 scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad(set_to_noneTrue) total loss.item() * accum_steps return total / len(loader) torch.no_grad() def validate(model, loader, criterion, device): model.eval() total, tp, fp, fn 0.0, 0, 0, 0 for img, msk in loader: img, msk img.to(device), msk.to(device) logits model(img) total criterion(logits, msk).item() pred (torch.sigmoid(logits) 0.5).float() tp ((pred 1) (msk 1)).sum().item() fp ((pred 1) (msk 0)).sum().item() fn ((pred 0) (msk 1)).sum().item() # 用全局累计的混淆矩阵算指标避免小目标被平均掉 dice 2 * tp / (2 * tp fp fn 1e-7) return total / len(loader), diceaccum_steps是梯度累积显存不够时用它把等效 batch 撑大。用累积的时候损失要除以accum_steps再 backward否则梯度是累加而不是平均等效学习率会被放大跑几次就飞了。clip_grad_norm_(max_norm5.0)对分割任务是个安全值加 Dice 损失后偶尔会出现梯度尖峰裁剪一下能救回来。学习率调度我一般用CosineAnnealingLR配 warmup前 200 步线性升温到初始学习率后面余弦衰减。Adam 的初始 lr 取1e-3配base32比较稳base64或者用 SGD 时降到1e-4。batch 大小翻倍、lr 跟着翻倍这条规则在分割上不太适用因为 batch 一大BN 统计量准了但每步的有效样本多样性没变我一般是 batch 加倍 lr 只乘 1.4。早停计数器配合验证 Dice 用连续 15 个 epoch 没刷新 best 就停同时保留 best 权重。别按 train loss 挑模型加了 Dice 损失之后 train loss 一直在降但验证 Dice 早就见顶了这是第二个高频误判。4.3 IoU 与 Dice 的口径别被逐图平均值骗了聚合指标有两种算法。一种是每张图各算一次 Dice 再取平均另一种是把整个验证集的 TP/FP/FN 累加后算一次。两者在大目标上差别不大但当验证集里有大量空 mask 图整张图没有前景时前者会给出虚高的数字。想想看一张完全没有前景的图如果模型也预测全背景按2TP/(2TPFPFN)算分子分母都是 0我上面的代码加了1e-7兜底结果会被算成 0拉低平均如果你在每图计算时给空图跳过那平均 Dice 可能虚高十几个点。更麻烦的是有些实现会把空图的 Dice 记成 1.0这时候你看到 0.9 的验证分数实际在非空图上可能只有 0.7。我的做法是始终报三个数全局累计 Dice、全局累计 IoU、以及非空图的平均 Dice。上线决策只看全局累计因为它和实际部署时的整体表现最接近。另外阈值也会影响这个数字——0.5 只是默认真正调优时要在验证集上扫一遍 0.1 到 0.9选全局 Dice 最高那个点这个后面还会说。5. 避坑与排查五个让 Dice 卡在 0.6 不动的现场这一节是我攒下来的血泪经验每一条都对应过至少一次“看起来哪儿都对就是不收敛”的排查过程。按现象、原因、解决三段来记。5.1 数据与标签层面的三个坑坑一loss 一路降到 0.05验证 Dice 死活过不了 0.6。现象是训练集上的 Dice 能到 0.95验证集却在 0.6 附近横盘且验证 loss 从第三个 epoch 开始缓慢上升。原因通常是训练集和验证集划分时按“裁剪块”随机分同一张原图裁出来的多个 patch 同时进了训练和验证模型等于见过答案。解决办法是按原图 ID 分组划分同一张图的全部 patch 只能落在一边划分比例按图数而不是 patch 数算。改完之后验证 Dice 一般会掉 5~10 个点但那才是真实水平。坑二预测掩膜上细长目标断裂或者干脆整片消失。现象是面积大的目标分割得很好宽度只有几个像素的细线状缺陷完全丢。原因是RandomResizedCrop或Resize把图缩到 512 之后原本 8 像素宽的缺陷只剩 2 像素再经过池化就没了。解决有两条路一是把输入尺寸提到 768 或 1024代价是显存翻倍二是把大图裁成 patch 训练而不是整图缩放。我一般优先选裁 patch因为缺陷分割本来就不需要看全图512x512的 patch 里目标占比反而更高正负样本更均衡。坑三验证集 Dice 波动极大同一个模型跑两次差 0.15。现象是每次重新训练结果差异巨大没有可复现性。原因有三没设随机种子、num_workers0时每个 worker 的随机状态不同、以及数据增强里的RandomResizedCrop每次都变。前两个好解决设torch.manual_seed并在worker_init_fn里给每个 worker 设种子。第三个是玄学重灾区验证集绝对不能用随机增强验证流程必须确定性。我见过有人图省事训练和验证共用一个 transform验证分数每次都不一样调参完全失去参照系。5.2 训练与推理不一致的两个坑坑四训练时一切正常部署单张推理结果抖动同一张图两次输出不一样。现象是同一张图连续推理两次掩膜边缘有细微差别有时候连通域数量都变了。原因就一个忘了model.eval()。BN 层在 train 模式下用的是当前 batch 的均值和方差单张图推理时这个统计量就是它自己和训练时累积的滑动平均完全不同输出自然乱。修法是推理前固定写model.eval()加torch.no_grad()并且把这两行封装进一个predict()函数里别指望每次手动记得写。坑五输入尺寸不是 16 的倍数时右下角预测全黑。现象是任意尺寸图片推理输出掩膜右侧和下侧边缘几十个像素永远是背景。原因在 2.3 节说过池化取整导致尺寸回不来。解决办法是推理前 pad 到 32 的倍数推理后裁回原尺寸同时在测试集上专门挑几张非 16 倍数的图做回归测试防止后续改动把这块弄坏。这条坑最阴的地方在于如果验证集是统一 resize 过的你根本发现不了上线遇到原始尺寸的图才爆。6. 进阶滑窗推理、阈值搜索与 ONNX 一致性验证模型训完只是半程。真实场景里图片动辄 4000×3000直接 resize 到 512 会把小目标抹平正确做法是滑窗推理按512x512切块步长取 256 让相邻块重叠一半每块单独预测后按位置拼回大图重叠区域取平均。重叠比例不能太小步长等于块大小时块与块的接缝处会有一条明显的判定跳变看起来像被人用刀切过。import numpy as np import torch torch.no_grad() def sliding_window_predict(model, img, device, patch512, stride256): img: [1,3,H,W] 已归一化的张量返回同尺寸概率图 model.eval() _, _, H, W img.shape prob torch.zeros((1, 1, H, W), devicedevice) count torch.zeros((1, 1, H, W), devicedevice) for y in range(0, max(H - patch, 0) 1, stride): for x in range(0, max(W - patch, 0) 1, stride): # 靠边时把窗口往内推保证每块都是完整的 patch y0 min(y, max(H - patch, 0)) x0 min(x, max(W - patch, 0)) tile img[:, :, y0:y0 patch, x0:x0 patch] p torch.sigmoid(model(tile.to(device))) prob[:, :, y0:y0 patch, x0:x0 patch] p count[:, :, y0:y0 patch, x0:x0 patch] 1 return prob / count.clamp_min(1)stridepatch//2是重叠 50%显存允许的话可以到patch//4接缝会更平滑代价是推理时间成倍增长。clamp_min(1)防止图片比 patch 还小时除零。阈值别用 0.5 就完事。在验证集上把阈值从 0.1 扫到 0.9步长 0.05画一条 Dice-阈值曲线取最高点。缺陷分割因为正样本少最优阈值经常落在 0.3~0.4 一带用 0.5 会漏检一批边缘弱响应的目标。扫完之后记得把这个阈值写进配置别留在 notebook 里。最后一步是导出 ONNX 做一致性校验这一步经常被跳过然后线上换个推理引擎就翻车torch.onnx.export( model, torch.randn(1, 3, 512, 512, devicedevice), unet.onnx, input_names[input], output_names[logits], dynamic_axes{input: {0: b, 2: h, 3: w}, logits: {0: b, 2: h, 3: w}}, opset_version12, )导出后必须拿同一张图分别跑 PyTorch 和 ONNX Runtime比对输出的最大绝对误差正常情况下应该在 1e-4 量级。如果差到 1e-2多半是动态轴没设对或者网络里有F.interpolate用了align_cornersTrue这种和导出引擎不一致的参数。从那以后我每次导出模型都强制走一遍这套回归固定 5 张覆盖不同尺寸和不同前景占比的图PyTorch 与 ONNX 各跑一次误差超阈值就阻断发布。这个习惯帮我拦下过两次因为换 opset 版本导致的静默精度下降那种问题不报错只是指标悄悄掉几个点靠肉眼根本发现不了。希望帮到你。本文还有配套的精品资源点击获取

UNet图像分割实战：PyTorch从掩膜生成到滑窗推理与ONNX校验

相关推荐

JSP相册管理系统实战：图片上传、缩略图与分页鉴权

JAVA第三周学习感悟

css盒子模型讲解

更多资讯

NLP与LightGBM构建高效客服投诉分类系统

VuePress vdoing 私密文章权限控制：可见性收敛与路由门禁

Windows x64下libhackrf静态编译与USB信号链路稳定实践

开发者如何高效记忆编程语法：认知科学与实战技巧

单链表数据结构：原理、实现与优化技巧

PSO优化RBF神经网络在电力负荷预测中的应用

今日更新

PLFM_RADAR：从被动告警到主动感知的微服务健康监控系统

SPI隔离设计实战：MAX14431BAEE+与STM32F407工业级应用

数字隔离器如何根治SPI通信地环路干扰：MAX14431与TM4C129实战

本周精选

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

本月热门

面向复杂光照与角度的宠物面部识别鲁棒性增强策略

基于DRV8818与STM32的双极步进电机驱动方案实战解析

DRV8818PWPR+PIC18F45K80双极步进电机控制方案实战

读完文章，想聊聊您的网站？