【yolov5】pytorch模型导出为onnx模型

港控/mmm° 2023-10-02 21:51 59阅读 0赞

> 博主想拿官网的yolov5训练好pt模型，然后转换成rknn模型，然后在瑞芯微开发板上调用模型检测。但是官网的版本对npu不友好，所以采用改进结构的版本：
> 
> 1.  将Focus层改成Conv层
> 2.  将Swish激活函数改成Relu激活函数
> 
> 自带的预训练模型是预测80类CoCo数据集的yolov5s改进结构，下面就带大家一起转换模型！

1、首先部署好yolov5的环境，保证可以运行detect.py进行检测，将自己训练好的pt模型放到weights目录下，我这里命名是best.pt。

![在这里插入图片描述][fdb68122d7ed4727a3dca2fed90b4ab6.png_pic_center]

2、`pip install onnx`安装好onnx库

![在这里插入图片描述][3d71e2956fad46ab943e2b4f642fd0f6.png_pic_center]  
3、输入下面命令导出模型（后面的weights、img和batch参数可以不加，设置下默认参数即可）

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

然后会分解网络层，预览模型层数和参数量，最后保存为`./weights/best.onnx`。

(yolov5) dzh@dzh-Lenovo-Legion-Y7000:~/airockchip-yolov5$ python models/export.py --weights ./weights/best.pt --img 640 --batch 1
    Namespace(batch_size=1, img_size=[640, 640], weights='./weights/best.pt')
    Fusing layers... 
    Model Summary: 140 layers, 7.2627e+06 parameters, 0 gradients
    
    Starting TorchScript export with torch 1.12.1+cu102...
    /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages/torch/jit/_trace.py:967: TracerWarning: Encountering a list at the output of the tracer might cause the trace to be incorrect, this is only valid if the container structure does not change based on the module's inputs. Consider using a constant container instead (e.g. for `list`, use a `tuple` instead. for `dict`, use a `NamedTuple` instead). If you absolutely need this and know the side effects, pass strict=False to trace() to allow this behavior.
      module._c._create_method_from_trace(
    TorchScript export success, saved as ./weights/best.torchscript.pt
    
    Starting ONNX export with onnx 1.12.0...
    /home/dzh/airockchip-yolov5/./models/yolo.py:103: 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 augment:
    /home/dzh/airockchip-yolov5/./models/yolo.py:128: 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 profile:
    /home/dzh/airockchip-yolov5/./models/yolo.py:143: 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 profile:
    /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py:621: UserWarning: You are trying to export the model with onnx:Resize for ONNX opset version 10. This operator might cause results to not match the expected results by PyTorch.
    ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).
    We recommend using opset 11 and above for models using this operator.
      warnings.warn(
    ONNX export success, saved as ./weights/best.onnx
    
    Export complete (4.98s). Visualize with https://github.com/lutzroeder/netron.
    (yolov5) dzh@dzh-Lenovo-Legion-Y7000:~/airockchip-yolov5$

4、导出的onnx模型可以用[Netron网站][Netron]查看网络结构：

![在这里插入图片描述][22ff3ab031a2435fa761282e17919cd9.png_pic_center]  
可以看到转换后的结构比较乱，有些结点其实可以跳过，减小模型。这里可以使用[修改onnx模型的工具][onnx]，说明都在Readme中。

![在这里插入图片描述][a5054eaaf45f448181a127b182066ff5.png_pic_center]

这里我们使用深度学习的库onnx-simplifier，通过pip安装即可。

Collecting onnx-simplifier
      Downloading onnx_simplifier-0.4.8-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (2.0 MB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.0/2.0 MB 203.7 kB/s eta 0:00:00
    Requirement already satisfied: onnx in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx-simplifier) (1.12.0)
    Collecting rich
      Downloading rich-12.6.0-py3-none-any.whl (237 kB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 237.5/237.5 kB 191.1 kB/s eta 0:00:00
    Requirement already satisfied: numpy>=1.16.6 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx->onnx-simplifier) (1.23.3)
    Requirement already satisfied: protobuf<=3.20.1,>=3.12.2 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx->onnx-simplifier) (3.19.6)
    Requirement already satisfied: typing-extensions>=3.6.2.1 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnx->onnx-simplifier) (4.4.0)
    Collecting commonmark<0.10.0,>=0.9.0
      Downloading commonmark-0.9.1-py2.py3-none-any.whl (51 kB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 51.1/51.1 kB 215.0 kB/s eta 0:00:00
    Collecting pygments<3.0.0,>=2.6.0
      Downloading Pygments-2.13.0-py3-none-any.whl (1.1 MB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.1/1.1 MB 142.8 kB/s eta 0:00:00
    Installing collected packages: commonmark, pygments, rich, onnx-simplifier
    Successfully installed commonmark-0.9.1 onnx-simplifier-0.4.8 pygments-2.13.0 rich-12.6.0
    (yolov5) dzh@dzh-Lenovo-Legion-Y7000:~/airockchip-yolov5$ python -m onnxsim ./weights/red.onnx ./weights/red2.onnx
    Installing onnxruntime by `/home/dzh/anaconda3/envs/yolov5/bin/python -m pip install --user onnxruntime`, 
    please wait for a moment..
    Collecting onnxruntime
      Downloading onnxruntime-1.12.1-cp38-cp38-manylinux_2_27_x86_64.whl (4.9 MB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.9/4.9 MB 52.8 kB/s eta 0:00:00
    Collecting coloredlogs
      Downloading coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 46.0/46.0 kB 74.5 kB/s eta 0:00:00
    Requirement already satisfied: protobuf in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (3.19.6)
    Collecting sympy
      Downloading sympy-1.11.1-py3-none-any.whl (6.5 MB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.5/6.5 MB 72.3 kB/s eta 0:00:00
    Requirement already satisfied: packaging in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (21.3)
    Requirement already satisfied: numpy>=1.21.0 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from onnxruntime) (1.23.3)
    Collecting flatbuffers
      Downloading flatbuffers-22.9.24-py2.py3-none-any.whl (26 kB)
    Collecting humanfriendly>=9.1
      Downloading humanfriendly-10.0-py2.py3-none-any.whl (86 kB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 86.8/86.8 kB 129.3 kB/s eta 0:00:00
    Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /home/dzh/anaconda3/envs/yolov5/lib/python3.8/site-packages (from packaging->onnxruntime) (3.0.9)
    Collecting mpmath>=0.19
      Downloading mpmath-1.2.1-py3-none-any.whl (532 kB)
         ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 532.6/532.6 kB 128.3 kB/s eta 0:00:00
    Installing collected packages: mpmath, flatbuffers, sympy, humanfriendly, coloredlogs, onnxruntime

然后输入`python -m onnxsim ./weights/red.onnx ./weights/red2.onnx`优化onnx模型。可以看到优化前后的表。

![在这里插入图片描述][74479216aecc46cb9df0045fe15ff474.png_pic_center]

--------------------

总结可能遇到的问题：

<table> 
 <thead> 
  <tr> 
   <th>ONNX export failure: No module named ‘onnx’</th> 
  </tr> 
 </thead> 
 <tbody> 
  <tr> 
   <td>原因：没有安装onnx库导致找不到该模块</td> 
  </tr> 
 </tbody> 
</table>

<table> 
 <thead> 
  <tr> 
   <th>RuntimeError: Given groups=1, weight of size [32, 12, 3, 3], expected input[1, 3, 640, 640] to have 12 channels, but got 3 channels instead</th> 
  </tr> 
 </thead> 
 <tbody> 
  <tr> 
   <td>模型预期的图像通道是12，但是现在输入是3个通道的图像，所以在图片输入之前改成模型需要的通道数</td> 
  </tr> 
 </tbody> 
</table>

[fdb68122d7ed4727a3dca2fed90b4ab6.png_pic_center]: https://img-blog.csdnimg.cn/fdb68122d7ed4727a3dca2fed90b4ab6.png#pic_center
[3d71e2956fad46ab943e2b4f642fd0f6.png_pic_center]: https://img-blog.csdnimg.cn/3d71e2956fad46ab943e2b4f642fd0f6.png#pic_center
[Netron]: https://netron.app/
[22ff3ab031a2435fa761282e17919cd9.png_pic_center]: https://img-blog.csdnimg.cn/22ff3ab031a2435fa761282e17919cd9.png#pic_center
[onnx]: https://github.com/saurabh-shandilya/onnx-utils
[a5054eaaf45f448181a127b182066ff5.png_pic_center]: https://img-blog.csdnimg.cn/a5054eaaf45f448181a127b182066ff5.png#pic_center
[74479216aecc46cb9df0045fe15ff474.png_pic_center]: https://img-blog.csdnimg.cn/74479216aecc46cb9df0045fe15ff474.png#pic_center