當前位置：首頁 > 编程资源 > 编程问答 >内容正文

编程问答

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

發布時間：2025/5/22 编程问答 18 豆豆

生活随笔收集整理的這篇文章主要介紹了 YOLOv5的pytorch模型文件转换为ONNX文件小編覺得挺不錯的,現在分享給大家,幫大家做個參考.

YOLOv5

- - - 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

總結

以上是生活随笔為你收集整理的YOLOv5的pytorch模型文件转换为ONNX文件的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。

上一篇： Raspberry Pi 4B安装 Op
下一篇： OpenVINO 部署 YOLOv5 转