简介面向工业质检领域深度学习入门与实战的钢材表面缺陷检测项目基于UNet网络和NEU-DET数据集构建完整检测流程覆盖划痕、凹坑、夹杂等典型缺陷的自动识别与分割。资源包共2000个文件包含1800个xml标注文件和180个jpg图像样本以及Python训练脚本、ipynb分析笔记、ONNX模型文件等方便用户独立复现模型训练与推理。压缩包整体约122.16MB目录结构清晰适合需要掌握图像分割项目全流程的开发者参考学习。目前已吸引892人浏览学习项目中包含数据预处理、训练集/验证集/测试集划分、损失函数选取与参数优化、模型部署维护等关键环节的完整思路可帮助读者从数据准备到实际部署快速构建缺陷检测系统为自动化质量控制提供可直接借鉴的技术方案。1. 钢材表面缺陷检测项目UNet NEU-DET 这套组合好在哪钢材表面的裂纹、夹杂、斑块、麻点、氧化铁皮压入和划伤直接决定一块钢板能不能出厂。用肉眼做质检的工人一天要看上千张图疲劳之后的漏检率并不低。NEU-DET 数据集就是为这类问题准备的公开样本集1800 张 200×200 的灰度图标注精度足够支撑训练一个像样的分割模型。UNet 的优势在于它不需要特别大的数据量编码器、解码器加跳跃连接的结构在小样本场景下反而比很多重模型稳定让它成为入门缺陷检测项目时最顺手的模型之一。这篇文章记录的是如何用 NEU-DET 数据集在 PyTorch 框架下把 UNet 从标注解析、掩码生成、模型训练到推理验证一次跑通。适合刚接触缺陷检测的算法工程师也适合正在评估工业质检落地方案的从业者。2. NEU-DET 数据集与 UNet 原理搞懂输入输出才能动手写代码2.1 NEU-DET 数据集的六类缺陷与标注格式NEU-DET 总共有 1800 张 200×200 像素的灰度图把热轧带钢表面的缺陷分成六类每类 300 张crazing裂纹、inclusion夹杂物、patches斑块、pitted surface麻点、rolled-in scale氧化铁皮压入和 scratches划伤。这六类的形态差异很明显裂纹和划伤是又细又长的线条型夹杂和氧化铁皮压入是块状麻点则是密集的小点。形态差异直接决定了分割任务的难度线条型目标占像素比例极小训练时对损失函数的敏感度极高块状目标虽然占得多但标注框和真实区域之间的空隙也更大。我在开始处理这个数据集之前会先做一个抽样统计把每个类别的标注框面积和对应图像的尺寸做个比值。比如裂纹类的框内面积占比往往只有 1% 到 3%而夹杂、氧化铁皮压入这类块状缺陷能到 8% 以上。像素占比直接影响到模型收敛的方向占比过低的类别会让交叉熵损失里背景梯度占绝对优势模型很容易退化成“全部输出背景”的偷懒状态。所以拿到数据集第一时间算一下这个比例后面选择损失函数心里就有数了。标注格式上NEU-DET 每张图配一个 XML 文件里面记录文件名、图像尺寸、每个目标的类别名和边界框坐标整体结构接近 PASCAL VOC 的标注体系。这意味着网上能找到的 VOC 解析代码稍作修改就能用。但有一个关键点需要先想清楚标题里写的是缺陷检测NEU-DET 原始给的是检测框而不是像素级掩码而 UNet 要做的是分割逐像素判断每个点是不是缺陷。所以需要把检测框转换成分割标签最常见做法是框内像素填 1、框外填 0把外接矩形当成目标的粗略近似。这个近似有误差尤其细长缺陷的框里大部分像素反而是背景项目后期要对这部分做专门处理。解析 XML 是整个流程里最基础的一步我习惯用一个独立函数把标注读进内存返回文件名和目标列表。下面这段代码可以直接复用import xml.etree.ElementTree as ET def parse_voc_xml(xml_path): tree ET.parse(xml_path) root tree.getroot() filename root.find(filename).text objs [] for obj in root.findall(object): name obj.find(name).text box obj.find(bndbox) xmin int(round(float(box.find(xmin).text))) ymin int(round(float(box.find(ymin).text))) xmax int(round(float(box.find(xmax).text))) ymax int(round(float(box.find(ymax).text))) objs.append({name: name, bbox: [xmin, ymin, xmax, ymax]}) return filename, objs这段代码用 ElementTree 把 XML 里每个 object 节点的类别名和边界框提取出来返回一个字典列表。注意 xmin、ymin 这些坐标值在 XML 里是字符串要转成 int如果原始标注里有小数坐标需要先转 float 再取整直接 int() 一个数字字符串没问题但碰到带小数点的值会直接抛异常这里做了容错。文件名通常对应图像名后续拿它和图像目录做对齐能及时发现标注文件和图像文件缺失的批次问题。2.2 UNet 编解码结构为什么适合钢材缺陷分割UNet 名字来自网络形状像一个 U。左侧编码器逐层卷积加下采样把 256×256 的输入一路压缩到 16×16 左右通道数从第一层的 64 往上涨到 1024空间信息被压缩的同时语义信息变得更抽象。右侧解码器再把特征图一步步上采样还原到输入分辨率。在对称的两侧之间编码器每一层输出的特征图都会通过跳跃连接拼到解码器对应层的输入上这个拼接动作是 UNet 区别于普通编码器-解码器网络的精髓。下采样虽然让模型看到更大的感受野却也丢失了裂纹这种细目标的边缘细节而跳跃连接把浅层的高分辨率细节直接送到解码器让上采样恢复像素级分类时能参考到真实边缘。很多想入门缺陷检测的开发者一上来就选 YOLO 或 Faster R-CNN但如果你第一次做这个项目我更推荐先用 UNet 把分割基线跑通。原因是钢材表面缺陷和自然场景目标有一个明显不同缺陷没有固定的长宽比例裂纹可能从左上角延伸到右下角划伤就是一道细线麻点是散布的点群。检测框模型需要为每个目标预测一个框对细长目标天然不友好分割模型直接逐像素输出不需要强行用一个矩形去拟合一条曲线。后续要做缺陷面积统计、严重程度分级时像素级结果也更容易换算。UNet 的参数量和训练开销在分割模型里属于偏小的。标准 UNet 的编码器加解码器总共大约 3100 万个参数比大多数分类骨干网络都轻单张消费级显卡跑 256×256 输入、batch size 8显存占用大概在 4GB 到 6GB。NEU-DET 只有 1800 张图数据量不大UNet 这种浅层的跳跃连接结构在少量数据下不容易过拟合训练几个 epoch 就能看到明显效果。等你用 UNet 跑通基准线再想提升精度可以在 UNet 基础上换预训练骨干或者换成 DeepLabV3 这类带空洞卷积的模型但那是后话。有一点必须提前说明NEU-DET 虽然有六类缺陷但做分割时我不建议第一版就上六分类。六类缺陷的像素占比差异悬殊多分类会让类别不均衡问题成倍放大而且初始阶段你其实只需要知道“有没有缺陷”不关心具体是哪种。所以典型做法是先按二分类分割训练输出一个通道像素值大于阈值判定为缺陷区域。等二分类的 IoU 稳定了再扩展成多分类版本把六类标签加进训练流程。这个顺序能省掉大量排错时间也符合工业现场先告警再分类的实际诉求。2.3 数据预处理灰度图、归一化与标签编码的细节NEU-DET 图像是单通道灰度图读取方式决定后面所有张量的形状。用 OpenCV 的cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)读出来的是 H×W 的二维数组值域 0 到 255。送入模型前需要把它变成 1×H×W 的通道维度PyTorch 卷积层的输入规定是 N×C×H×W所以要在第一维插入一个通道轴用np.expand_dims(img, axis0)或torch.unsqueeze(0)都可以。如果有人图省事用三通道方式读灰度图等于把同一个灰度值复制三遍参数量没变但输入体积大了三倍没必要。归一化这里有两种常见路线。一种是直接除以 255把像素映射到 0 到 1 区间再配合 BatchNorm训练也能稳定收敛这是最快上手的做法。另一种是用均值 0.5、方差 0.5 做标准化把像素映射到 -1 到 1数值分布更对称梯度更新更平稳代价是推理时也要套同一套参数。我项目里一般会用后者因为实际感受是同样学习率下标准化后损失曲线波动明显更小尤其当 batch size 很小的时候BatchNorm 统计不稳定带来的干扰会被标准化抵消一部分。标签编码是另一个容易埋雷的地方。二分类分割的 mask 只有两个值背景是 0前景缺陷是 1数据类型用 int64PyTorch 的 LongTensor。如果误把 mask 存成 float后面配合CrossEntropyLoss用会被直接报错因为交叉熵要求标签是整数索引。如果用二分类配合BCEWithLogitsLoss则要求预测值不经过 Sigmoid、直接和 float 类型的 mask 计算标签得转成 float。这个类型匹配问题不搞清楚训练代码会在第一个 batch 就报错或者更糟糕的是不报错但结果错误。我见过有同学把 mask 转成了 float、模型输出只有一个通道、却用了CrossEntropyLoss结果训练了十几个 epoch 损失一直在 1.7 附近下不去后来一查才发现是标签编码和损失函数不匹配。最后是输入尺寸。标准 UNet 下采样四层之后还会再经过一层 bottleneck总共五次尺度减半也就是说输入图像的宽高最好是 16 的倍数而且实际下采样用的是向下取整的池化一旦宽高不能被整除解码器上采样回来和跳跃连接拼特征图时就会遇到尺寸对不上的报错。NEU-DET 原始图是 200×200不能被 16 整除所以数据加载时需要统一 resize 到 256×256 或者 192×192。我习惯用 256×256因为在多种公开实现里这是默认配置遇到问题网上容易搜到对照显存紧张的机器可以降到 192×192但需要确认模型里没有硬编码的尺寸常量。3. 从原始数据集到可训练样本数据准备阶段的完整流程3.1 依赖环境与项目目录组织开始写数据处理代码之前先把环境和目录结构准备好。Python 版本建议 3.8 到 3.10PyTorch 版本用 1.10 以上同时安装和 PyTorch 严格对应的 torchvision不要混装不同版本。图像读取用 opencv-python数据增强用 albumentations进度展示用 tqdm。没有 GPU 也能跑只是训练速度会慢三四倍NEU-DET 这个数据量 CPU 训练一轮一百多张图大概十几分钟跑 30 个 epoch 也能接受但每 epoch 时间较长调参效率很低有条件还是建议用 GPU。依赖库版本建议用途Python3.8 ~ 3.10基础运行环境PyTorch 1.10模型定义与训练torchvision与 PyTorch 同版本模型层与基础变换opencv-python 4.5图像读取与可视化albumentations 1.2数据增强与 mask 同步变换numpy 1.21数组与矩阵操作tqdm任意训练进度显示项目目录我一般按下面这个结构组织数据、代码、模型权重分开避免把训练产物混进源码目录neudet-unet/ ├── data/ │ ├── images/ # NEU-DET 原始图像 │ └── annotations/ # NEU-DET 原始 XML 标注 ├── src/ │ ├── dataset.py # 数据集类与预处理 │ ├── model.py # UNet 模型定义 │ ├── train.py # 训练脚本 │ └── predict.py # 推理脚本 └── checkpoints/ # 模型权重保存目录目录组织看起来简单但重要性不低。后续在产线上换数据、换模型版本、跑实验对比时一个清晰的目录能让你少踩很多坑。尤其是标注和图像分开放清洗数据时只需要动一个目录。3.2 解析 NEU-DET 标注并生成掩码数据准备的核心动作是把 XML 里的检测框转换成和图像尺寸一致的 mask 数组。上一步parse_voc_xml已经拿到了边界框这里要做的是新建一个全零矩阵把框内像素赋成前景值。但直接整框填充有个明显的缺点NEU-DET 的标注框是目标外接矩形对裂纹、划伤这类细目标来说框内背景占比可能超过 80%。如果 mask 完全按框填充等于告诉模型“框内全是缺陷”模型学到的分割结果会比真实缺陷胖一圈推理时把大量背景误报成缺陷。我一般会在框填充之后加一个边缘腐蚀操作把掩码边缘向内收缩几个像素缓解框和真实目标的误差。腐蚀核大小需要按目标尺寸调整3×3 对 200×200 的图风险偏低5×5 会明显压缩边缘。具体实现如下import cv2 import numpy as np def box_to_mask(image_shape, bbox, class_idx1, erode_kernel3): mask np.zeros(image_shape, dtypenp.uint8) xmin, ymin, xmax, ymax bbox mask[ymin:ymax, xmin:xmax] class_idx if erode_kernel 1: kernel np.ones((erode_kernel, erode_kernel), np.uint8) mask cv2.erode(mask, kernel, iterations1) return maskbox_to_mask接收图像尺寸、边界框和类别索引返回和原图等大的 mask。class_idx1表示前景类别编号二分类场景固定用 1多分类时传入对应类别编号。erode_kernel是腐蚀核边长传 3 意味着 3×3 核执行一次腐蚀前景区域的外圈像素被抹掉一圈这个参数需要按你数据里标注框的紧致程度来调。腐蚀的副作用是会把本来就小的目标削得更小如果缺陷像素占比本身很低可以跳过腐蚀直接使用整框标签用损失函数来平衡偏差后文会讲到。mask 生成之后要和图像一一对应检查。常见做法是在生成完所有 mask 后随机抽 10 张把原图和 mask 叠加可视化肉眼确认框和目标轮廓是否大致吻合。这一步虽然不产生代码产出但能发现标注文件里坐标错位、框超出图像边界等问题是最便宜的质检手段。坐标超出边界时NumPy 切片不会报错但会把 mask 形状截断后续训练张量拼接时尺寸对不上到那时候排查成本已经很高了。3.3 训练验证集划分与同步数据增强数据划分按类别分层抽样。NEU-DET 每类 300 张如果随机划分不控制类别比例可能某一类在训练集里只有 200 张验证集有 100 张同一类别在两个集合里的占比差异会导致验证指标失真。我一般用sklearn.model_selection.train_test_split按类别标签做分层划分训练集 1440 张、验证集 360 张比例 8:2每类的 240 张进训练集、60 张进验证集保证类别分布一致。数据增强用 albumentations因为它的分割增强 API 能同步变换图像和 mask避免自己写同步逻辑出错。针对钢材表面缺陷我常用的增强组合包括水平翻转、垂直翻转、旋转 90 度、亮度对比度微调。翻转和旋转对小数据集扩充作用明显亮度对比度模拟不同光照下钢材表面反射差异。这里有一个必须注意的点mask 在增强时要用最近邻插值不能用线性插值。mask 里只有 0 和 1 两种值线性插值会在边缘产生 0.3、0.7 这种中间值即使后面取整也会造成边缘锯齿错位。albumentations 默认对 mask 使用最近邻插值但如果你手动写 transform必须显式指定。数据集类封装成 PyTorch 的Dataset重点是把预处理、mask 生成、resize、增强全放在__getitem__里import os import cv2 import torch import numpy as np from torch.utils.data import Dataset from albumentations import Compose, HorizontalFlip, VerticalFlip, Rotate, Resize class SteelDataset(Dataset): def __init__(self, img_dir, ann_dir, img_names, img_size256, augmentFalse): self.img_dir img_dir self.ann_dir ann_dir self.img_names img_names self.img_size img_size self.augment augment if augment: self.transform Compose([ HorizontalFlip(p0.5), VerticalFlip(p0.5), Rotate(limit90, p0.3), Resize(img_size, img_size, interpolationcv2.INTER_LINEAR) ]) else: self.transform Compose([ Resize(img_size, img_size, interpolationcv2.INTER_LINEAR) ]) def __len__(self): return len(self.img_names) def __getitem__(self, idx): name self.img_names[idx] img_path os.path.join(self.img_dir, name) xml_path os.path.join(self.ann_dir, name.replace(.jpg, .xml)) img cv2.imread(img_path, cv2.IMREAD_GRAYSCALE) _, objs parse_voc_xml(xml_path) mask np.zeros(img.shape, dtypenp.uint8) for obj in objs: xmin, ymin, xmax, ymax obj[bbox] mask[ymin:ymax, xmin:xmax] 1 transformed self.transform(imageimg, maskmask) img transformed[image] mask transformed[mask] img img.astype(np.float32) / 255.0 img torch.from_numpy(img).unsqueeze(0) mask torch.from_numpy(mask.astype(np.int64)) return img, mask这个数据集类的核心逻辑是读灰度图、解析 XML、按边界框生成 mask、统一 resize、最后转成 PyTorch 张量。img.astype(np.float32) / 255.0把像素归一化到 0 到 1unsqueeze(0)增加通道维度变成 1×H×W。mask 用int64是因为它对应损失函数里的类别索引。augment 参数控制训练集和验证集是否走同一套变换验证集只做 resize 不做增强确保验证集反映真实分布。4. 用 PyTorch 搭 UNet 训练流程模型结构、损失函数与训练循环4.1 UNet 模型定义编码器、解码器与跳跃连接UNet 的代码实现有不少开源版本我常用的是一个结构清晰、适合改造成多分类的版本。核心模块是两个卷积加 BatchNorm 加 ReLU 的双卷积块编码器通过最大池化下采样解码器通过转置卷积上采样跳跃连接用torch.cat在通道维拼接。代码里把双卷积块抽成独立类避免重复写四段同样的结构import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels1, num_classes1): super().__init__() self.enc1 DoubleConv(in_channels, 64) self.enc2 DoubleConv(64, 128) self.enc3 DoubleConv(128, 256) self.enc4 DoubleConv(256, 512) self.pool nn.MaxPool2d(2) self.bottleneck DoubleConv(512, 1024) self.up4 nn.ConvTranspose2d(1024, 512, 2, stride2) self.dec4 DoubleConv(1024, 512) self.up3 nn.ConvTranspose2d(512, 256, 2, stride2) self.dec3 DoubleConv(512, 256) self.up2 nn.ConvTranspose2d(256, 128, 2, stride2) self.dec2 DoubleConv(256, 128) self.up1 nn.ConvTranspose2d(128, 64, 2, stride2) self.dec1 DoubleConv(128, 64) self.out nn.Conv2d(64, num_classes, 1) def forward(self, x): e1 self.enc1(x) e2 self.enc2(self.pool(e1)) e3 self.enc3(self.pool(e2)) e4 self.enc4(self.pool(e3)) b self.bottleneck(self.pool(e4)) d4 self.dec4(torch.cat([self.up4(b), e4], dim1)) d3 self.dec3(torch.cat([self.up3(d4), e3], dim1)) d2 self.dec2(torch.cat([self.up2(d3), e2], dim1)) d1 self.dec1(torch.cat([self.up1(d2), e1], dim1)) return self.out(d1)DoubleConv接收输入通道和输出通道内部做两次 3×3 卷积每次卷积后接 BatchNorm 和 ReLU这是 UNet 每一层的基本处理单元。编码器部分从输入单通道开始逐层把通道变成 64、128、256、512每层之间用 2×2 最大池化把空间尺寸减半。bottleneck 是网络最深的地方通道数 1024。解码器部分先用转置卷积把特征图尺寸翻倍再和编码器对应层的特征图在通道维拼接拼接后通道数等于解码器通道加编码器通道所以dec4的输入通道写成 1024实际是 512 拼接 512。in_channels和num_classes这两个参数让模型既能处理单通道灰度图也能应对多分类输出后面扩展六分类时只需把num_classes改为 7 并调整损失函数。4.2 损失函数选择交叉熵、Dice 与 Focal 的取舍损失函数是分割项目里最影响收敛质量的部件。NEU-DET 这类缺陷数据集的显著特点是前景像素占比普遍偏低直接使用普通交叉熵背景像素的梯度贡献会压过前景模型倾向于输出全背景。这不是模型能力问题而是损失函数的偏好问题。所以第一版建议直接用 Dice Loss 或者 Focal Loss。Dice Loss 的思路是按预测区域和真实区域的交叠比例计算损失不依赖像素总数对类别不平衡天然不敏感。它的公式是 1 减去两倍的预测与真实交集除以预测加真实我在项目里常这样实现import torch import torch.nn as nn class DiceLoss(nn.Module): def __init__(self, smooth1.0): super().__init__() self.smooth smooth def forward(self, pred_logits, target): pred torch.sigmoid(pred_logits) pred_flat pred.reshape(-1) target_flat target.reshape(-1) intersection (pred_flat * target_flat).sum() return 1 - (2.0 * intersection self.smooth) / ( pred_flat.sum() target_flat.sum() self.smooth )smooth是平滑项默认 1.0作用是避免前景和背景都为 0 时除以零同时让损失曲线更平滑。pred_logits是模型输出的未激活 logits先经过 Sigmoid 再和目标计算。reshape(-1) 把二维特征图拉平成向量方便计算全局交集。Dice Loss 的梯度特性是当预测和真实完全不重叠时梯度最大能有效推开模型输出全背景的局部最优。实际训练中我发现 Dice Loss 单独用在小目标上容易震荡一个 batch 里如果恰好没有缺陷目标Dice 会直接变成 0梯度剧烈变化。所以更稳的组合是 Dice Loss 加二分类交叉熵按权重相加比如total_loss 0.5 * dice_loss 0.5 * bce_loss这样既有像素级梯度的稳定又有 Dice 的类别均衡能力。如果你发现缺陷区域特别小比如裂纹类占比不到 2%可以改用 Focal Loss它通过调制因子降低易分样本权重强行把模型注意力拉到困难样本上。权重比例需要实验我从 0.5:0.5 起步观察验证集 IoU 变化再调整。4.3 训练循环、学习率调度与模型保存训练循环本身不复杂但在工程上要注意三个细节学习率策略、验证集评估和最优模型保存。训练脚本里每一步都要能看到当前 loss最好每个 epoch 打印一次训练集平均损失和验证集 IoU而不是等到最后才看结果。下面是一个完整的训练循环import torch from torch.utils.data import DataLoader from tqdm import tqdm def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0 for imgs, masks in tqdm(loader): imgs imgs.to(device) masks masks.to(device) preds model(imgs).squeeze(1) loss criterion(preds, masks.float()) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * imgs.size(0) return total_loss / len(loader.dataset) def validate(model, loader, device): model.eval() iou_sum 0 with torch.no_grad(): for imgs, masks in loader: imgs imgs.to(device) masks masks.to(device) preds torch.sigmoid(model(imgs).squeeze(1)) pred_binary (preds 0.5).float() inter (pred_binary * masks).sum(dim(1, 2)) union ((pred_binary masks) 0).sum(dim(1, 2)) iou_sum (inter / (union 1e-6)).sum().item() return iou_sum / len(loader.dataset) model UNet(in_channels1, num_classes1).to(device) criterion DiceLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-4) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemax, patience5, factor0.5 ) best_iou 0 for epoch in range(50): train_loss train_one_epoch( model, train_loader, criterion, optimizer, device ) val_iou validate(model, val_loader, device) scheduler.step(val_iou) if val_iou best_iou: best_iou val_iou torch.save(model.state_dict(), checkpoints/best_unet_steel.pth) print(fepoch {epoch}: loss{train_loss:.4f}, val_iou{val_iou:.4f})训练循环里preds model(imgs).squeeze(1)是关键一步。因为二分类模型输出是 N×1×H×Wsqueeze(1) 把通道维度去掉变成 N×H×W才能和同样是 N×H×W 的 mask 计算损失。如果忘了 squeezeBCEWithLogitsLoss 会报维度不匹配这个错误在代码调试中高频出现。验证函数里用阈值 0.5 把预测概率转成二值 mask计算每个样本的 IoU 再求平均inter / (union 1e-6)加了一个极小值防止除零。学习率调度器用ReduceLROnPlateau当验证集 IoU 连续 5 个 epoch 不提升时学习率减半这种策略比固定学习率更省心不用人工盯损失曲线反复调 lr。模型保存只存state_dict而不是整个 model 对象这样权重文件小、跨版本加载更稳妥部署时也方便重新搭建网络后再灌入权重。IoU 最高点对应的权重才保存不要用最后一轮的权重因为训练后期模型可能出现轻微过拟合验证集上表现未必是最高点。5. UNet 钢材缺陷训练避坑指南五个高频问题排查记录5.1 张量尺寸不匹配200×200 输入为什么让标准 UNet 崩溃现象把 NEU-DET 原始 200×200 图像直接送进 4.1 节定义的标准 UNet前向传播在解码器拼接处报错提示The size of tensor a (26) must match the size of tensor b (25)之类的维度不匹配。原因标准 UNet 下采样四次加一个 bottleneck共有五次 2 倍池化输入尺寸必须是 2 的 4 次方的整数倍即 16 的倍数。200 除以 16 等于 12.5向下取整后编码器特征和下采样回来后的尺寸对不上跳跃连接拼接时必然报错。解决在数据加载时统一Resize(256, 256)或Resize(192, 192)并在模型代码入口加一个断言函数输入长宽必须能被 16 整除提前拦截预防后续错误。实际项目里我在UNet.forward开头写过固定检查不仅能拦尺寸问题还能在 batch 数据异常时快速定位是数据层还是模型层的问题。5.2 缺陷像素太少模型只输出全背景现象训练 20 个 epoch 后验证集 IoU 停留在 0.05 以下可视化预测结果发现整张图都是黑色模型完全没有预测出任何缺陷区域。原因这是类别不平衡的典型表现。NEU-DET 的裂纹、划伤等细长缺陷在 256×256 图像里只占几百个像素背景占比超过 98%。普通交叉熵损失会把大量梯度分配给背景像素模型发现“预测全背景”已经能让损失很低于是停在局部最优解。解决换用 Dice Loss 或 Focal Loss并按 0.5 权重叠加交叉熵。我实际对比过单用 Dice Loss 在缺陷占比低于 2% 的类别上收敛慢但配合交叉熵后稳定性好很多。另一种思路是损失函数不变改为在采样阶段做在线难例挖掘每个 batch 里强制包含至少一张有小目标的图但实现成本比换损失函数高不推荐第一版就这样做。5.3 标注框到掩码转换后的边缘偏肥现象分割结果和真实缺陷轮廓大致吻合但整体像被充了气比标注框还宽出一圈误报区域集中在缺陷边缘周围几个像素。原因生成 mask 时用了整框填充把框内的背景像素全标成了前景。模型在训练时把这些背景当成正样本学习推理时自然会把缺陷外围的背景像素预测成前景。尤其裂纹、划伤这类细长目标框内背景占比高边缘偏肥非常明显。解决在box_to_mask里加腐蚀操作3×3 核缩边一圈。但要注意标注框边缘紧贴缺陷时腐蚀会把真实缺陷边缘也削掉所以腐蚀后需要抽样可视化对比折中处理。另一种可行方案是用 GrabCut 这类半自动交互式分割方法对框内区域做前景背景分离生成更精细的掩码但工作量较大适合精标注需求高的场景。5.4 增强操作把标签搞歪了现象训练损失正常下降但验证集指标忽高忽低抽几张训练图看增强后的效果发现 mask 里的缺陷区域和原图缺陷错位甚至出现 mask 变形的极端情况。原因数据增强同时作用于图像和 mask如果 mask 用的插值算法和图像不一致比如 mask 用了双线性插值会在边缘产生中间值虽然代码里转成 int 后不影响但多次旋转和翻转叠加放大后导致错位。还有另一种原因某些 transform 默认对 mask 不做处理比如随机裁剪如果只裁了图像而没同步裁 mask两者尺寸直接对不上。解决统一用 albumentations 的 Compose 管理所有增强操作它保证图像和 mask 走同一套变换且 mask 默认用最近邻插值。自己写 transform 时务必在代码里显式指定interpolationcv2.INTER_NEAREST不要省略。增强强度也要控制旋转角度别超过 90 度钢材表面缺陷大多有方向性旋转过度反而影响模型对划伤方向的判断。5.5 显存不足与训练中断现象训练中途报CUDA out of memory或者刚开始第二个 epoch 就崩batch size 调小后能跑但训练速度明显下降。原因UNet 的跳跃连接会把每层编码器特征都留在内存里供解码器拼接使用所以显存占用比同参数量纯分类模型高出不少。256×256 输入、batch size 8、初始通道 64在 8GB 显存上勉强能跑但验证时还要额外算一次模型前向传播如果验证集没放进 no_grad 上下文里内存占用会进一步增加。解决三条路同时走。第一验证阶段用torch.no_grad()包住验证不占梯度内存这是最容易忽略的优化点。第二把输入尺寸降到 192×192分辨率减少约 44%显存占用大幅下降。第三训练时加torch.cuda.amp.autocast()混合精度同时把模型和优化器用GradScaler包装。混合精度对 UNet 这种卷积密集的模型提速明显且对分割精度影响很小是工程里最划算的优化手段。6. 模型推理与结果验证从权重加载到检测效果评估6.1 单张图像推理加载权重、阈值分割与连通域过滤训练完成后的部署验证是从“模型能跑”到“项目能用”之间必须跨过的一步。推理脚本的核心动作是加载权重、单张图前向传播、Sigmoid 转概率、阈值分割、后处理。代码能直观地串起整个流程model UNet(in_channels1, num_classes1).to(device) model.load_state_dict(torch.load(checkpoints/best_unet_steel.pth, map_locationdevice)) model.eval() img cv2.imread(data/images/demo.jpg, cv2.IMREAD_GRAYSCALE) img cv2.resize(img, (256, 256)) x torch.from_numpy(img.astype(np.float32) / 255.0).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): logits model(x) prob torch.sigmoid(logits).squeeze(0).squeeze(0).cpu().numpy() pred_mask (prob 0.5).astype(np.uint8) n_labels, labels, stats, _ cv2.connectedComponentsWithStats(pred_mask, connectivity8) min_area 10 for i in range(1, n_labels): if stats[i, cv2.CC_STAT_AREA] min_area: pred_mask[labels i] 0权重加载时用map_locationdevice是为了兼容在 GPU 上训练、在 CPU 上推理的环境也兼容 CPU 上训练后加载到 GPU 的情况。model.eval()必须调用它会关闭 Dropout 和 BatchNorm 的训练模式统计否则推理结果会出现莫名其妙的波动。阈值 0.5 是默认值如果验证集上发现误报多可以提高到 0.6 或 0.7这一步会直接影响最终精度和召回。连通域过滤把面积小于 10 像素的孤立预测点去掉这类小点往往是噪声而不是真实缺陷。后处理后用叠加方式把预测轮廓画在原图上既方便人工质检也能快速判断模型对新样本的泛化能力。6.2 量化检测效果IoU 与 Dice 指标的计算方法指标不能只看损失曲线分割任务的标准评估指标是 IoU 和 Dice。IoU 的计算是预测二值图与真实 mask 的交集面积除以并集面积取值范围 0 到 1越接近 1 越好。Dice 是 2 倍交集除以预测面积加真实面积数值上通常略高于 IoU。项目里我一般同时算这两个并分别统计六类缺陷各自的 IoU而不是只在全局算一个平均值。因为全局 IoU 会被占比较高的类别拉高掩盖划伤这类小目标效果差的问题。验证脚本里validate函数已经按样本计算了平均 IoU如果要做类别细分则需要在验证时带上每个样本的类别标签按缺陷类别分组统计。实际操作可以在循环里维护一个字典键是类别名值是 IoU 列表最后统一取均值。评估时保持验证集固定不变不要每次跑都用新随机种子否则不同模型版本之间的对比就没有意义。我用过一段经验把验证集的每张图和预测结果保存成 PNG 拼接图按类别归类摆放方便肉眼对比模型在不同缺陷类型上的表现比只看数字直观得多。6.3 把模型用在实际产线前的三个调整从验证集表现不错到真正落地产线中间还差几步工程化工作。第一个调整是推理速度。UNet 前向传播在 GPU 上处理 256×256 图像大约需要 20 到 50 毫秒如果产线相机帧率是每秒 10 张需要确认单张推理耗时在上限内必要时把输入降到 192×192 或者用 TensorRT 做加速。第二个调整是阈值选择。验证集上试多个阈值画 P-R 曲线挑出误报和漏报平衡点而不是死守 0.5。钢材表面缺陷分类漏检一块缺陷的损失远大于误报一次阈值可以适当降低让模型偏保守。第三个调整是数据分布漂移。NEU-DET 的每类缺陷形态相对固定但实际产线可能遇到不同钢种、不同光照、不同氧化程度用公开数据集训练出的模型往往在新产线数据上效果打折。常见做法是冷启动阶段先用这个模型跑自动标注把产线采集的图片用模型预测结果辅助人工标注积累一批真实数据后做增量训练微调。这个流程能让模型快速适应现场数据同时也是持续迭代的正确路径。我在多个项目里养成了一个习惯训练结束第一时间打印输入输出张量的 shape、归一化参数和阈值连同最优权重一起存档。这样即使隔了三周回来调模型也能快速回忆当时的配置不至于对着一个 weight 文件发愁。模型推理这个环节看着简单翻车都翻在这些不起眼的细节上。希望我上面记录的数据格式、损失选择、尺寸处理和推理后处理这些经验能帮你在类似的项目里少走一段弯路。本文还有配套的精品资源点击获取
