1. 模型转换背景与需求解析
在边缘计算和嵌入式AI领域,将训练好的深度学习模型转换为特定硬件平台支持的格式是部署流程中的关键环节。YOLOv6作为目标检测领域的高效模型,其官方提供的best.pt权重文件需要经过特定转换才能在算能科技(Sophgo)的BM1684/BM1688等芯片上运行。这种转换过程不仅涉及文件格式的变化,更需要考虑计算图优化、算子兼容性等一系列技术细节。
我最近在部署一个智能安防项目时,就遇到了将YOLOv6模型转换为bmodel格式的需求。经过多次实践和踩坑,总结出一套稳定可靠的转换方案。与常见的ONNX中间格式转换不同,bmodel转换需要特别注意芯片架构特性,比如对特定算子的支持程度、内存访问模式优化等。下面就从环境准备到最终生成,详细拆解整个转换流程。
2. 环境配置与工具链搭建
2.1 基础环境准备
转换工作需要在Linux环境下进行(推荐Ubuntu 18.04/20.04),主要依赖以下组件:
- Python 3.8+ 环境
- PyTorch 1.10+(需与原始模型训练版本匹配)
- Sophgo SDK(需从官方获取对应芯片版本的SDK包)
- OpenCV(用于预处理验证)
安装核心依赖的命令示例:
bash复制conda create -n yolo6_bmodel python=3.8
conda activate yolo6_bmodel
pip install torch==1.10.0 torchvision==0.11.1
pip install opencv-python onnx
重要提示:Sophgo SDK的版本必须与目标芯片严格对应。例如BM1684X芯片需要使用SDK版本v23.07.01,而BM1688则需要v23.10.01及以上版本。
2.2 模型转换工具安装
Sophgo提供的模型转换工具链主要包括:
model_transform.py:模型格式转换工具bmneto:编译生成bmodel的核心工具calibration_tool:量化校准工具
这些工具通常位于SDK的tools/目录下。建议将工具路径加入环境变量:
bash复制export PATH=$PATH:/path/to/sophgo_sdk/tools
3. 模型转换全流程详解
3.1 PyTorch到ONNX的转换
首先需要将.pt模型转换为ONNX中间格式。YOLOv6的导出需要特别注意输出节点的命名:
python复制import torch
from models.yolo import Model
# 加载预训练模型
ckpt = torch.load('best.pt', map_location='cpu')
model = Model(ckpt['model'].yaml, ch=3, nc=80) # nc为类别数
model.load_state_dict(ckpt['model'].float().state_dict())
# 设置导出参数
input_shape = (1, 3, 640, 640) # BS,C,H,W
dummy_input = torch.randn(input_shape)
# 导出ONNX
torch.onnx.export(
model,
dummy_input,
'yolov6.onnx',
opset_version=11,
input_names=['images'],
output_names=['output'],
dynamic_axes={
'images': {0: 'batch'},
'output': {0: 'batch'}
}
)
关键参数说明:
opset_version=11:确保使用支持的ONNX算子集dynamic_axes:保留batch维度动态性以便适配不同批大小- 输出节点必须命名为单输出(YOLOv6默认多输出需要修改)
3.2 ONNX模型优化
生成的ONNX模型需要经过以下优化步骤:
- 图结构优化:
bash复制python -m onnxsim yolov6.onnx yolov6_sim.onnx
- 节点检查:
bash复制polygraphy inspect model yolov6_sim.onnx --mode=basic
常见问题处理:
- 如果出现
GridSample算子不支持,需要修改模型结构替换为可支持算子 - 遇到
Resize算子版本不兼容时,可尝试固定为opset 11的实现方式
3.3 生成FP32 BModel
使用Sophgo工具链进行转换:
bash复制model_transform.py \
--model_name yolov6 \
--model_def yolov6_sim.onnx \
--input_shapes [[1,3,640,640]] \
--mean 0,0,0 \
--scale 0.003921568627,0.003921568627,0.003921568627 \
--keep_aspect_ratio \
--pixel_format rgb \
--output_names output \
--mlir yolov6.mlir
bmneto --model=yolov6.mlir \
--target=BM1684 \
--outdir=./fp32 \
--opt=2 \
--cmp=true \
--shapes=[1,3,640,640]
参数解析:
--keep_aspect_ratio:保持图像原始宽高比处理--opt=2:启用高级优化级别--target:指定目标芯片型号--cmp=true:启用计算图压缩优化
4. 模型量化与INT8优化
4.1 校准集准备
准备约500-1000张具有代表性的图片,存放在calib_data/目录下。建议使用如下格式的校准集:
code复制calib_data/
├── images/
│ ├── 000001.jpg
│ ├── 000002.jpg
│ └── ...
└── calib_list.txt
其中calib_list.txt每行包含图片路径和尺寸信息:
code复制images/000001.jpg 640,480
images/000002.jpg 1280,720
4.2 量化校准执行
bash复制calibration_tool \
--model=yolov6.mlir \
--data_list=calib_data/calib_list.txt \
--data_path=calib_data \
--output_name=yolov6_cali_table \
--target=BM1684
4.3 生成INT8 BModel
bash复制bmneto --model=yolov6.mlir \
--target=BM1684 \
--outdir=./int8 \
--opt=2 \
--cmp=true \
--shapes=[1,3,640,640] \
--calibration_table=yolov6_cali_table \
--quantize INT8
量化效果验证:
- 通常INT8模型体积减小为FP32的1/4
- 推理速度提升2-3倍
- 精度损失应控制在1% mAP以内(可通过官方测试工具验证)
5. 转换过程中的常见问题与解决方案
5.1 算子不支持错误
典型报错:
code复制Unsupported operator: GridSample
解决方案:
- 修改YOLOv6模型源码,将
RepVGGBlock中的GridSample替换为可支持算子 - 使用Sophgo提供的自定义算子插件(需申请获取)
5.2 精度异常问题
现象:转换后模型检测结果异常或漏检
排查步骤:
- 验证ONNX模型在CPU上的推理结果是否正常
- 检查预处理参数(mean/scale)是否与训练时一致
- 确认输入图像格式(RGB/BGR)匹配
- 检查校准集是否具有代表性
5.3 性能优化技巧
-
多batch优化:
生成时指定--shapes=[4,3,640,640]可优化多batch推理性能 -
内存布局调整:
添加--input_layout NHWC可能提升部分芯片的访存效率 -
算子融合:
在mlir生成阶段添加--fuse_preprocess可将归一化操作融合到模型中
6. 模型部署验证
6.1 基础推理测试
使用Sophgo提供的测试工具验证bmodel:
bash复制bmrt_test --bmodel ./int8/compilation.bmodel --input input.npy --output output.npy
6.2 Python接口调用示例
python复制import bmruntime as bm
import numpy as np
# 初始化环境
handle = bm.runtime.Bmruntime()
net = bm.runtime.Net(handle, "compilation.bmodel")
# 准备输入
input_data = np.random.rand(1, 3, 640, 640).astype(np.float32)
input_tensor = bm.runtime.Tensor(input_data)
# 执行推理
output_tensor = net.forward([input_tensor])[0]
detections = output_tensor.asnumpy()
# 后处理
# ... (实现解码和NMS)
6.3 性能基准测试
典型性能指标(BM1684X芯片):
| 模型类型 | 推理时延(ms) | 峰值内存(MB) | 计算利用率(%) |
|---|---|---|---|
| FP32 | 15.2 | 342 | 78 |
| INT8 | 6.8 | 210 | 92 |
实测建议:在最终部署环境中进行端到端测试,考虑视频流处理时的流水线优化