自己训练一个 ONNX 检测模型,比你想象中简单得多
从标注到毫秒级推理——一个 YOLOv8 ONNX 的实战全记录附完整工具链写在前面为什么我没有直接用多模态大模型最近要做一个目标检测需求从图片里把「手机」和「吸烟」识别出来。我身边不少做视觉处理的朋友一上来就接一个多模态大模型GPT-4V、Qwen-VL、Gemini 之类理由是「不用训练发张图直接问就完事了」。说实话这确实是当下最省事的做法——把图丢进去用自然语言描述你想要什么模型就能给你答案。但我想说一句大实话对于固定场景、固定目标的重复识别自己训练一个轻量 ONNX 模型往往比直接上多模态大模型更划算。这篇文章就是我用一晚上把这件事跑通的完整记录从「连数据都没有」到「网页上传图片实时出框」全部基于 YOLOv8 ONNX。我会把每一步用到的工具、踩过的坑、最终的模型规格都写清楚照着做你也能复现。大家可以看一下最终效果图任务基本上就完成了一、你到底需要哪些工具工具链全景很多人卡在「不知道从哪下手」本质是没理清这条流水线。其实就三段每段都有成熟工具1. 打标注——给图片贴标签模型不会凭空学得先告诉它「这张图里哪是手机、哪是烟」。这一步叫标注。常用工具工具适合谁特点LabelImg个人、小规模本地运行、最轻量YOLO 模式直接导出.txt标签文件Label Studio团队协作网页端、支持多模态、可多人同时标CVAT专业标注功能强、支持视频、多边形、追踪Roboflow想省事到底在线标注 自动划分训练/验证集 数据增强一条龙我这次数据集已经就绪dataset_yolo/下 3074 张图、3074 个标签结构是标准的 YOLO 格式dataset_yolo/ ├── data.yaml # 类别与路径配置 ├── images/ │ ├── train/ # 2767 张训练图 │ └── val/ # 307 张验证图 └── labels/ ├── train/ # 与图片一一对应的 .txt └── val/每个.txt一行一个目标格式是class_id cx cy w h归一化到 0~1 的中心点坐标。比如1 0.4 0.7 0.1 0.2表示「类别1(smoke)中心在图的 40%宽、70%高框占图宽 10%、高 20%」。,2. 训练框架——把数据变成权重Ultralytics (YOLOv8)一行 API 完成训练、验证、导出不用自己写训练循环。PyTorch CUDA本地 GPU 加速。我用的 GTX 16504GB 显存足够跑yolov8n这种轻量模型。onnx / onnxslim把训练好的.pt导出成通用格式的.onnx顺便做图优化slim。3. 部署与推理——让模型跑起来onnxruntime真正干活的推理引擎支持 CPU 和 GPU毫秒级响应。Flask几行代码把模型包成一个网页服务方便演示和接入业务。二、环境准备一次性安装# 建议用虚拟环境隔离 python -m venv .venv .venv\Scripts\activate pip install torch torchvision ultralytics onnx onnxslim flask # 如果要用 GPU 加速推理再装 pip install onnxruntime-gpu # 否则默认装 onnxruntime(CPU版) 也够用我实测的环境版本torch带 CUDA、ultralytics 8.4.124、onnx 1.22.0、onnxslim 0.1.96、flask 3.1.3。一个坑先说在前面本机装的 onnxruntime 是CPU 版没有 CUDAExecutionProvider。推理时如果让框架自动选 GPU会报CUDAExecutionProvider加载失败。解决方法是推理时显式devicecpu。后面细说。三、训练几行代码的事data.yaml是整个流程的地图指定类别和数据集位置path: D:/PythonProject/PythonProject2/dataset_yolo train: images/train val: images/val nc: 2 names: [phone, smoke]训练脚本核心就这么几行from ultralytics import YOLO # 从官方预训练权重出发它自带 COCO 80 类的基础能力我们在此之上微调 model YOLO(yolov8n.pt) # 训练50 轮输入 640batch 84GB 显存稳妥 model.train( datadata.yaml, epochs50, imgsz640, batch8, device0, # 用 GPU若报错可改成 cpu projectruns, nameyolo_phone_smoke ) # 训练完直接导出 ONNXsimplify 让图更干净 model.export(formatonnx, simplifyTrue, imgsz640)跑一句话启动后台跑日志落盘python train_yolov8.py --epochs 50 runs_train.log 21我的机器上GTX 1650实测约77 分钟跑完 50 轮显存占用约 2GB。训练指标如下指标数值precision精确率0.821recall召回率0.744mAP0.50.825mAP0.5:0.950.436mAP0.50.825 说明「框得对」的程度不错mAP0.5:0.95更严格的 IoU 平均0.436 偏低对 yolov8n 仅 50 轮 类别不平衡smoke 占绝大多数属于正常 baseline。想再提精度可以加 epoch、开数据增强或做采样均衡。踩坑实录都是真实发生的路径错了。原data.yaml指向旧桌面路径训练直接报找不到数据集。改成实际路径D:/PythonProject/PythonProject2/dataset_yolo即可。脚本多拼了一层目录。训练脚本里BASE已经是 dataset_yolo 根目录结果又拼了一次dataset_yolo/data.yaml导致dataset_yolo/dataset_yolo/data.yaml不存在。把DATA/PROJECT的拼接改对就过了。后台进程联网会挂死。Ultralytics 默认会做遥测/版本检查后台任务没网就卡住。先把yolov8n.pt下载到本地绝对路径并在训练里关掉联网相关检查才顺利跑完。ONNX 推理要锁 CPU。如前所述强制devicecpu解决。四、导出与校验 ONNX训练脚本跑完会自动生成runs/yolo_phone_smoke/weights/best.onnx12.3 MB。但「能导出」不等于「能用」我做了三步硬校验import onnx, onnxruntime as ort, numpy as np from PIL import Image p runs/yolo_phone_smoke/weights/best.onnx m onnx.load(p); onnx.checker.check_model(m) # 1. 图结构合法 sess ort.InferenceSession(p, providers[CPUExecutionProvider]) x (np.asarray(Image.open(test.jpg).convert(RGB).resize((640,640))).astype(float32)/255.0).transpose(2,0,1)[None] y sess.run([output0], {images: x})[0] # 2. 真能跑 print(y.shape) # 3. 看输出维度 - (1, 6, 8400)模型的真实规格给同事对接用输入张量名images形状[1, 3, 640, 640]NCHWfloat32需要把图缩放到 640×640 并除以 255 归一化。输出张量名output0形状[1, 6, 8400]每列是[cx, cy, w, h, score_phone, score_smoke]。分数已经是 0~1 概率导出时已 sigmoid无需再激活但 8400 个候选框需自己做 NMS 去重并把坐标从 640 空间还原到原图尺寸。五、做成网页 Demo端口 8096为了让非技术同事也能用我用 Flask 包了一个极简页面上传图片 → 后端加载best.onnx推理 → 把画好框的图返回。前后端同一个端口 8096Flask 既托管 HTML 又提供/upload接口。核心逻辑from flask import Flask, request, jsonify from ultralytics import YOLO import base64, io from PIL import Image app Flask(__name__) model YOLO(runs/yolo_phone_smoke/weights/best.onnx) # 加载 ONNX app.route(/upload, methods[POST]) def upload(): img Image.open(request.files[image]) res model.predict(sourceimg, conf0.25, imgsz640, devicecpu, verboseFalse) # 把 res[0].plot() 生成带框图转 base64 返给前端 ...实测一张验证图检测出smoke (0.87)CPU 推理约0.08 秒/张。六、进阶把「原版模型」和「自训模型」一起用YOLOv8 官方yolov8n.pt自带 COCO 80 类人、车、动物、日常物我训的best.onnx只有 phone/smoke。两者是独立的权重文件可以一起加载、各跑一遍、结果合并。更聪明的玩法是级联先用原版找出画面里的person把每个人裁剪放大再只在「人」区域里跑自训模型判断是否抽烟/玩手机。这样既提高小目标检出率又省算力——非常契合「监控里看人是否在违规」的场景。七、观点为什么「自训 ONNX」值得而不只是图省事上多模态大模型这是我想借这篇实战讲清楚的核心。现在做视觉很多人「图省事」直接调多模态大模型VLM发图、写 prompt、拿结果。确实快上手。但当你面对的是固定场景、固定目标、高频重复的识别时自训一个轻量 ONNX 模型其实更优1. 推理是毫秒级的不是「请求-等待」。多模态大模型走的是「上传图片 → 云端推理 → 返回」少则几百毫秒多则几秒还受网络、并发、限流制约。而 ONNX onnxruntime 在本地 CPU 就能做到几十毫秒一张图可以做实时视频流、可以塞进边缘设备摄像头、工控机、手机完全离线。2. 固定场景专用模型更准。多模态大模型是「通才」什么都能聊但对你那个具体场景未必最准你自己的数据训出来的模型是「专才」。比如「吸烟」这种需要区分「手里拿的是烟还是笔」的细活用你的数据微调后在固定摄像头角度、固定光照下准确率可以远高于通用大模型。3. 成本结构完全不同。VLM 按调用计费量一大账单就上去了自训是「一次性投入」之后无限次免费推理且没有隐私外泄风险——数据完全不出本地。4. 可控、可迭代、可解释。模型在你手里换数据重训即可升级输出是确定性坐标和置信度能接入下游系统报警、统计、对账做自动化而不是一句「这张图好像有人抽烟」的自然语言。5. 门槛真没那么高。正如上面全程所示标注用 LabelImg训练一行model.train()导出model.export(formatonnx)部署用 Flask。整个链路开源、本地、免费。你不需要是算法专家「会装环境、会改 yaml、会看懂几个指标」就够了。当然多模态大模型在「开放场景、零样本、需要理解语义」时依然无可替代。但当你的需求是「每天从一万张监控帧里挑出抽烟的人」——别再为每个请求花几毛钱去问大模型了训一个 ONNX让它毫秒级、零成本、永远在线地替你盯着。结语从「一张图都没有」到「网页实时出框」中间只是标注 → 一份 yaml → 几行训练代码 → 一次导出 → 一个 Flask。工具都是现成的难点从来不是技术而是「我先动手」。如果你也想试把上面七步跑一遍半天就能拥有自己的第一个 ONNX 模型。要不要我把这份流程连同模型规格MODEL_SPEC.md一起打到包里方便直接交接同事

相关新闻