mirror of https://github.com/open-mmlab/mmyolo.git
57 lines
2.1 KiB
Markdown
57 lines
2.1 KiB
Markdown
|
# MMYOLO 模型 ONNX 转换
|
|||
|
|
|||
|
## 环境依赖
|
|||
|
|
|||
|
- [onnx](https://github.com/onnx/onnx)
|
|||
|
|
|||
|
```shell
|
|||
|
pip install onnx
|
|||
|
```
|
|||
|
|
|||
|
[onnx-simplifier](https://github.com/daquexian/onnx-simplifier) (可选,用于简化模型)
|
|||
|
|
|||
|
```shell
|
|||
|
pip install onnx-simplifier
|
|||
|
```
|
|||
|
|
|||
|
## 使用方法
|
|||
|
|
|||
|
[模型导出脚本](./projects/easydeploy/tools/export.py)用于将 `MMYOLO` 模型转换为 `onnx` 。
|
|||
|
|
|||
|
### 参数介绍:
|
|||
|
|
|||
|
- `config` : 构建模型使用的配置文件,如 [`yolov5_s-v61_syncbn_fast_8xb16-300e_coco.py`](./configs/yolov5/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](https://github.com/daquexian/onnx-simplifier),默认关闭。
|
|||
|
- `--opset`: 指定导出 `onnx` 的 `opset`,默认为 `11` 。
|
|||
|
- `--backend`: 指定导出 `onnx` 用于的后端 id,`ONNXRuntime`: `1`, `TensorRT8`: `2`, `TensorRT7`: `3`,默认为`1`即 `ONNXRuntime`。
|
|||
|
- `--pre-topk`: 指定导出 `onnx` 的后处理筛选候选框个数阈值,默认为 `1000`。
|
|||
|
- `--keep-topk`: 指定导出 `onnx` 的非极大值抑制输出的候选框个数阈值,默认为 `100`。
|
|||
|
- `--iou-threshold`: 非极大值抑制中过滤重复候选框的 `iou` 阈值,默认为 `0.65`。
|
|||
|
- `--score-threshold`: 非极大值抑制中过滤候选框得分的阈值,默认为 `0.25`。
|
|||
|
|
|||
|
例子:
|
|||
|
|
|||
|
```shell
|
|||
|
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` 等
|