Spaces:
Running
on
T4
MMYOLO 模型 ONNX 转换
1. 导出后端支持的 ONNX
环境依赖
-
pip install onnx
onnx-simplifier (可选,用于简化模型)
pip install onnx-simplifier
*** 请确保您在 MMYOLO
根目录下运行相关脚本,避免无法找到相关依赖包。***
使用方法
模型导出脚本用于将 MMYOLO
模型转换为 onnx
。
参数介绍:
config
: 构建模型使用的配置文件,如yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py
。checkpoint
: 训练得到的权重文件,如yolov5s.pth
。--work-dir
: 转换后的模型保存路径。--img-size
: 转换模型时输入的尺寸,如640 640
。--batch-size
: 转换后的模型输入batch size
。--device
: 转换模型使用的设备,默认为cuda:0
。--simplify
: 是否简化导出的onnx
模型,需要安装 onnx-simplifier,默认关闭。--opset
: 指定导出onnx
的opset
,默认为11
。--backend
: 指定导出onnx
用于的后端名称,ONNXRuntime
:onnxruntime
,TensorRT8
:tensorrt8
,TensorRT7
:tensorrt7
,默认为onnxruntime
即ONNXRuntime
。--pre-topk
: 指定导出onnx
的后处理筛选候选框个数阈值,默认为1000
。--keep-topk
: 指定导出onnx
的非极大值抑制输出的候选框个数阈值,默认为100
。--iou-threshold
: 非极大值抑制中过滤重复候选框的iou
阈值,默认为0.65
。--score-threshold
: 非极大值抑制中过滤候选框得分的阈值,默认为0.25
。--model-only
: 指定仅导出模型 backbone + neck, 不包含后处理,默认关闭。
例子:
python ./projects/easydeploy/tools/export.py \
configs/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py \
yolov5s.pth \
--work-dir work_dir \
--img-size 640 640 \
--batch 1 \
--device cpu \
--simplify \
--opset 11 \
--backend 1 \
--pre-topk 1000 \
--keep-topk 100 \
--iou-threshold 0.65 \
--score-threshold 0.25
然后利用后端支持的工具如 TensorRT
读取 onnx
再次转换为后端支持的模型格式如 .engine/.plan
等。
MMYOLO
目前支持 TensorRT8
, TensorRT7
, ONNXRuntime
后端的端到端模型转换,目前仅支持静态 shape 模型的导出和转换,动态 batch 或动态长宽的模型端到端转换会在未来继续支持。
端到端转换得到的 onnx
模型输入输出如图:
输入名: images
, 尺寸 640x640
输出名: num_dets
, 尺寸 1x1,表示检测目标数量。
输出名: boxes
, 尺寸 1x100x4,表示检测框的坐标,格式为 x1y1x2y1
。
输出名: scores
, 尺寸 1x100,表示检测框的分数。
输出名: labels
, 尺寸 1x100,表示检测框的类别 id。
可以利用 num_dets
中的个数对 boxes
, scores
, labels
进行截断,从 100 个检测结果中抽取前 num_dets
个目标作为最终检测结果。
2. 仅导出模型 Backbone + Neck
当您需要部署在非 TensorRT
, ONNXRuntime
等支持端到端部署的平台时,您可以考虑使用--model-only
参数并且不要传递 --backend
参数,您将会导出仅包含 Backbone
+ neck
的模型,模型的部分输出如图:
这种导出方式获取的 ONNX
模型具有如下优点:
- 算子简单,一般而言只包含
Conv
,激活函数等简单算子,几乎不存在无法正确导出的情况,对于嵌入式部署更加友好。 - 方便不同算法之间对比速度性能,由于不同的算法后处理不同,仅对比
backbone
+Neck
的速度更加公平。
也有如下缺点:
- 后处理逻辑需要单独完成,会有额外的
decode
+nms
的操作需要实现。 - 与
TensorRT
相比,由于TensorRT
可以利用多核优势并行进行后处理,使用--model-only
方式导出的模型性能会差很多。
使用方法
python ./projects/easydeploy/tools/export.py \
configs/yolov5/yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py \
yolov5s.pth \
--work-dir work_dir \
--img-size 640 640 \
--batch 1 \
--device cpu \
--simplify \
--opset 11 \
--model-only
使用 model-only
导出的 ONNX 进行推理
模型推理脚本用于推理导出的 ONNX
模型,需要安装基础依赖环境:
pip install onnxruntime
pip install opencv-python==4.7.0.72 # 建议使用最新的 opencv
参数介绍:
img
: 待检测的图片路径或图片文件夹路径。onnx
: 导出的model-only
ONNX 模型。--type
: 模型名称,目前支持yolov5
,yolox
,yolov6
,ppyoloe
,ppyoloep
,yolov7
,rtmdet
,yolov8
。--img-size
: 转换模型时输入的尺寸,如640 640
。--out-dir
: 保存检测结果的路径 。--show
: 是否可视化检测结果。--score-thr
: 模型检测后处理的置信度分数 。--iou-thr
: 模型检测后处理的 IOU 分数 。
使用方法
cd ./projects/easydeploy/examples
python main_onnxruntime.py \
"image_path_to_detect" \
yolov5_s_model-only.onnx \
--out-dir work_dir \
--img-size 640 640 \
--show \
--score-thr 0.3 \
--iou-thr 0.7
注意!!!
当您使用自定义数据集训练得到的模型时,请修改 config.py
中 CLASS_NAMES
和 CLASS_COLORS
,如果是 yolov5
或者 yolov7
基于 anchor
的模型请同时修改 YOLOv5_ANCHORS
和 YOLOv7_ANCHORS
。
numpy_coder.py
是目前所有算法仅使用 numpy
实现的 decoder
,如果您对性能有较高的要求,可以参照相关代码改写为 c/c++
。