简介：面向工业质检领域深度学习入门与实战的钢材表面缺陷检测项目，基于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, axis=0)或torch.unsqueeze(0)都可以。如果有人图省事用三通道方式读灰度图，等于把同一个灰度值复制三遍，参数量没变但输入体积大了三倍，没必要。

归一化这里有两种常见路线。一种是直接除以 255，把像素映射到 0 到 1 区间，再配合 BatchNorm，训练也能稳定收敛，这是最快上手的做法。另一种是用均值 0.5、方差 0.5 做标准化，把像素映射到 -1 到 1，数值分布更对称，梯度更新更平稳，代价是推理时也要套同一套参数。我项目里一般会用后者，因为实际感受是同样学习率下，标准化后损失曲线波动明显更小，尤其当 batch size 很小的时候，BatchNorm 统计不稳定带来的干扰会被标准化抵消一部分。

标签编码是另一个容易埋雷的地方。二分类分割的 mask 只有两个值，背景是 0，前景缺陷是 1，数据类型用 int64（PyTorch 的 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.10，PyTorch 版本用 1.10 以上，同时安装和 PyTorch 严格对应的 torchvision，不要混装不同版本。图像读取用 opencv-python，数据增强用 albumentations，进度展示用 tqdm。没有 GPU 也能跑，只是训练速度会慢三四倍，NEU-DET 这个数据量 CPU 训练一轮一百多张图大概十几分钟，跑 30 个 epoch 也能接受，但每 epoch 时间较长，调参效率很低，有条件还是建议用 GPU。

项目目录我一般按下面这个结构组织，数据、代码、模型权重分开，避免把训练产物混进源码目录：

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_idx=1, erode_kernel=3): mask = np.zeros(image_shape, dtype=np.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, iterations=1) return mask

box_to_mask接收图像尺寸、边界框和类别索引，返回和原图等大的 mask。class_idx=1表示前景类别编号，二分类场景固定用 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_size=256, augment=False): 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(p=0.5), VerticalFlip(p=0.5), Rotate(limit=90, p=0.3), Resize(img_size, img_size, interpolation=cv2.INTER_LINEAR) ]) else: self.transform = Compose([ Resize(img_size, img_size, interpolation=cv2.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, dtype=np.uint8) for obj in objs: xmin, ymin, xmax, ymax = obj['bbox'] mask[ymin:ymax, xmin:xmax] = 1 transformed = self.transform(image=img, mask=mask) 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 到 1，unsqueeze(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, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, in_channels=1, num_classes=1): 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, stride=2) self.dec4 = DoubleConv(1024, 512) self.up3 = nn.ConvTranspose2d(512, 256, 2, stride=2) self.dec3 = DoubleConv(512, 256) self.up2 = nn.ConvTranspose2d(256, 128, 2, stride=2) self.dec2 = DoubleConv(256, 128) self.up1 = nn.ConvTranspose2d(128, 64, 2, stride=2) 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], dim=1)) d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1)) d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1)) d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1)) 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, smooth=1.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_channels=1, num_classes=1).to(device) criterion = DiceLoss() optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='max', patience=5, factor=0.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(f'epoch {epoch}: loss={train_loss:.4f}, val_iou={val_iou:.4f}')

训练循环里preds = model(imgs).squeeze(1)是关键一步。因为二分类模型输出是 N×1×H×W，squeeze(1) 把通道维度去掉，变成 N×H×W，才能和同样是 N×H×W 的 mask 计算损失。如果忘了 squeeze，BCEWithLogitsLoss 会报维度不匹配，这个错误在代码调试中高频出现。验证函数里用阈值 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 时，务必在代码里显式指定interpolation=cv2.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_channels=1, num_classes=1).to(device) model.load_state_dict(torch.load('checkpoints/best_unet_steel.pth', map_location=device)) 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, connectivity=8) 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_location=device是为了兼容在 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 文件发愁。模型推理这个环节看着简单，翻车都翻在这些不起眼的细节上。希望我上面记录的数据格式、损失选择、尺寸处理和推理后处理这些经验，能帮你在类似的项目里少走一段弯路。

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