U-Net工业缺陷分割实战：从数据标注到ONNX INT8量化部署

做分割模型最怕什么？不是准确率上不去，而是项目验收那天模型还在训练笔记本里躺着，产线工控机上根本没有一条能直接跑通的路。我这次被朋友拉去过一道类似的关：某工厂的注塑件外观质检要上自动视觉系统，缺陷分割是核心算法环节，要求从数据采集一路做到部署在工控机上，而不是在 Jupyter 里秀个 mIoU 就完事。整个过程踩了不少坑，也沉淀下一条完整的“数据 → 训练 → 部署”链路。这一篇我就用这个项目当例子，把 U-Net 实战中真正重要的细节拆开讲，前面几篇讲了网络结构和损失函数，这篇我们直接串起整个项目。

1. 项目开题：先搞清分割对象和数据底账，再谈网络

很多人做分割项目第一反应是去翻论文、调网络结构，但我建议顺序反过来。先把“到底要分割什么、数据长什么样、业务上什么指标说了算”这三件事弄明白，再回来看网络选型。否则后面每走一步都会返工。

1.1 一个从产线落下来的真实需求

这个项目的背景很直接：注塑车间生产一批外壳零件，表面可能出现划痕、黑点、缺料、飞边四类缺陷，原来靠人工目检，效率低且标准不统一。要上的自动化系统结构大概是：工业相机拍照，工控机跑图像算法，输出“有缺陷 / 无缺陷”的判定结果，同时把缺陷区域圈出来供后续机械臂分选。

注意这里有个容易被忽视的关键点——业务要的是像素级分割掩膜，不是简单地输出一个“有缺陷”的标签。因为产线后续要根据缺陷的面积、位置做返修决策，如果只给一个图像级分类结果，设备没法判断缺陷到底在哪。所以模型的输出设计从一开始就定了：输出每个像素的缺陷概率图，然后再由后处理决定怎么转换成最终判定。

业务指标也很明确：

• 缺陷召回率要优先保证，漏掉一个带缺陷的部件，比多标一个正常部件严重得多。
• 单张图像在工控机上的推理时间要在 500ms 以内，工控机没有独立显卡，纯 CPU 跑。
• 初版只做“缺陷 / 正常”二分类，后续再细分成缺陷类型，降低上线风险。

1.2 数据收集、标注和数据集切分

数据来自产线上两套不同机台的工业相机，总共收集了 12800 张图像，原始分辨率 1280×1024，BMP 格式。这里有个细节：同一个零件在生产线上会被多个工位的相机拍到，不同帧之间高度相似。如果切割数据集时不注意，很容易把同一零件的照片同时分进训练集和测试集，最后验证指标虚高，部署时立刻露馅。

所以数据划分时必须按产品批次划分，而不是按单张图随机划分。同一批次的产品只会出现在训练或验证中的一侧，这样测出来的指标才接近真实分布。

标注也是个冗长过程。我们用开源的标注工具逐像素标注，标完后再做一轮“标注仲裁”。四个人标同一批图，差异大的样本拿出来重新讨论，以高年资质检员的结论为准。这一步很花时间，但值得做，因为 U-Net 这类像素级模型对掩膜噪声非常敏感，标注边界差几个像素都会直接拉低收敛上限。

1.3 项目的整体设定与数据底账

这里先把本文用到的项目配置列出来，后面所有分析都围绕这套参数展开：

数据底账我建议做一个简单表格记录：每一批图像的采集时间、机台编号、光照情况、缺陷标注比例。后面训练结束回头看，能帮你定位不少“为什么验证集指标高但现场效果差”的诡异问题。

2. 从 U-Net 基线到缺陷场景改造：注意力、深度监督和损失函数

2.1 为什么我仍从 U-Net 基线起步

分割领域现在可选择的东西很多，Transformer 类模型、基于掩膜的 SAM 系列都很热闹。但在这个项目里，我坚定地先把标准 U-Net 作为第一版跑通。原因是：第一，缺陷分割属于典型的密集小目标预测任务，U-Net 的编码器-解码器结构加上跳跃连接，天然能把浅层细节和深层语义融合起来，对这种“既要看到全局、又要抠得住局部边缘”的任务非常友好。第二，工业项目的核心不是刷榜单，而是稳定可复现，U-Net 的假设少、结构简单，出问题时容易排查。第三，当时数据集不算很大，预训练权重的作用仍然非常关键。用 Transformer 为主干的模型在小数据量下容易过拟合，而且推理代价高。

基线网络我用了 ResNet34 做编码器的 U-Net。为什么不从编码器-解码器全是随机初始化的纯 U-Net 开始？因为纯 U-Net 从头训练对数据量的要求更高，收敛也慢，在只有一万多张图的项目里不划算。用了 ImageNet 预训练的 ResNet34 后，前几个 epoch 收敛速度肉眼可见地快。

2.2 小目标缺陷场景下的关键改造：注意力门控

标准 U-Net 跑完第一版，发现一个典型问题：大块缺陷分割得很好，小划痕和低对比度黑点却经常漏掉或断裂。原因在于上采样路径从深层拿到的语义特征在逐层恢复分辨率时，会把背景区域的响应也一并放大，小目标的细节被稀释了。

我的改进思路不是推翻 U-Net，而是在跳跃连接处加入注意力门控。简单说，就是让上采样路径给跳跃连接一个“筛选信号”，把与目标区域（缺陷）相关的特征放行，无关背景特征压下去。类似下面这个简化模块：

import torch import torch.nn as nn class AttentionGate(nn.Module): def __init__(self, in_channels, gating_channels, out_channels): super().__init__() self.W_g = nn.Conv2d(gating_channels, out_channels, kernel_size=1) self.W_x = nn.Conv2d(in_channels, out_channels, kernel_size=1) self.psi = nn.Conv2d(out_channels, 1, kernel_size=1) self.relu = nn.ReLU(inplace=True) def forward(self, x, g): # x 是跳跃连接特征，g 是来自更深层的门控信号 g1 = self.W_g(g) x1 = self.W_x(x) alpha = torch.sigmoid(self.psi(self.relu(g1 + x1))) return x * alpha

这里门控信号来自解码路径更高层的特征，它已经包含了“哪里更像缺陷”的语义判断，用它去加权浅层细节特征，相当于让模型在看边缘细节时心里同时有个全局答案。实际用下来，小目标召回率提升明显，Params 增加却很少，比直接把网络换成更深的版本划算得多。

2.3 深度监督：让梯度直通浅层

另一个有效改造是深度监督。我在解码器的三个不同分辨率输出上各接了一个 1×1 卷积头，训练时同时计算三个头的损失，推理时只保留最终输出。这样中间层的特征也能直接收到监督信号，而不是等损失穿过多层卷积之后才模模糊糊传过来。

深度监督特别适合小目标分割：小缺陷像素占比极低，单靠最后一层交叉熵时梯度容易被背景淹没。中间层有了直接监督，模型会更早学会在低分辨率下“锁定疑似区域”。训练结束后可以去掉辅助头，对推理速度没有任何影响。

2.4 损失函数是怎么定下来的

损失函数我试过多种组合：BCE、Dice、Focal、以及 BCE+Dice 混合。

先说结论：

• 纯 BCE：效果一般，因为背景像素占绝大多数，模型学到最后倾向于输出全零掩膜，局部最优很难挣脱。
• 纯 Dice：整体分割效果好，但小目标仍然不稳，因为 DICE 的梯度在目标很小的时候波动很大，一个两个像素的错误就会引起明显震荡。
• Focal Loss：理论上是调不平衡的，对难样本更敏感，但超参数（alpha、gamma）比较敏感，需要额外调。
• BCE + Dice：稳定好用。BCE 对每个像素独立监督，Dice 负责从整体区域重叠上纠偏，两者互补。我的权重设为 0.4 : 0.6。

下面是我实际用的损失函数实现，去掉了项目里的私有处理，保留核心思路：

class BCEWithDiceLoss(nn.Module): def __init__(self, bce_weight=0.4, dice_weight=0.6): super().__init__() self.bce = nn.BCEWithLogitsLoss() self.bce_weight = bce_weight self.dice_weight = dice_weight def forward(self, logits, targets): bce = self.bce(logits, targets) prob = torch.sigmoid(logits) smooth = 1e-6 # 只在 batch 维度计算 Dice，避免小目标被大图稀释 intersection = (prob * targets).sum(dim=(2, 3)) union = prob.sum(dim=(2, 3)) + targets.sum(dim=(2, 3)) dice = 1 - (2 * intersection + smooth) / (union + smooth) dice = dice.mean() return self.bce_weight * bce + self.dice_weight * dice

这里有一个容易被忽视的小细节：Dice 损失的计算要按每个样本单独算，不能把整个 batch 的所有像素合在一起算。因为如果一个 batch 里绝大多数样本是纯背景，那整体 Dice 基本由背景决定，小目标样本的监督信号完全被淹没。按样本算 Dice 能保证每个样本都得到相对均衡的梯度贡献。

3. 训练阶段复盘：优化器策略、损失曲线与过拟合处理

3.1 数据增强：工业场景下请“保守”一点

图像分割项目里数据增强几乎是必做的，但增强策略不能照搬自然场景。一开始我把旋转、翻转、随机缩放、色彩抖动全打开，结果在验证集上指标不升反降。回溯后发现，问题出在色调和饱和度扰动上。产线相机的光照是相对固定的，缺陷的颜色特征（划痕发亮、黑点发暗）是重要判别依据，把图像色相大幅扰动，等于给模型学了不存在的“变色缺陷”。

最后保留的增强组合是：水平翻转、垂直翻转、90 度旋转、随机缩放（0.8~1.2）、亮度对比度小幅度扰动、轻微高斯噪声。而色彩相关的增强全部关闭或调成极小幅度。

归一化这里也提醒一句：如果编码器用了 ImageNet 预训练权重，输入必须使用 ImageNet 的 mean/std 归一化，否则预训练特征分布直接被破坏，效果可能不如不迁移。

3.2 优化器与学习率策略

优化器用的是 AdamW，不是普通 Adam。主要原因不是玄学，而是 AdamW 在权重衰减处理上更规范，和余弦退火学习率配合起来更稳。在一万张图像规模的数据集上，直接用 3e-4 的学习率微调整个网络容易把预训练编码器的特征冲掉。我的实际做法是分两个阶段：

• 阶段一（前 10 个 epoch）：冻结编码器参数，只训练解码器。此时学习率走 warmup，从 1e-4 升到 3e-4。
• 阶段二（第 10 个 epoch 之后）：解冻编码器，全部参数一起微调，学习率按余弦退火降到 1e-5 附近。

这样做的理由是：解码器初期对缺陷特征的理解是混乱的，如果同时更新编码器，预训练权重的有效信息很容易被噪声梯度打乱。先把解码器训练到能基本框出目标区域，再让编码器去微调特征，收敛路径稳得多。

实际损失曲线大致趋势是：前 3 个 epoch Dice 从接近 0 快速爬到 0.6 左右，第 10 个 epoch 解冻编码器后有小幅波动，然后继续上升，30 个 epoch 之后逐渐平稳。到第 45 个 epoch 左右验证集 Dice 基本不再变化，再往下训练就开始出现训练损失继续下降、验证损失回升的过拟合迹象。

3.3 训练中的过拟合信号和应对手段

过拟合在这个项目里来得比想象中早，原因是缺陷像素占比太低，模型很容易把训练集里某些背景纹理当作缺陷记忆下来。我观察到的一个明显信号是：训练集 Dice 一路上升到 0.92 以上，但验证集 Dice 在 0.86 附近徘徊，同时验证集上把部分正常反光区域误判成缺陷的样本变多了。

应对手段除了常规的早停，还做了两件事：

1. 在训练集中按缺陷占比做重采样：把完全没有缺陷的纯背景图部分下调采样，保留一部分（大约 30%），因为模型需要见过足够多“干净”背景才能学会压制假阳性。但如果全是干净图，模型会往“全输出背景”的方向偷懒。
2. 使用 label smoothing 处理标注噪声：分割掩膜是人工标注的，边缘几个像素本来就有不确定性，把它们从硬标签 0/1 改成 0.9/0.1，能有效抑制模型在边缘处过度自信的尖峰响应。

3.4 评价指标的经验之谈

项目里最终看的不是 mIoU 一个数。mIoU 对背景占比极其敏感，在这个数据集上背景占 97% 以上，一个输出全零的“傻瓜模型”也能拿到 0.94 以上的 mIoU，所以它只能做参考。真正被我盯死的是缺陷召回率和每张图的误报缺陷数。

我们在验收时定下了一个口径：同一条阈值下，模型在所有标注有缺陷的零件上至少召回 95%，同时允许把少量正常零件判成返修，但误报率要低于 10%。这个口径直接指导了后面部署阶段的阈值选择。

4. 部署链路实落：ONNX 导出、INT8 量化与推理服务拆分

4.1 为什么生产环境不再直接跑 PyTorch

训练和推理在不同环境下的诉求完全不一样。PyTorch 在训练时动态图、灵活的 Python 控制流都是优点，但到了纯 CPU 的工控机上，动态图的解释开销、算子调度开销就成了实打实的性能负担。我当时把 PyTorch 的 FP32 模型直接放到工控机上跑，单张 512×512 推理大约要 120ms 以上，虽然勉强达标，但 CPU 稍微忙一点就会飙到 200ms+，风险很大。

标准做法是先把 PyTorch 模型导出为 ONNX。ONNX 是静态图，推理框架可以做算子融合、常量折叠、内存复用等优化，而且后续换推理引擎也方便。导出代码大致如下：

import torch torch.onnx.export( model, dummy_input, # 形状为 (1, 3, 512, 512) 的随机张量 "unet_defect.onnx", input_names=["images"], output_names=["logits"], dynamic_axes={"images": {0: "batch"}, "logits": {0: "batch"}}, opset_version=17 )

导出后最好做一个精度比对：随便挑几十张图，对比 PyTorch 推理和 ONNX Runtime 推理的输出结果，确认最大绝对误差在 1e-4 量级，再继续往下走。这一步能拦住很多奇奇怪怪的算子兼容问题。

4.2 INT8 量化怎么做才不掉点

ONNX 导出后的 FP32 模型在 ONNX Runtime 上单张大概能压到 80~90ms，但离“稳定、低负载”还差一口气。于是对模型做 INT8 动态量化。

所谓动态量化，是只把权重量化到 INT8，推理时会根据每层输入动态计算量化参数。实现非常简单：

from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_input="unet_defect.onnx", model_output="unet_defect_int8.onnx", weight_type=QuantType.QInt8, op_types_to_quantize=["Conv", "MatMul", "Gemm"], per_channel=True )

量化后模型体积从约 48MB 降到 13MB，单张推理时间从 80ms 左右降到 19ms 左右，效果非常明显。但量化会带来精度损失，主要损失集中在缺陷边缘和微小划痕上。我测下来 FP32 的验证集 IoU 大约是 0.853，INT8 动态量化后降到了 0.836，还能接受。如果你遇到更明显的掉点，可以考虑保留最后几层为 FP32，只量化前面的卷积层；更稳妥的方案是做量化感知训练（QAT），让模型在训练时就去适应量化误差，但那会多花不少实验时间。

4.3 推理服务拆分成三段，别把预处理和后处理混在一起

模型只是整条链路的一环。工控机上跑的推理服务不能只做“调用 ONNX Session”，而是要把三块逻辑分开：

1. 预处理：读图、resize 到 512×512、转 RGB、按 ImageNet 归一化、转成 NCHW 张量。
2. 推理：ONNX Runtime 加载 INT8 模型，跑一次 inference，输出 512×512 的概率图。
3. 后处理：把概率图上采样回 1280×1024 原图尺寸，按阈值转成二值掩膜，做连通域分析，过滤掉面积过小或长宽比异常的噪声区域，最后输出缺陷掩膜编码和判定结果。

这里有个容易踩的坑：ONNX Runtime 的 session 不要每次请求都重新创建。Session 初始化会加载模型、分配内存，耗时可能比一次 inference 还长。正确做法是进程启动时加载一次 session，之后所有请求复用。同时可以设置 intra_op_num_threads 和 inter_op_num_threads，在工控机 8 核 CPU 上我分别设为 8 和 1，实测比默认配置快约 20%。

后处理里另一个重要细节是阈值。训练时我们通常默认用 0.5 作为二值化阈值，但业务上更关心召回率，所以上线前我把阈值从 0.5 逐步降到 0.3，观察召回率和误报率变化。最终选在了 0.35：召回率从 91% 升到 96%，误报率仍然控制在 10% 以内。这个阈值以后要作为模型产物的一部分一起保存，否则版本迭代时配置容易丢失。

推理服务我用 FastAPI 包了一层，接口返回掩膜的 RLE 编码或 PNG base64。示例代码简化后如下：

import cv2 import numpy as np from fastapi import FastAPI from pydantic import BaseModel import onnxruntime as ort import base64 app = FastAPI() sess = ort.InferenceSession("unet_defect_int8.onnx", providers=["CPUExecutionProvider"]) class Req(BaseModel): image_base64: str def decode_b64_to_image(raw: str) -> np.ndarray: data = base64.b64decode(raw) arr = np.frombuffer(data, np.uint8) img = cv2.imdecode(arr, cv2.IMREAD_COLOR) return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) def preprocess(image: np.ndarray) -> np.ndarray: image = cv2.resize(image, (512, 512)) image = image.astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) image = (image - mean) / std return np.transpose(image, (2, 0, 1))[None, ...] @app.post("/segment") def segment(req: Req): image = decode_b64_to_image(req.image_base64) x = preprocess(image) logits = sess.run(None, {"images": x})[0][0, 0] prob = 1.0 / (1.0 + np.exp(-logits)) mask = (prob > 0.35).astype(np.uint8) mask_full = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST) # 这里可以继续做连通域分析和噪声过滤 return {"defect": bool(mask_full.max()), "mask": base64.b64encode(mask_full).decode()}

4.4 灰度上线的思路

模型部署上线不稳定是常态，直接切断人工巡检线风险太大。我们的做法是影子模式：工控机同时接模型输出和人工判断，前两周只记录模型结果和人工结果的差异，不参与实际分选。累计了两周数据后，挑选一致性高的批次逐步切到自动分选。灰度过程中还要记录推理延迟的波动曲线，因为加载了其他质检工序（比如 OCR 读取、几何测量）之后，CPU 负载会显著变化，绝不能在空载环境下测一次延迟就以为大功告成。

5. 复盘笔记：磨坏过赛乎的真坑与沉淀下的可靠方法论

5.1 中途踩过的三个典型坑

第一坑是训练数据里混进了重复图像。产线相机拍图时可能会在短时间内连续抓拍同一个零件多次，整理数据时如果不做图像去重，某些高度相似的图像会同时出现在训练和验证集中，验证指标虚高。后来我把所有图像计算感知哈希，设定阈值把相似度超过 97% 的图像全部标记出来，再按批次做一次人工确认，才彻底清掉这个隐患。

第二坑是标注掩膜和原图错位。标注工具里不小心把标注文件导出顺序弄乱，导致掩膜和图像不是一一对应。这种错位非常隐蔽，训练时损失也能下降，但验证集会莫名其妙的出现一个条目对不齐、mIoU 震荡很大的现象。后来给每个样本增加了一一对应的校验逻辑，训练前自动检查尺寸、通道数、文件名哈希，从流程上杜绝。

第三坑是纯背景负样本占比失衡。一开始我把所有背景图全部丢弃，只保留带缺陷的图像训练。结果模型在真实产线上的误报率非常高，看到任何不熟悉的反光纹理都当成缺陷。后来留了部分干净背景图进训练集，模型才学会“这不是缺陷”的边界。其实这类问题在分割项目里非常常见，处理原则是：背景样本不能没有，但也不能过剩到淹没前景监督。

5.2 我给这个项目留下的几条判断标准

模型选型上，不是结构越新越好，而是要和数据规模、部署约束匹配。在这个项目里，ResNet34 编码器 + 注意力 U-Net 是目前性价比最高的组合：参数少、推理快、预训练好找、部署稳定。

损失函数上，BCE + Dice 的组合足够应对 90% 的缺陷分割场景，先用简单方案把基线和数据流程跑通，再根据具体短板去换 Focal 或 Lovasz，顺序很重要。

部署方案上，ONNX Runtime 加 INT8 动态量化是 CPU 推理的最短路径。如果项目预算够、现场有显卡，再考虑 TensorRT；但 CPU 场景下先把 ONNX Runtime 的性能吃透，大部分问题已经能解决。

5.3 后续迭代还可以往哪里走

这个项目目前跑的是粗粒度“缺陷/正常”二分类，下一步是把缺陷类型细分。我的想法是先用一个模型输出总掩膜，再做四个细分类模型分别判断缺陷类别，而不是一上来就训练一个五分类分割模型。后者面对的类间不平衡会更严重，调参成本高，而且初期细分类别标准还没稳定，容易反复返工。

数据层面，产线上每天还在产生大量未标注图像。用已标注的小批量数据加上半监督方法去给新数据做伪标签，再人工抽查修正，是提升模型泛化能力性价比很高的路径。标注成本能省下一大半。

如果让我重新做这个项目，我会把更多的精力花在数据版本管理和实验记录上。哪怕只是把所有“图像-标注-模型权重-阈值”打包成一个带标签的目录，后续迭代也能省掉大量回忆时间。至于模型本身，U-Net 这套东西看起来“老”，但在工业视觉这种数据量不算特别大、目标边缘极其讲究的场景里，依然非常能打。真正的差距，往往不在模型结构有多新颖，而在整个数据到部署的闭环打得多扎实。