YOLOv5的pytorch模型文件转换为ONNX文件

文章目录[隐藏]

YOLOv5下载与测试运行
导出ONNX格式文件
ONNX转为为IR中间格式

环境：

Windows 10
Anaconda 2.0.4
OpenVINO 工具包 2021.2
Python 3.6.13
torch 1.9.0
onnx 1.10.1
YOLOv5

YOLOv5下载与测试运行

YOLOv5是第二个非官方的YOLO对象检测版本，也是第一个Pytorch实现的YOLO对象检测版本。Github地址：https://github.com/ultralytics/yolov5

克隆到本地

git clone https://github.com/ultralytics/yolov5.git

安装YOLOv5所有依赖

pip install -r requirements.txt

导出ONNX格式文件

OpenVINO 工具包 2021.2 可以直接读取ONNX格式文件，所以我们既可以通过脚本直接导出onnx格式文件，直接给OpenVINO调用，也可以对得到ONNX文件通过OpenVINO的模型转换生成IR中间格式(.bin文件与.xml文件)。

Pytorch的YOLOv5项目本身已经提供了转换脚本，命令行运行方式如下：

(python37) C:\Program Files (x86)\Intel\openvino_2021.2.185\bin> setupvars.bat
Python 3.7.10

[setupvars.bat] OpenVINO environment initialized

cd 到YOLOv5项目所在目录下

使用 YOLOv5 提供的 export.py 将 yolov5s.pt 转换为 ONNX。

python models/export.py --weights yolov5s.pt --img 640 --batch 1

(python37) M:\python\OpenCV\yolov5\yolov5-master>python models/export.py --weights yolov5s.pt --img 640 --batch 1
Namespace(batch_size=1, device=‘cpu’, dynamic=False, half=False, img_size=[640, 640], include=[‘torchscript’, ‘onnx’, ‘coreml’], inplace=False, opset_version=12, optimize=False, simplify=False, train=False, weights=‘yolov5s.pt’)
YOLOv5 2021-5-18 torch 1.8.1+cpu CPU
Fusing layers…
Model Summary: 224 layers, 7266973 parameters, 0 gradients
PyTorch: starting from yolov5s.pt (14.8 MB)
TorchScript: starting export with torch 1.8.1+cpu…
M:\python\OpenCV\yolov5\yolov5-master\models\yolo.py:51: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can’t record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
TorchScript: export success, saved as yolov5s.torchscript.pt (29.4 MB)
ONNX: starting export with onnx 1.10.1…
ONNX: export success, saved as yolov5s.onnx (29.2 MB)
CoreML: starting export with coremltools 4.1…
Tuple detected at graph output. This will be flattened in the converted model.
Converting graph.
Adding op ‘1’ of type const
Adding op ‘2’ of type const
Adding op ‘3’ of type const
…
…
…
Adding op ‘729’ of type add
Converting op 730 : select
Converting Frontend ==> MIL Ops: 87%|██████████████████████████████████████▏ | 604/695 [00:01<00:00, 557.66 ops/s]
CoreML: export failure:
Export complete (10.18s). Visualize with https://github.com/lutzroeder/netron.

在这里插入图片描述
生成的yolov5s.onnx文件在YOLOv5目录下。

ONNX转为为IR中间格式

Windows 10 下 torch模型转换为 OpenVINO需要的IR文件：https://blog.csdn.net/qq_44989881/article/details/119488209

使用 OpenVINO 工具包提供的 mo_onnx.py文件，对模型进行转换。
在这里插入图片描述

管理员模式打开Anaconda，启动虚拟环境的终端，cd进openVINO转换工具目录，执行转换代码：

cd C:\Program Files (x86)\Intel\openvino_2021.2.185\deployment_tools\model_optimizer

python mo_onnx.py --input_model M:\python\OpenCV\yolov5\yolov5-master\yolov5s.onnx

(pytorch) C:\Program Files (x86)\Intel\openvino_2021.2.185\deployment_tools\model_optimizer>python mo_onnx.py --input_model M:\python\OpenCV\yolov5\yolov5-master\yolov5s.onnx
Model Optimizer arguments:
Common parameters:
- Path to the Input Model: M:\python\OpenCV\yolov5\yolov5-master\yolov5s.onnx
- Path for generated IR: C:\Program Files (x86)\Intel\openvino_2021.2.185\deployment_tools\model_optimizer.
- IR output name: yolov5s
- Log level: ERROR
- Batch: Not specified, inherited from the model
- Input layers: Not specified, inherited from the model
- Output layers: Not specified, inherited from the model
- Input shapes: Not specified, inherited from the model
- Mean values: Not specified
- Scale values: Not specified
- Scale factor: Not specified
- Precision of IR: FP32
- Enable fusing: True
- Enable grouped convolutions fusing: True
- Move mean values to preprocess section: None
- Reverse input channels: False
ONNX specific parameters:
Model Optimizer version: 2021.2.0-1877-176bdf51370-releases/2021/2
[ SUCCESS ] Generated IR version 10 model.
[ SUCCESS ] XML file: C:\Program Files (x86)\Intel\openvino_2021.2.185\deployment_tools\model_optimizer.\yolov5s.xml
[ SUCCESS ] BIN file: C:\Program Files (x86)\Intel\openvino_2021.2.185\deployment_tools\model_optimizer.\yolov5s.bin
[ SUCCESS ] Total execution time: 26.05 seconds.
It’s been a while, check for a new version of Intel® Distribution of OpenVINO™ toolkit here https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit/choose-download.html?cid=other&source=Prod&campid=ww_2021_bu_IOTG&content=upg_pro&medium=organic_uid_agjj or on the GitHub*

转换成功后，在转换工具 model_optimizer 目录下生成了bin和xml文件，然后就可以用 OpenVINO部署了。

.xml - 描述网络拓扑
.bin - 包含权重和偏差二进制数据。
在这里插入图片描述

遇到的问题：
缺少 onnx 库和 coremltools库

(Python37) M:\python\OpenCV\yolov5\yolov5-master>python models/export.py --weights yolov5s.pt --img 640 --batch 1
Namespace(batch_size=1, device=‘cpu’, dynamic=False, half=False, img_size=[640, 640], include=[‘torchscript’, ‘onnx’, ‘coreml’], inplace=False, opset_version=12, optimize=False, simplify=False, train=False, weights=‘yolov5s.pt’)
YOLOv5 2021-5-18 torch 1.8.1+cpu CPU
Fusing layers…
Model Summary: 224 layers, 7266973 parameters, 0 gradients
PyTorch: starting from yolov5s.pt (14.8 MB)
TorchScript: starting export with torch 1.8.1+cpu…
M:\python\OpenCV\yolov5\yolov5-master\models\yolo.py:51: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can’t record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
TorchScript: export success, saved as yolov5s.torchscript.pt (29.4 MB)
ONNX: export failure: No module named ‘onnx’
CoreML: export failure: No module named ‘coremltools’
Export complete (5.22s). Visualize with https://github.com/lutzroeder/netron.

安装：onnx

pip install onnx

(python37) C:\Program Files (x86)\Intel\openvino_2021.2.185\bin>pip install onnx
Collecting onnx
Downloading onnx-1.10.1-cp37-cp37m-win_amd64.whl (11.4 MB)
|████████████████████████████████| 11.4 MB 1.1 MB/s
Requirement already satisfied: typing-extensions>=3.6.2.1 in h:\anacondanavigator\anaconda\envs\python37\lib\site-packages (from onnx) (3.7.4.3)
Requirement already satisfied: six in h:\anacondanavigator\anaconda\envs\python37\lib\site-packages (from onnx) (1.16.0)
Requirement already satisfied: numpy>=1.16.6 in h:\anacondanavigator\anaconda\envs\python37\lib\site-packages (from onnx) (1.20.2)
Requirement already satisfied: protobuf in h:\anacondanavigator\anaconda\envs\python37\lib\site-packages (from onnx) (3.17.0)
Installing collected packages: onnx
Successfully installed onnx-1.10.1

安装：coremltools

pip install coremltools

(python37) C:\Program Files (x86)\Intel\openvino_2021.2.185\bin>pip install coremltools
Collecting coremltools
Downloading coremltools-4.1.tar.gz (783 kB)
|████████████████████████████████| 783 kB 726 kB/s
Collecting numpy<1.20,>=1.14.5
Downloading numpy-1.19.5-cp37-cp37m-win_amd64.whl (13.2 MB)
|████████████████████████████████| 13.2 MB 1.3 MB/s
Requirement already satisfied: protobuf>=3.1.0 in h:\anacondanavigator\anaconda\envs\python37\lib\site-packages (from coremltools) (3.17.0)
Requirement already satisfied: six>=1.10.0 in
…
…
…
installed. This behaviour is the source of the following dependency conflicts.
imgaug 0.4.0 requires opencv-python-headless, which is not installed.
labelme 4.5.7 requires matplotlib<3.3, but you have matplotlib 3.4.2 which is incompatible.
Successfully installed attr-0.3.1 coremltools-4.1 mpmath-1.2.1 numpy-1.19.5 sympy-1.8