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

博主想拿官网的yolov5训练好pt模型，然后转换成rknn模型，然后在瑞芯微开发板上调用模型检测。但是官网的版本对npu不友好，所以采用改进结构的版本：

1. 将Focus层改成Conv层
2. 将Swish激活函数改成Relu激活函数

自带的预训练模型是预测80类CoCo数据集的yolov5s改进结构，下面就带大家一起转换模型！

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

2、pip install onnx安装好onnx库

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网站查看网络结构：

可以看到转换后的结构比较乱，有些结点其实可以跳过，减小模型。这里可以使用修改onnx模型的工具，说明都在Readme中。

这里我们使用深度学习的库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模型。可以看到优化前后的表。

总结可能遇到的问题：