简介:面向计算机视觉研究者、深度学习开发者及目标检测与实例分割应用人员,这套 Ultralytics YOLOv8 项目覆盖模型训练、验证到部署全流程,可用于物体识别、图像分析、安全监控、无人机视觉与自动驾驶等场景的科研实验和原型开发。压缩包共788个文件、约26.63MB,以Markdown文档、Python源码、预编译文件、YAML配置为主,并含Notebook、C++/Rust推理示例、预训练权重和多套Dockerfile,便于算法复现、参数调整与跨环境部署。已有528人学习。读者可据此梳理数据准备、训练评估到推理部署的流程,结合配置文件和脚本完成自定义数据集训练,并通过多平台容器配置了解CPU、GPU、Jetson等环境的部署思路,减少从零搭建工程的时间成本。
1. 从一次“数错零件”的返工说起
前阵子有个做仓储盘点的开发者拿了个 Demo 来找我,摄像头底下堆着几十个形状相近的金属件。用纯目标检测框出来能数个数,但一遇到重叠就露馅——框叠框,到底是一个还是两个全靠猜。换成 YOLOv8 实例分割之后,每个零件既有外接框又有像素级轮廓,重叠部分按 mask 面积一算就清楚了。
这份 YOLOv8 实例分割 + 目标检测的实战包,干的就是这条线:数据标注、模型训练、导出部署,全部用一套 Python API 串起来,代码量不大,但该有的环节一个不少。适合手上已经有一批自己的图片、想训一个分割模型却不想从论文和 loss 公式啃起的开发者;也适合本来就跑通了检测、想给项目补上“轮廓”能力的熟手。下面按我拆包的顺序讲,从目录骨架一直聊到怎么把 mask 变成一个业务能直接用的数字。
2. 目录结构与技术选型:先把工程骨架摸清楚
拿到一个压缩包,我习惯先不急着pip install,而是把目录树打出来看一遍。原因很简单:这类项目最容易出问题的地方往往不是模型本身,而是“数据放哪、权重放哪、输出往哪写”这套约定没对齐。骨架看明白了,后面每一步操作都能对上号。
2.1 目录骨架:数据、配置、训练、推理各在哪
包里的结构大致是这样,我用tree之类工具打印过一遍,做了点归类:
# 树形结构示意,重点是数据与代码分区 yolov8-seg-project/ ├── datasets/ # 数据集根目录,对应 data.yaml 里的 path │ ├── images/ │ │ ├── train/ # 训练集图片 │ │ └── val/ # 验证集图片 │ ├── labels/ │ │ ├── train/ # 与图片同名的 .txt 标签 │ │ └── val/ │ └── data.yaml # 数据集描述文件,训练入口靠它找数据 ├── models/ │ └── yolov8n-seg.yaml # 模型结构配置,改类别数在这里 ├── weights/ │ └── yolov8n-seg.pt # 预训练权重,用于迁移学习 ├── train.py # 训练脚本 ├── predict.py # 单图/批量推理 ├── export.py # 导出 ONNX / TensorRT ├── split_dataset.py # 原始图片划分 train/val └── requirements.txt几个约定值得单独说一句。images和labels是平行的两棵树,文件名必须一一对应,abc.jpg对应abc.txt,扩展名不同但主名相同,这是这套框架找标签的唯一依据,改不了。data.yaml里的path指向数据集根,train/val是相对这个根的路径,写绝对路径也能跑,但换台机器就废,我不建议。
weights目录和models目录容易被搞混:前者存的是.pt权重文件,里面既有结构也有参数;后者存的是.yaml结构描述,只有网络长什么样,没有训练出来的数值。日常改类别数改的是models/下的 yaml,迁移学习加载的是weights/下的 pt,两个都要动,但动的位置不一样。
2.2 技术栈与版本组合:为什么是 PyTorch + ultralytics
这套项目的技术栈不复杂,都是围绕 PyTorch 生态转的:
| 组件 | 作用 | 备注 |
|---|---|---|
| PyTorch | 训练与推理的底层框架 | 版本要和本机 CUDA 驱动对齐 |
| ultralytics | 提供 YOLO 模型、训练器、导出器 | 一条model.train()全包 |
| OpenCV | 读写图片、画框画 mask | ultralytics 会自动装 |
| NumPy | 数组运算、后处理 | 后处理阶段离不开 |
| onnx / onnxruntime | 导出与跨平台推理 | 部署环节才用得上 |
| labelme 类工具 | 多边形标注 | 产出的 json 需要转换 |
选 ultralytics 这套 API 的核心理由是省事:模型定义、数据加载、增强、学习率调度、验证指标、权重保存,全在一个train()里,配置文件用 yaml 描述,改参数不用碰源码。代价是黑盒感比较强,出了问题得知道往哪翻日志——这件事后面第 5 章会专门讲。
版本上我一般会锁死,不吃“最新版”的亏。Python 3.10、PyTorch 2.1 上下的 LTS 版本、ultralytics 8.x 系列,这个组合我跑过很多次,比较稳。ultralytics 8.x 内部迭代很快,小版本之间 API 偶有变动,比如某些参数名调整过,所以requirements.txt里最好写成ultralytics==8.x.y这种固定版本,别留>=。
2.3 环境安装与第一次跑通
安装这一步,血泪经验是:先装 PyTorch,再装 ultralytics。顺序反了的话,ultralytics 会顺手给你拉一个 CPU 版的 torch,等到训练时发现torch.cuda.is_available()返回 False,再回头卸了重装,很折腾。
# 建独立环境,别污染系统 Python conda create -n seg python=3.10 -y conda activate seg # 第一步:按本机 CUDA 版本装 PyTorch,cu121 只是示例,按实际改 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121 # 第二步:装 ultralytics,它会带出 opencv、numpy、matplotlib 等依赖 pip install ultralytics # 第三步:部署相关的包,训练阶段用不上,先装着备用 pip install onnx onnxruntime onnxsim装完之后先别急着训练,跑一个自检命令,把环境、CUDA、依赖版本全打出来:
# 输出 Python 版本、torch 版本、CUDA 是否可用、依赖是否齐全 yolo checks这一步的输出里重点看三行:torch版本、CUDA是否 available、ultralytics版本。CUDA 显示不可用但你有显卡,九成是驱动版本和 torch 的 CUDA 版本对不上,去 PyTorch 官网查兼容表重装,别在这里硬扛。
环境没问题后,用官方权重跑一次推理,确认整条链路是通的:
from ultralytics import YOLO # 首次运行会自动下载 yolov8n-seg.pt 到当前目录,也可以提前放进 weights/ model = YOLO("weights/yolov8n-seg.pt") # 对单张图做推理,返回的是 Results 对象列表,一张图一个 results = model("datasets/images/val/sample.jpg", imgsz=640, conf=0.25) # 打印检测框坐标和对应的 mask 数量 r = results[0] print(r.boxes.xyxy.shape) # (n, 4),n 是检出的实例数 print(r.masks.data.shape) # (n, h, w),每个实例一张二值 mask r.save("runs/demo.jpg") # 把框和 mask 画到图上存下来这段代码里,imgsz=640是推理分辨率,和训练时保持一致最稳;conf=0.25是置信度阈值,低于它的框直接丢。r.masks.data是分割任务独有的,形状是(实例数, 高, 宽),值在 0 到 1 之间,后面做面积统计就靠它。如果这一步能出图,说明环境、权重、OpenCV 都没问题,可以进数据环节了。
3. 数据标注与配置文件:实例分割的标签不是矩形框
检测和分割最容易翻车的地方不在模型,在标签。很多人拿检测的数据集直接改个后缀就想训分割,结果 loss 一路降、mAP 死活上不去,最后发现标签格式压根不对。这一章把标签长什么样、data.yaml 每个字段什么意思、数据集怎么划分讲透。
3.1 分割标签的 txt 到底长什么样
检测任务的标签是一行五个数:class_id cx cy w h,中心点加宽高,全部归一化到 0 到 1。分割任务的标签是变长的:class_id x1 y1 x2 y2 ... xn yn,后面跟的是多边形所有顶点的归一化坐标,顶点数不固定,几个到几百个都可能。
# 检测标签:一行 5 个数 0 0.512 0.438 0.221 0.316 # 分割标签:一行 1 + 2n 个数,n 是多边形顶点数 0 0.482 0.391 0.531 0.402 0.548 0.467 0.497 0.489 0.451 0.452标注工具产出的通常是 json,需要转成这种 txt。用 labelme 这类多边形标注工具的话,转换脚本大概是这样:
import json from pathlib import Path def labelme_to_yolo(json_path, txt_path, class_map, img_w, img_h): """把 labelme 的 json 转成 YOLO 分割格式的 txt""" data = json.load(open(json_path, encoding="utf-8")) lines = [] for shape in data["shapes"]: label = shape["label"] # 类别不在映射表里就跳过,避免脏标签污染训练 if label not in class_map: continue pts = shape["points"] # 矩形标注转成四点多边形,兼容一部分只画框的旧数据 if shape["shape_type"] == "rectangle": (x1, y1), (x2, y2) = pts pts = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]] coords = [] for x, y in pts: # 归一化并裁剪到 [0,1],防止标注越界导致训练报错 coords.append(f"{min(max(x / img_w, 0.0), 1.0):.6f}") coords.append(f"{min(max(y / img_h, 0.0), 1.0):.6f}") lines.append(f"{class_map[label]} " + " ".join(coords)) Path(txt_path).write_text("\n".join(lines), encoding="utf-8")逻辑上有三处值得说。第一,坐标除以图片宽高做归一化,这一步不能省,模型输入是归一化后的,标签不归一化会导致 loss 爆炸。第二,min(max(...))的裁剪是保命的,标注时鼠标稍微拖出边界一点点,坐标就会小于 0 或大于 1,训练时直接给你一个 NaN。第三,矩形转四点多边形只是兜底,真要做分割,多边形顶点密度要够,一个圆形零件至少十几个点,否则 mask 出来是棱角分明的多边形,边缘全是锯齿。
调用的时候按图片实际尺寸传参,别用固定值:
from PIL import Image img = Image.open("raw/images/part_001.jpg") w, h = img.size labelme_to_yolo( "raw/jsons/part_001.json", "datasets/labels/train/part_001.txt", class_map={"bolt": 0, "nut": 1, "washer": 2}, img_w=w, img_h=h, )class_map里的数字就是后续data.yaml里names的顺序,两边必须严格一致,这是最容易出错的地方——标注时bolt是 0,yaml 里写成了 1,训练不会报错,但模型学出来的东西全是错位的。
3.2 data.yaml 每个字段的含义
这个文件是训练入口的“地图”,路径错了训练直接起不来。一份能跑的最小配置长这样:
# 数据集根目录,train/val 都相对它解析 path: ../datasets # 训练图片目录,相对 path train: images/train # 验证图片目录,相对 path val: images/val # 类别名,键就是标签文件里的 class_id names: 0: bolt 1: nut 2: washer关于names有一个坑要提醒:早期版本要求显式写nc: 3表示类别数,新版本可以省略,框架会从names的长度自动推。但如果你的names用了列表写法(names: ['bolt', 'nut', 'washer']),顺序就是 id 顺序,从 0 开始。两种写法都行,别混着来。
path用相对路径时,是相对于data.yaml所在目录解析的,不是相对于你执行命令的目录。这一点很多人搞错:在项目根目录执行python train.py,yaml 里写path: datasets,框架会去找项目根/datasets,但如果 yaml 放在configs/子目录里,就会变成项目根/configs/datasets,找不到数据。我的习惯是 yaml 直接放数据集根目录旁边,路径写成../datasets这种一眼能对上的形式。
3.3 数据集划分脚本:别让同一张图跨集合
划分看着简单,随机打乱按 8:2 分就完了,但这里有个隐蔽的坑:如果你的图片是从视频里抽帧来的,相邻帧长得几乎一样,随机划分会让训练集和验证集里出现“双胞胎”,验证指标虚高得离谱,上线一跑全崩。正确做法是按视频或按采集批次分组划分。
import random import shutil from pathlib import Path random.seed(42) # 固定种子,保证每次划分结果可复现 src_img = Path("raw/images") src_lbl = Path("raw/labels") dst = Path("datasets") # 按文件名前缀分组,同一批采集的图片进同一个集合 groups = {} for img in sorted(src_img.glob("*.jpg")): key = img.stem.split("_")[0] # 例如 batchA_001 -> batchA groups.setdefault(key, []).append(img) keys = list(groups.keys()) random.shuffle(keys) split = int(len(keys) * 0.8) train_keys = set(keys[:split]) for key, imgs in groups.items(): subset = "train" if key in train_keys else "val" for img in imgs: lbl = src_lbl / f"{img.stem}.txt" # 没有对应标签的图片直接跳过,避免训练时报 missing label if not lbl.exists(): continue shutil.copy(img, dst / "images" / subset / img.name) shutil.copy(lbl, dst / "labels" / subset / lbl.name)这段代码的核心是key = img.stem.split("_")[0],把同一批采集的图绑在一起划分,杜绝数据泄漏。另外那个lbl.exists()判断也别省,标注漏掉的图片混进去,训练时会在 dataloader 里抛异常,报错信息还不一定指得清楚是哪张图。
划分完建议再跑一个统计,看看每个类别的实例数和图片数,类别严重不平衡的话,训练前就得想好补偿策略:
from collections import Counter from pathlib import Path counter = Counter() for txt in Path("datasets/labels/train").glob("*.txt"): for line in txt.read_text(encoding="utf-8").splitlines(): if line.strip(): counter[line.split()[0]] += 1 # 每行第一个字段是 class_id print(counter) # 例如 Counter({'0': 1200, '1': 800, '2': 150})如果某个类别的实例数只有其他类的十分之一,先别急着上模型,回头补标注比调参有效得多。
4. 训练核心:从模型 yaml 到 train() 的每个参数
数据齐了就该动模型了。这一章我按“改结构 → 写训练脚本 → 读指标”的顺序拆,每一段都对应包里实际的文件。
4.1 模型结构配置:yolov8n-seg.yaml 里改什么
打开models/yolov8n-seg.yaml,最上面几行是你要动的:
# 类别数,这是唯一必须改的一行 nc: 3 # 模型缩放系数,n 表示 nano 版本,深度 0.33、宽度 0.25 scales: n: [0.33, 0.25, 1024] # 骨干网络,负责提特征,一般不动 backbone: - [-1, 1, Conv, [64, 3, 2]] # ... 中间省略若干层 # 检测头,分割任务在这里多一个 mask 分支 head: # ... 上采样与特征融合层 - [[15, 18, 21], 1, Segment, [nc, 32, 256]]nc必须改成你自己的类别数,改成 3 之后,检测头的输出通道会从 80 类的配置自动收缩。scales里那个n决定模型大小,n 最小最快,s/m/l/x 依次变大,显存够、精度要求高就往大选,但训练时间也跟着涨。
最后一行Segment是分割任务和检测任务的分水岭,[nc, 32, 256]里的32是 mask 系数的维度——模型不是直接输出一张 mask,而是输出 32 个系数,再和一组原型 mask 做线性组合得到最终结果。这个设计是分割头能轻量的关键,后处理时必须按这个逻辑解码,第 6 章会写具体代码。
如果不想手改 yaml,直接用YOLO("yolov8n-seg.pt")也能跑,框架会自动把最后一层的类别数替换掉。但我一般还是走 yaml 这条路,改了什么心里有数,出问题好排查。
4.2 train.py 逐行拆解
训练脚本是整包的核心,参数多,但真正影响结果的没几个。先看完整版:
from ultralytics import YOLO # 从结构文件构建模型,再加载预训练权重做迁移学习 # 这一步等价于“新网络 + 旧知识”,比从零训收敛快很多 model = YOLO("models/yolov8n-seg.yaml").load("weights/yolov8n-seg.pt") results = model.train( data="datasets/data.yaml", # 数据集配置,路径要对 epochs=100, # 总轮数,小数据集 100 起步够用 imgsz=640, # 训练分辨率,和推理保持一致 batch=8, # 批大小,显存不够就往下调 device=0, # 0 表示第一块 GPU,CPU 写 "cpu" workers=4, # 数据加载线程,Windows 上别超过 4 project="runs/seg", # 输出根目录 name="parts_v1", # 本次实验名,结果存 runs/seg/parts_v1 patience=30, # 30 轮指标不涨就早停,省时间 optimizer="SGD", # 优化器,小数据集 SGD 比 AdamW 稳 lr0=0.01, # 初始学习率 lrf=0.01, # 最终学习率 = lr0 * lrf,余弦退到底 momentum=0.937, # SGD 动量 weight_decay=0.0005, # 权重衰减,抑制过拟合 warmup_epochs=3.0, # 前 3 轮热启动,学习率从低往上爬 mosaic=1.0, # 马赛克增强概率,分割任务建议保留 close_mosaic=10, # 最后 10 轮关掉 mosaic,让模型适应真实分布 copy_paste=0.3, # 实例复制粘贴增强,分割任务的关键参数 mixup=0.0, # 混合增强,分割任务慎用,容易糊 mask cache=False, # 缓存图片到内存,数据集小可以开 amp=True, # 混合精度,省显存还能快一点 resume=False, # 断点续训时改成 True exist_ok=False, # 同名实验目录是否覆盖 )参数逐个说重点。batch=8是个保守值,8G 显存的卡跑 640 分辨率差不多这个数,爆显存就降到 4,别硬撑,amp=True已经帮你省了一半。workers在 Windows 上设太大容易卡死,4 是个安全线,Linux 上可以拉到 8。
copy_paste=0.3是分割任务里我最看重的增强,它把图里已有的实例抠出来贴到别的位置,等于凭空造出更多重叠场景。零件堆叠这类任务,开了它 mAP 能涨好几个点。但注意它依赖 mask 质量,标签边缘糙的话,贴出来的图也是脏的。
close_mosaic=10和mosaic=1.0是一对。mosaic 把四张图拼成一张,能显著提升小目标召回,但拼出来的图分布和真实场景差得远,所以最后 10 轮关掉,让模型在“正常图”上收个尾。这个操作对最终指标的影响比调学习率还明显。
patience=30是早停,验证指标连续 30 轮不涨就停。小数据集上这个参数能帮你省掉大量无效训练时间,但也别设太小,比如 10 轮,模型可能还没爬出平台期就被掐了。
跑起来之后,终端会实时打印每一轮的 loss 和 mAP,同时往runs/seg/parts_v1/下写日志、权重、混淆矩阵和验证集的可视化结果。weights 目录里会有best.pt和last.pt,前者是验证指标最好的一版,后者是最后一轮,部署用best.pt。
4.3 训练日志与指标怎么读
日志里一堆缩写,第一次看容易懵。按重要程度排一下:
| 指标 | 含义 | 怎么看 |
|---|---|---|
| box_loss | 检测框回归损失 | 应持续下降,后期趋平 |
| seg_loss | mask 分割损失 | 下降慢于 box_loss 是正常的 |
| cls_loss | 分类损失 | 波动大说明类别不平衡 |
| dfl_loss | 分布焦点损失 | 影响框的精细度,一般不用管 |
| mAP50(B) | 检测框 IoU=0.5 的平均精度 | 检测能力的直观指标 |
| mAP50-95(B) | 检测框多阈值平均 | 更严格,看这个更真实 |
| mAP50(M) | mask IoU=0.5 的平均精度 | 分割任务的核心指标 |
| mAP50-95(M) | mask 多阈值平均 | 上线前重点看这个 |
判断训练健不健康,我一般看三条线。第一,seg_loss和box_loss是不是同步在降,如果 box 降而 seg 不降,多半是标签多边形的点太少,mask 本身就画不准,模型学不到细节。第二,mAP50(B)和mAP50(M)的差距,正常情况下 mask 比 box 低 5 到 10 个点,差距过大说明轮廓质量有问题。第三,训练集指标和验证集指标的分叉,验证集早早走平而训练集还在涨,就是过拟合了,该加增强或者减模型容量。
还有个小技巧:训练结束后去runs/seg/parts_v1/下翻val_batch*.jpg,那是验证集的可视化结果,框和 mask 都画在图上。指标看不出的问题,肉眼一眼就能发现——比如某个类别的 mask 全部偏左,八成是标注时坐标转换错了。
5. 常见问题排查:五个真金白银踩出来的坑
这一章单独拎出来,是因为下面这几条我在不同项目里反复遇到,每条都按“现象 → 原因 → 解决”写,照着对号入座就行。
5.1 训练一启动就报类别越界
现象:刚跑train.py没几秒就抛异常,提示标签里的类别 id 超出了数据集定义的类别数。
原因:data.yaml里names只写了 3 个类(0、1、2),但某个标签文件里出现了3。常见于标注工具里多建了一个类别,或者转换脚本的class_map写漏了。
解决:跑一段脚本扫一遍所有标签,打印出现过的最大 class_id,和 yaml 里的类别数对一下:
from pathlib import Path max_id = 0 for txt in Path("datasets/labels").rglob("*.txt"): for line in txt.read_text(encoding="utf-8").splitlines(): if line.strip(): max_id = max(max_id, int(line.split()[0])) print("最大类别 id:", max_id) # 应该小于 len(names)对不上就改 yaml 或者修标签,别指望框架自动兜底。
5.2 loss 跑到一半变成 NaN
现象:前十几轮正常,突然 loss 打印nan,之后再也不恢复。
原因:九成是学习率太大或者标签坐标越界。lr0=0.01对小数据集偏大,另外标注时坐标超出 [0,1] 也会让梯度炸掉。
解决:先把lr0降到 0.001 重跑,同时用 3.1 节的脚本重新扫一遍标签,确认所有坐标都在 0 到 1 之间。如果两者都没问题,检查amp=True在个别显卡上会不会溢出,关掉混合精度试试。
5.3 显存爆了但 batch 已经很小
现象:batch=4还是报 CUDA out of memory。
原因:imgsz太大、workers太多、或者开了cache=True把整个数据集读进内存。
解决:按影响大小依次调——先把imgsz从 640 降到 512,再把workers降到 2,最后确认cache=False。另外amp=True必须开着,能省将近一半显存。如果这些都不够,说明模型选大了,把yolov8n-seg换成更小的配置,或者换成yolov8n-seg的 n 版本从头训。
5.4 mask 边缘全是锯齿,或者中间有空洞
现象:验证集可视化里,mask 轮廓像被狗啃过,圆形零件出来是多边形。
原因:标注多边形顶点太少,或者训练分辨率不够。分割头输出的原型 mask 是输入分辨率的四分之一,640 输入对应 160×160 的原型,小目标在这一步就丢细节了。
解决:两条路一起走。标注端提高顶点密度,圆形至少 16 个点;训练端把imgsz提到 960 或 1024,牺牲速度换边缘质量。如果目标本身很小,还可以考虑切图训练,把大图裁成小块再喂。
5.5 导出 ONNX 后结果和 PyTorch 对不上
现象:.pt推理正常,导出的 ONNX 用 onnxruntime 跑出来框全偏了,或者 mask 全黑。
原因:三种可能——预处理没对齐(letterbox 的填充比例算错)、后处理漏了 mask 解码、opset 版本太低不支持某些算子。
解决:先确认导出参数,用opset=12以上、simplify=True;再检查输入张量的预处理,PyTorch 侧是 RGB、归一化到 0-1、letterbox 填充到 640×640,ONNX 侧必须一模一样;最后确认 mask 后处理有没有做“32 个系数 × 原型 mask”的解码,这一步漏了就是全黑。第 6 章给了解码的核心代码。
6. 二次开发:把 mask 变成业务能用的数字
工程跑通只是起点,真正交付的时候,业务方要的不是一张花花绿绿的图,而是“这块料面积多少、有没有超出公差、一共几个”。这一段落到改造思路上,给你两个最常用的改法。
第一个改法是加后处理。分割模型输出的是 mask,转成面积只需要按像素统计再乘标定系数:
import numpy as np def masks_to_area(result, pixel_per_mm=1.0, min_area_px=50): """把 ultralytics 的 Results 转成每个实例的面积(平方毫米)""" areas = [] if result.masks is None: return areas # masks.data 形状 (n, h, w),值在 0-1 之间 masks = result.masks.data.cpu().numpy() for m in masks: binary = (m > 0.5).astype(np.uint8) # 二值化,阈值 0.5 px = int(binary.sum()) if px < min_area_px: # 过滤噪点,太小的不要 continue areas.append(px / (pixel_per_mm ** 2)) return areaspixel_per_mm是相机标定出来的换算系数,拍一张已知尺寸的标定板算一次就行。min_area_px这个过滤别省,模型偶尔会在背景上吐出一小块假 mask,面积小得离谱,不滤掉会污染统计结果。
第二个改法是换任务。这套代码的复用性比想象中好,换类别只需要动三处:重标数据、改data.yaml的names、改模型 yaml 的nc,然后从头训。想换成检测任务,把模型文件换成yolov8n.yaml和yolov8n.pt,标签格式改成class cx cy w h,其余训练脚本一行不用改。
如果要做跨平台部署,导出 ONNX 之后的后处理是关键:
import numpy as np import onnxruntime as ort sess = ort.InferenceSession("best.onnx") # 输入名和输出名用 sess.get_inputs() / get_outputs() 打印确认 outputs = sess.run(None, {"images": blob}) # blob: 1x3x640x640, float32 preds, protos = outputs[0], outputs[1] # (1, 116, 8400) 和 (1, 32, 160, 160) preds = np.squeeze(preds).T # (8400, 116) boxes, scores, coeffs = preds[:, :4], preds[:, 4:4 + 3], preds[:, 4 + 3:] # 每个实例的 mask = sigmoid(32 个系数 @ 32x160x160 原型) flat = protos.reshape(32, -1) masks = 1 / (1 + np.exp(-(coeffs @ flat))) # NMS 之后再算,别全算维度上要留个心眼:116 = 4 个框参数 + 3 个类别分数 + 32 个 mask 系数,类别数变了这个数也跟着变。protos是 160×160,对应 640 输入的四分之一分辨率。这套解码逻辑和 PyTorch 侧是同一套,导出的模型只是把网络搬过去了,后处理还得自己写。
验证改造有没有效,我一般用两条线交叉确认:一条是拿同一批图分别过.pt和.onnx,对比框坐标的误差是否在 1 像素以内;另一条是统计面积的分布,和人工测量值对一下,偏差超过 5% 就回头查标定系数。
从那以后我每次拿到一个新数据集,都强制先跑一遍标签扫描和类别统计,再开训。这个习惯帮我省掉的返工,比任何调参技巧都值钱。希望帮到你。
本文还有配套的精品资源,点击获取