1. YOLOv6模型转换背景与需求解析
在计算机视觉领域,YOLO系列模型因其高效的实时目标检测能力而广受欢迎。YOLOv6作为该系列的最新演进版本,在精度和速度上都有显著提升。然而在实际工业部署中,我们常常需要将训练好的PyTorch模型转换为特定硬件平台支持的格式。本文将详细介绍如何将YOLOv6的PyTorch模型(best.pt)转换为Sophgo BM1684X芯片支持的bmodel格式。
注意:转换过程中最大的技术难点在于处理ONNX中的Mod算子,因为部分推理引擎(包括Sophgo的工具链)可能不支持该算子。我们需要通过数学等价替换的方式解决这个问题。
2. 环境准备与初始ONNX导出
2.1 基础环境配置
首先确保已安装以下环境:
- Anaconda或Miniconda
- Python 3.7+
- PyTorch 1.8+
- Ultralytics YOLOv6官方仓库
建议使用conda创建独立环境:
bash复制conda create -n yolo26 python=3.8
conda activate yolo26
pip install torch torchvision
git clone https://github.com/ultralytics/yolov6
cd yolov6
pip install -r requirements.txt
2.2 模型导出为ONNX
使用YOLOv6官方提供的导出脚本将best.pt转换为ONNX格式:
bash复制yolo export model=best.pt format=onnx opset=13 imgsz=640 simplify=True nms=False
关键参数说明:
opset=13:指定ONNX算子集版本,13是较稳定的版本imgsz=640:与训练时一致的输入尺寸simplify=True:启用ONNX图优化nms=False:不导出后处理中的NMS部分,通常在部署时单独实现
导出完成后,建议立即备份原始ONNX文件:
bash复制cp yolo26x_640.onnx yolo26x_640.onnx.bak
3. Mod算子处理与ONNX优化
3.1 Mod算子问题分析
YOLOv6模型中可能包含Mod(取模)算子,而Sophgo的BM1684X芯片工具链对此支持有限。我们需要将其替换为基本数学运算组合。数学原理如下:
code复制y = a % b 等价于 y = a - floor(a / b) * b
在ONNX中,我们可以用Div、Mul和Sub三个算子的组合来实现等效计算。
3.2 自动替换Mod算子
以下Python脚本将自动扫描并替换ONNX中的Mod算子:
python复制import onnx
from onnx import helper
src = "yolo26x_640.onnx"
dst = "yolo26x_640_nomod.onnx"
model = onnx.load(src)
graph = model.graph
new_nodes = []
replaced = 0
for node in graph.node:
if node.op_type != "Mod":
new_nodes.append(node)
continue
# 解构Mod算子的输入输出
a, b = node.input[0], node.input[1]
y = node.output[0]
# 生成中间节点名称
div_out = y + "_div"
mul_out = y + "_mul"
# 构建替代算子
div = helper.make_node("Div", [a, b], [div_out], name=(node.name or "Mod") + "_Div")
mul = helper.make_node("Mul", [div_out, b], [mul_out], name=(node.name or "Mod") + "_Mul")
sub = helper.make_node("Sub", [a, mul_out], [y], name=(node.name or "Mod") + "_Sub")
new_nodes.extend([div, mul, sub])
replaced += 1
# 更新计算图
del graph.node[:]
graph.node.extend(new_nodes)
# 验证并保存新模型
onnx.checker.check_model(model)
onnx.save(model, dst)
print(f"Replaced {replaced} Mod nodes")
print("Saved modified ONNX:", dst)
3.3 验证Mod算子替换
运行以下脚本确认Mod算子已被完全替换:
python复制import onnx
from collections import Counter
model = onnx.load("yolo26x_640_nomod.onnx")
ops = Counter(n.op_type for n in model.graph.node)
print("Mod算子存在情况:", "Mod" in ops, "数量=", ops.get("Mod", 0))
预期输出应为:
code复制Mod算子存在情况: False 数量= 0
4. 模型转换全流程
4.1 准备工作目录
建议创建清晰的目录结构:
bash复制mkdir -p /dataSharing/otherTrain/bmodel_build/yolo26x_640
cd /dataSharing/otherTrain/bmodel_build/yolo26x_640
cp /path/to/yolo26x_640_nomod.onnx ./yolo26x_640.onnx
4.2 生成MLIR中间表示
使用Sophgo提供的model_transform.py工具:
bash复制model_transform.py \
--model_name yolo26x_640 \
--model_def yolo26x_640.onnx \
--input_shapes [[1,3,640,640]] \
--mlir yolo26x_640.mlir
关键参数说明:
--input_shapes: 必须与训练时完全一致--mlir: 输出的MLIR文件路径- 其他可选参数包括均值/标准差归一化等
4.3 编译为FP16 bmodel
针对BM1684X芯片进行编译:
bash复制model_deploy.py \
--mlir yolo26x_640.mlir \
--chip bm1684x \
--quantize F16 \
--model yolo26x_640_fp16.bmodel
编译选项说明:
--chip: 指定目标芯片型号--quantize F16: 使用FP16量化,保持较高精度- 也可选择INT8量化,但需要校准数据集
5. 验证与性能测试
5.1 模型文件检查
编译完成后,检查生成的bmodel文件:
bash复制ls -lh *.bmodel
预期会看到类似输出:
code复制-rw-r--r-- 1 user group 124M Jun 1 15:30 yolo26x_640_fp16.bmodel
5.2 推理测试
使用Sophgo提供的推理工具进行测试:
bash复制bmrt_test --bmodel yolo26x_640_fp16.bmodel --input input.npy --output output.npy
实操技巧:可以使用OpenCV生成测试用的随机输入:
python复制import numpy as np
import cv2
# 生成随机测试图像
img = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
# 预处理(与训练时一致)
input_data = cv2.dnn.blobFromImage(img, 1/255.0, (640, 640), swapRB=True)
np.save("input.npy", input_data)
6. 常见问题与解决方案
6.1 ONNX导出失败
问题现象:导出时报错"Unsupported operator: Mod"
解决方案:
- 确保使用最新版YOLOv6
- 尝试降低opset版本(如12)
- 手动修改模型代码,避免使用取模运算
6.2 MLIR生成错误
问题现象:model_transform.py报shape不匹配
排查步骤:
- 使用Netron可视化ONNX模型,检查输入输出shape
- 确认--input_shapes参数与ONNX模型一致
- 检查是否有动态shape,需要固定为具体值
6.3 推理精度下降
可能原因:
- FP16量化导致的精度损失
- Mod算子替换引入的计算误差
验证方法:
- 对比ONNX和bmodel在相同输入下的输出
- 逐步缩小范围,定位精度下降的具体层
7. 性能优化建议
-
INT8量化:使用校准数据集进行INT8量化,可显著提升推理速度
bash复制
model_deploy.py \ --mlir yolo26x_640.mlir \ --chip bm1684x \ --quantize INT8 \ --calibration_table yolo26x_640_cali_table \ --model yolo26x_640_int8.bmodel -
多batch优化:根据实际场景调整input_shapes中的batch大小
bash复制--input_shapes [[4,3,640,640]] # 同时处理4张图像 -
模型剪枝:在转换前对模型进行剪枝,减少参数量
在实际部署中,我发现BM1684X对FP16的支持非常完善,大多数情况下可以保持接近原始FP32的精度。对于Mod算子的替换,虽然理论上会引入微小误差,但在YOLOv6的实际检测任务中几乎不会影响mAP指标。建议在关键应用场景中仍然进行严格的精度验证测试。