做工业视觉的同学应该都听过 NEU-DET 这个数据集。东北大学收集的热轧带钢表面六类常见缺陷总共 1800 张带标注图片尺寸统一 200×200类别覆盖裂纹、夹杂、斑块、麻点、氧化铁皮压入和划伤。这几乎是钢材表面缺陷检测领域最经典的入门门槛也是很多人接触工业缺陷检测的第一个数据集。而 YOLOv5s 是这个数据集上最省事的模型之一模型小单卡就能训练推理速度快从训练到模型部署的整条链路在社区里已经被踩得很熟。这篇教程就把完整流程过一遍包括 NEU-DET 的数据集格式转换、YOLOv5s 环境配置、训练与评估、最后把 .pt 模型导出成 ONNX 并封装成 HTTP 服务全程给出可以直接复制的命令和代码。说是保姆级是因为很多细节我不会跳过去类别顺序为什么不能乱、XML 转 txt 那一步为什么容易翻车、训练参数是怎么根据显存算出来的、部署时为什么框会偏移。这些坑单独看都不大连在一起却能卡住新手好几天。下面开始。1. 为什么选 YOLOv5s NEU-DET先搞清这套组合适合谁1.1 六类缺陷到底长什么样NEU-DET 里的六个类别分别是crazing裂纹表面细小的龟裂纹理灰度很低和背景容易混淆inclusion夹杂嵌入表面的颗粒状异物通常是暗色小块patches斑块大面积灰度不均区域边界模糊pitted_surface麻点表面密集的小凹坑类似点状噪声rolled-in_scale氧化铁皮压入铁皮在轧制时被压入表面呈条状或块状scratches划伤细长的高亮或暗色线条方向随机这六类缺陷的形态差异非常大。crazing 和 patches 属于纹理型缺陷细节少、边界模糊对模型的特征提取能力要求高scratches 和 rolled-in_scale 又是典型的细长目标长宽比很极端。NEU-DET 虽然只有 1800 张图但实际训练时会发现它的难点不在目标数量多而在某些缺陷和背景太像模型需要很强的上下文感知才能把它们找出来。1.2 YOLOv5s 在这条链路上的定位YOLOv5s 是 YOLOv5 系列里最小的模型权重文件只有十几 MB参数量在 YOLOv5 系列里最小。放在钢材表面缺陷检测这个场景里它的价值不是刷榜而是够用、快、好部署。工业质检项目对推理速度的要求往往比学术比赛高得多。一条产线的相机可能每秒出几十帧如果模型太大GPU 成本直接翻好几倍。YOLOv5s 在普通 1080Ti 上跑 640 分辨率单张推理大概 10ms 左右放到工控机上也能稳住实时性。而且 YOLOv5 的仓库生态非常完整训练、验证、导出 ONNX、OpenVINO、TensorRT 都有现成脚本不用自己造轮子。相比新出的各种检测框架YOLOv5 的项目资料也更丰富。遇到问题搜索一下基本都能找到答案。如果你不是做算法研究而是要给实际产线提供一套缺陷检测能力YOLOv5s NEU-DET 这个组合是性价比很高的选择。1.3 这篇教程覆盖的边界先说清楚适用前提默认你有 NVIDIA GPU6GB 显存以上的显卡都能跑。只有 CPU 的话训练 150 个 epoch 会非常煎熬建议直接用云 GPU 或者先把教程流程跑通再考虑本地训练。最终你会得到这些东西转好格式的 YOLO 数据集、一份 data.yaml、训练产生的 best.pt、导出后的 best.onnx以及一个基于 FastAPI 的 HTTP 检测接口。手机或者现场工位可以通过局域网调用这个接口上传图片做检测。至于 GUI 界面、PLC 联动、多相机并发这些产线集成内容不在本文范围但接口层做扩展是足够的。2. 数据集转换与环境安装NEU-DET 从 XML 到 YOLO txt 的完整处理2.1 环境依赖锁定建议用 Anaconda 创建独立环境避免把系统 Python 搞乱。我的固定组合是 Python 3.9 PyTorch 1.13 CUDA 11.7。命令如下conda create -n yolov5 python3.9 -y conda activate yolov5 pip install torch1.13.1cu117 torchvision0.14.1cu117 --extra-index-url https://download.pytorch.org/whl/cu117 git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt为什么锁 Python 3.9因为 YOLOv5 的旧版本在 Python 3.10/3.11 上OpenCV 和 NumPy 容易出现二进制兼容问题有些依赖库的预编译轮子只支持到 3.9。如果你用的 YOLOv5 版本比较新Python 3.10 也可以但没必要在环境上增加不确定性。安装完成后先跑一句python train.py --help能正常弹出参数说明说明依赖基本没问题。这时候不要急着训练先把数据集处理好。2.2 NEU-DET 的目录与标注格式下载后的 NEU-DET 一般是这样的目录结构NEU-DET/ IMAGES/ IMG001.jpg ... IMG1800.jpg ANNOTATIONS/ IMG001.xml ... IMG1800.xmlXML 格式跟 PASCAL VOC 很接近核心内容是这样annotation object namecrazing/name bndbox xmin30/xmin ymin40/ymin xmax160/xmax ymax180/ymax /bndbox /object /annotation图片文件名和 XML 文件名一一对应这一点处理起来很省事。但要注意有些网盘分享的版本目录结构可能和这里不完全一样先自己核对一下再写脚本。2.3 一个脚本把 VOC XML 转成 YOLO txtYOLO 训练用的是 txt 标注每行格式是class_id x_center y_center width height坐标全部是相对图片宽高的比例值范围在 0 到 1 之间。转换脚本如下import os import xml.etree.ElementTree as ET classes [crazing, inclusion, patches, pitted_surface, rolled-in_scale, scratches] image_dir NEU-DET/IMAGES xml_dir NEU-DET/ANNOTATIONS out_dir labels os.makedirs(out_dir, exist_okTrue) for xml_file in os.listdir(xml_dir): if not xml_file.endswith(.xml): continue tree ET.parse(os.path.join(xml_dir, xml_file)) root tree.getroot() img_name os.path.splitext(xml_file)[0] .jpg img_w int(root.find(size/width).text) img_h int(root.find(size/height).text) with open(os.path.join(out_dir, os.path.splitext(xml_file)[0] .txt), w) as f: for obj in root.findall(object): cls_name obj.find(name).text if cls_name not in classes: continue cls_id classes.index(cls_name) bbox obj.find(bndbox) x1 int(bbox.find(xmin).text) y1 int(bbox.find(ymin).text) x2 int(bbox.find(xmax).text) y2 int(bbox.find(ymax).text) x_center (x1 x2) / 2.0 / img_w y_center (y1 y2) / 2.0 / img_h w (x2 - x1) / img_w h (y2 - y1) / img_h x_center min(max(x_center, 0), 1) y_center min(max(y_center, 0), 1) w min(max(w, 0), 1) h min(max(h, 0), 1) f.write(f{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n)脚本里为什么最后要做一次 clip因为 NEU-DET 的标注框偶尔会超出图片边界比如 xmax 标成 201而图片宽度是 200。如果不处理YOLO 在计算锚框匹配时可能出现异常轻则警告重则导致部分目标参与不到训练。这个细节能挡掉很多莫名其妙的问题。2.4 数据划分与 data.yaml原始数据集没有官方划分必须自己分。我的习惯是 80% 训练、10% 验证、10% 测试固定随机种子保证可复现import os, random, shutil img_names os.listdir(NEU-DET/IMAGES) random.seed(42) random.shuffle(img_names) n len(img_names) train_names img_names[: int(n * 0.8)] val_names img_names[int(n * 0.8) : int(n * 0.9)] test_names img_names[int(n * 0.9) :] os.makedirs(images/train, exist_okTrue) os.makedirs(images/val, exist_okTrue) os.makedirs(images/test, exist_okTrue) os.makedirs(labels/train, exist_okTrue) os.makedirs(labels/val, exist_okTrue) os.makedirs(labels/test, exist_okTrue) for name in train_names: shutil.copy(fNEU-DET/IMAGES/{name}, images/train/ name) shutil.copy(flabels/{name.replace(.jpg, .txt)}, labels/train/ name.replace(.jpg, .txt)) for name in val_names: shutil.copy(fNEU-DET/IMAGES/{name}, images/val/ name) shutil.copy(flabels/{name.replace(.jpg, .txt)}, labels/val/ name.replace(.jpg, .txt)) for name in test_names: shutil.copy(fNEU-DET/IMAGES/{name}, images/test/ name) shutil.copy(flabels/{name.replace(.jpg, .txt)}, labels/test/ name.replace(.jpg, .txt))固定随机种子这件事很多人不在意但 NEU-DET 样本少不同划分带来的 mAP 波动可能有 2 到 3 个百分点甚至更高。如果每次跑结果都不一样你很难判断是模型改进了还是划分变了。接着在数据集根目录放一个 NEU-DET.yamlpath: dataset/NEU-DET train: images/train val: images/val test: images/test nc: 6 names: [crazing, inclusion, patches, pitted_surface, rolled-in_scale, scratches]这里必须再次强调names 的顺序一定要和转换脚本里的 classes 完全一致。类别名一样不代表顺序一样一旦顺序错位模型训练完所有类别都是错的而且不会报任何错误。3. 训练参数与 loss 曲线一个能复现的 YOLOv5s 训练过程3.1 预训练权重与超参数NEU-DET 只有 1800 张图从零训练效果会差很多。直接用 COCO 预训练权重做迁移学习让模型先用通用特征把网络权重初始化好再在钢材缺陷数据上微调收敛速度和最终精度都会好很多。启动训练的命令python train.py --img 640 --batch 16 --epochs 150 --data NEU-DET.yaml --weights yolov5s.pt --cache --device 0这几个参数值得展开讲。--img 640原始图片是 200×200训练时 YOLOv5 会先做 letterbox 缩放。640 不是越大越好但在这个缺陷目标偏小的数据集上高分辨率能明显提升召回率。6GB 显存跑 640 和 batch 16 会有点紧张可以降成 416或者把 batch 降到 8。--batch 16根据显存动态调整。8GB 显卡跑 640 batch 16 可能刚好卡在边界如果 OOM优先降到 8不要硬刚。--epochs 150小数据集 150 轮足够。跑太多轮验证集 loss 很容易回升反而过拟合。--cache把图片提前缓存到内存训练速度会快不少。缺点是吃内存机器内存小于 32GB 建议慎用。3.2 训练过程中的输出怎么看训练开始后终端会打印每个 epoch 的 box_loss、obj_loss、cls_loss、precision、recall、mAP0.5。我一般只看三个信号第一loss 是否稳定下降。前 20 轮如果 loss 几乎不动大概率是数据格式有问题先去检查 labels 是不是空文件或者类别下标是不是越界。第二val 指标是否持续上升。NEU-DET 不是复杂数据集正常 50 轮左右 mAP0.5 就能到 0.8 以上100 轮之后逐渐收敛mAP0.5 一般能到 0.85 以上。具体数字取决于数据划分和增强参数。第三val loss 是否出现明显回升。如果 train loss 还在下降而 val loss 开始反弹说明过拟合了。可以用 --patience 30 让训练在验证指标连续 30 轮不提升时自动早停。训练结束后runs/train/exp 目录下会有 weights/best.pt 和 weights/last.pt。best.pt 是按验证集 mAP 保存的最优权重后面评估和导出都用它。3.3 数据增强开关要不要动YOLOv5 默认开启 Mosaic、HSV 扰动、随机翻转等增强。对于 NEU-DET默认配置基本够用不需要专门关掉。有一个参数可以留意--hsv-v这是亮度随机扰动的范围。钢材表面缺陷在不同光照条件下差异很大适当保留亮度扰动有助于提升泛化能力。默认值 0.4 已经不错不用额外调大。如果某个类别召回率特别低先别急着改模型参数回去看看标注质量。NEU-DET 本身也有标注噪声比如 rolled-in_scale 的框经常会比实际缺陷大一圈。这类问题靠调超参解决不了只能重新清洗标注。4. 在部署前先学会看评估结果mAP 之外还有多少信息4.1 用 val.py 拿到每个类别的 AP训练完先用官方验证脚本跑一遍python val.py --weights runs/train/exp/weights/best.pt --data NEU-DET.yaml --img 640 --conf-thres 0.25 --iou-thres 0.5 --task val输出会给出每个类别的 AP50、AP50-95 和整体 mAP。NEU-DET 六类里crazing 的 AP 通常最低因为裂纹纹理和背景灰度太接近inclusion 和 scratches 通常最高边界相对清晰。这里提醒一句工业项目里不要只看整体 mAP。不同缺陷的后果不一样可能你只关心麻点和划伤那就算整体 mAP 很高如果麻点 AP 不达标模型照样不能用。按类别拆开看才能对应到产线质检标准上。4.2 混淆矩阵和 PR 曲线YOLOv5 会在 runs/train/exp 目录下生成 confusion_matrix.png 和 PR_curve.png。混淆矩阵里如果 crazing 大量被识别成背景说明这个类别的纹理特征确实弱如果 patches 被识别成 crazing说明两个类别的灰度分布重叠严重。看到这种情况常用的手段有三种一是给对应类别增加数据增强模拟更多纹理变化二是复制该类样本做简单过采样三是换更大的模型比如 yolov5m 或 yolov5l。第三种最直接但要评估部署端的算力能不能接受。我在一个项目里就碰到过 rolled-in_scale 的 AP50 只有 0.85其他类别都在 0.95 以上。排查后发现是标注框本身偏大模型学到的框比 GT 大导致 IoU 偏低。这种问题属于标注层面不是模型层面。4.3 评估阶段一定要看图片指标是辅助最终要回到图片上看预测结果。跑一下python detect.py --weights runs/train/exp/weights/best.pt --source test_images --conf-thres 0.25把几张典型图片的检测结果打印出来重点看漏检和误检。有些小缺陷在指标上不明显但人眼一眼就能看出模型把背景纹理当成了缺陷。这一步虽然原始却是部署前最可靠的质量检查。5. .pt 模型到在线服务本地推理、API 封装与边缘设备迁移5.1 导出成 ONNX为什么值得做训练完的 .pt 文件直接部署也可以但 PyTorch 推理环境很重生产环境不想为每个服务都装一套 torch。ONNX 是跨平台、跨语言的模型格式ONNX Runtime 相比直接跑 PyTorch 推理通常还有一定加速。导出命令python export.py --weights runs/train/exp/weights/best.pt --include onnx --opset 12 --simplify如果提示缺少 onnx-simplifier先执行pip install onnx onnxruntime onnx-simplifier导出的 ONNX 输出是原始预测张量形状通常是 [1, 25200, 12]其中 12 等于 4 个坐标 1 个 objectness 6 个类别置信度。25200 是三个尺度特征图上的候选框总数。到这一步YOLOv5 仓库还依赖着但真正部署时通常不能依赖整个仓库需要自己做预处理、解码和 NMS。5.2 ONNX Runtime 推理的最小实现下面是一份可以直接用的 ONNX 推理脚本包含 letterbox 预处理、解码和 NMSimport cv2 import numpy as np import onnxruntime as ort CLASSES [crazing, inclusion, patches, pitted_surface, rolled-in_scale, scratches] CONF_THRESH 0.25 IOU_THRESH 0.45 def letterbox(img, new_shape(640, 640), color(114, 114, 114)): shape img.shape[:2] r min(new_shape[0] / shape[0], new_shape[1] / shape[1]) new_unpad (int(round(shape[1] * r)), int(round(shape[0] * r))) dw, dh new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] dw, dh dw // 2, dh // 2 img cv2.resize(img, new_unpad, interpolationcv2.INTER_LINEAR) top, bottom int(round(dh - 0.1)), int(round(dh 0.1)) left, right int(round(dw - 0.1)), int(round(dw 0.1)) img cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, valuecolor) return img, r, (dw, dh) def nms(boxes, scores, iou_thres): indices cv2.dnn.NMSBoxes( [[float(b[0]), float(b[1]), float(b[2] - b[0]), float(b[3] - b[1])] for b in boxes], [float(s) for s in scores], score_thresholdCONF_THRESH, nms_thresholdiou_thres, ) return np.array(indices).flatten() if len(indices) 0 else [] sess ort.InferenceSession(best.onnx, providers[CPUExecutionProvider]) def detect(img0): img, ratio, pad letterbox(img0, (640, 640)) img img[:, :, ::-1].transpose(2, 0, 1) # BGR-RGB, HWC-CHW img np.ascontiguousarray(img).astype(np.float32) / 255.0 img np.expand_dims(img, 0) out sess.run(None, {sess.get_inputs()[0].name: img})[0] preds out[0] boxes, scores, cls_ids [], [], [] for p in preds: obj_conf p[4] if obj_conf CONF_THRESH: continue cls_conf p[5:] cls_id int(np.argmax(cls_conf)) score obj_conf * cls_conf[cls_id] if score CONF_THRESH: continue x1, y1, x2, y2 p[:4] x1 (x1 - pad[0]) / ratio y1 (y1 - pad[1]) / ratio x2 (x2 - pad[0]) / ratio y2 (y2 - pad[1]) / ratio boxes.append([x1, y1, x2, y2]) scores.append(score) cls_ids.append(cls_id) if len(boxes) 0: return [] keep nms(boxes, scores, IOU_THRESH) results [] for i in keep: results.append({bbox: boxes[i], score: scores[i], class: CLASSES[cls_ids[i]]}) return results img cv2.imread(test.jpg) print(detect(img))这段代码里最容易错的是 letterbox 还原。预测框是在缩放后、加了 padding 的图上算出来的回原图时一定要先减 padding 再除以缩放比。只减了 padding 却忘了除以 ratio框的位置会整体向右下偏移只除 ratio 却忘了 padding框会更靠近左上。这个错位问题在部署现场出现频率极高。如果你想把推理也加速可以在 onnxruntime 的 providers 里加上 CUDAExecutionProvider。注意要安装 onnxruntime-gpusess ort.InferenceSession(best.onnx, providers[CUDAExecutionProvider, CPUExecutionProvider])5.3 用 FastAPI 把模型包成 HTTP 服务本地脚本只能自己调要给现场或手机端调用最简单的方式是包一个 HTTP 接口。用 FastAPI 很轻量pip install fastapi uvicorn python-multipart接口代码from fastapi import FastAPI, UploadFile, File import numpy as np import cv2 from detect import detect app FastAPI() app.post(/detect) async def detect_image(file: UploadFile File(...)): data await file.read() img cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR) results detect(img) return {count: len(results), results: results} if __name__ __main__: import uvicorn uvicorn.run(app, host0.0.0.0, port8000)启动后同一局域网内的手机浏览器访问 http://电脑IP:8000/docs就能在 Swagger 页面上传图片测试。很多人问手机怎么调用电脑上部署的模型本质上就是手机发 HTTP 请求模型在电脑上跑完返回检测框坐标和置信度。工业现场接相机也一样把上传图片改成读取相机帧即可。生产环境使用时还需要加访问控制、图片大小限制、请求队列这些工程化处理但算法层的链路到这里已经完整了。5.4 边缘设备迁移TensorRT 和 OpenVINO 怎么选如果要部署到工业 PC 或 Jetson Orin Nano 这类边缘设备建议把 ONNX 再转成更高效的格式。Intel CPU 为主的工控机优先用 OpenVINO。YOLOv5 官方导出脚本支持python export.py --weights best.pt --include openvinoOpenVINO 在 Intel CPU 上通常比 ONNX Runtime CPU 快一到两倍而且部署包不依赖 PyTorch。NVIDIA GPU 设备包括 Jetson 系列可以转 TensorRT。用 trtexec 命令trtexec --onnxbest.onnx --saveEnginebest.engine --fp16TensorRT 加速效果非常明显但版本适配坑比较多opset 版本、ONNX 算子版本都可能导致转换失败。如果只是学习用 ONNX Runtime 就够了没必要一开始就陷进 TensorRT 的适配问题。部署层的核心原则只有一个预处理必须和训练完全一致。分辨率、padding、RGB/BGR、归一化方式任何一步不一致精度都会掉。后面优化做得再好也补不回预处理不一致带来的损失。6. 实战中踩过的坑与补救方案6.1 标签错位训练完所有类别 AP 都是 0我第一次带人跑 NEU-DET有人把 data.yaml 的 names 写成了和其他脚本不同的顺序但转换脚本里的 classes 是另一套顺序。模型照样能训练loss 也会下降但最终验证结果显示每个类别的 AP 都接近 0。这个问题的排查链路是先看 val 日志发现 precision 和 recall 一直很低再看验证集预测图片发现预测框类别和实际物体对不上最后逐行对比 data.yaml 和转换脚本才定位到类别顺序错位。这种坑最难发现的地方在于它不会报错只会让结果错得非常有规律。所以训练前一定先打印几个标签文件看一眼确认类 ID 和实际目标对应。6.2 显存溢出OOM 之后的正确做法训练到一半提示 CUDA out of memory最直接的办法是降低 batch size。YOLOv5 的 train.py 没有像一些库那样提供显式的梯度累积参数所以实践中最省事的方案就是先降 batch再降 img最后考虑关掉 --cache。原因是 --cache 会把图片缓存到内存虽然不直接占显存但会增加内存压力内存不足时同样会影响训练稳定性。我自己的经验是8GB 显存跑 640 分辨率batch 用 8 基本稳定优先从 batch 16 降到 batch 8而不是直接降到 4因为 batch 太小会带来明显的梯度噪声影响模型收敛。6.3 部署时框偏移letterbox 参数没还原ONNX 推理脚本很多人是复制来的复制完只改了模型路径没改 letterbox 的还原逻辑。如果训练时用 640部署时输入的是 1920×1080就必须在预处理代码里同步计算缩放比和 padding并在后处理时把坐标映射回原图。排查这个问题的快速方法是拿一张已知目标位置的图先打印 letterbox 返回的 ratio 和 pad再对比检测框和真实目标位置。如果偏移量固定说明还原公式有问题如果偏移量和目标大小有关说明缩放比没写对。6.4 验证集划分不稳定不同 run 之间 mAP 波动大NEU-DET 样本只有 1800 张随机划分的影响很大。同参数训练两次mAP 差三五个百分点是正常的不要一上来就怀疑模型不稳定。复盘实验或者写报告时一定要固定随机种子并且把数据划分结果保存下来。我习惯把划分后的 images 和 labels 目录直接纳入 git 管理这样每次训练用的验证集都一样指标对比才有效。最后分享一个我自己的交付经验如果是给产线做项目不要只交一个 .pt 或者 .onnx 文件。一定要连同 data.yaml、类别顺序、预处理代码、后处理代码、验证集划分说明一起交付。很多现场调试搞到最后都不是模型不行而是部署端对“输入图片应该怎么预处理”没有统一约定。把这条链路彻底固定下来YOLOv5s NEU-DET 这套流程才算真正跑通了。