yolov5 | 移动端部署yolov5s模型

移动端的部署有这么几条路：
（以yolov5s.pt模型为例）

pt文件 --> onnx文件/torchscript文件 --> ncnn --> 安卓端部署（android studio编写）
pt文件 --> onnx文件/torchscript文件 --> ML文件 --> ios端部署（需要mac系统运行xcode编写）
pt文件 --> onnx文件 /torchscript文件/wts文件–> openvino/tensorrt 等部署平台 --> 树莓派/jetson部署（最好是在linux环境下）

目前yolov5的部署可以说是把onnx文件/torchscript文件作为中间模型去转换到相应的模型。
部署需要一定的c++操作水平及工具

参考大佬项目

https://github.com/enazoe/yolo-tensorrt作者封装了模型的转换
https://github.com/wang-xinyu/tensorrtx
https://github.com/jkjung-avt/tensorrt_demos

1、pt转ONNX、ML、TorchScript

安装ONNX

在yolov5文件夹下运行

pip install -r requirements.txt onnx # install requirements.txt and ONNX

转换文件

在这里，我以转换官方的yolov5s.pt模型文件为例
在yolov5文件夹下运行

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

应该会多出现三个文件
分别是ONNX文件、ML文件、torchscript文件

使用Netron
使用Netron来查看你的ONNX文件模型
其实netron现已支持大部分格式的模型文件，都是可以查看的。
PaddlePaddle、OpenVINO、TensorFlow 、Caffe…
安装netron

pip install netron

进入python，运行netron.start()

import netron
netron.start('yolov5s.onnx')

2.Jetson nano的部署

下载yolov5、tensorrtx

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git

生成trt引擎

yolov5官方教程：https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

1.生成wts文件

这一步可以在x86系统上完成。
首先你得有yolov5的模型文件，我上一篇有讲，我这里使用yolov5s.pt训练出来的best.pt文件来进行下面的操作。
复制 tensorrtx/yolov5/gen_wts.py文件到 yolov5 文件夹中
5.0的tensorrtx修改了这一python文件
只需要输入一次文件路径即可生成对应文件。

python3 gen_wts.py weights/yolov5s.pt

这样就会在weights文件夹中生成一个yolov5s.wts文件。

2.build

这一步是生成部署引擎，必须来到部署该引擎的硬件设备上进行。
就像我这里模型是想在Jetson nano上部署，所以这一步就必须在Jetson nano上进行build。在此之前的有关yolov5的操作都可在x86的系统上完成。
将yolov5s.wts文件放到tensorrtx/yolov5文件夹中，
来到tensorrtx/yolov5文件夹中，查看yolov5.cpp中的模型是s。

在yolov5.cpp文件中
还可以修改fp16还是fp32 int8、device(选择哪一个GPU设备)、nms_thresh(nms的阈值)、conf_thresh(conf的置信度)、batch_size(批次大小)

打开yololayer.h文件，

修改num总数，根据你训练模型的类个数来
修改输入图片的尺寸，但必须是32的倍数。缩小输入尺寸可以一定程度上加快推理速度。

确保tensorrtx/yolov5文件下有你的wts文件，并且相应地修改了yolov5.cpp文件和yolovlayer.h文件。
然后在tensorrtx/yolov5文件夹下依次执行下面的代码

mkdir build
cd build
cmake ..
make -j6
sudo ./yolov5 -s ../yolov5s.wts yolov5s.engine s
#sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw]
#模型引擎生成的命令解释如上

如果你是通过修改每一层卷积的宽度和深度来达到自定义的模型的话，就可以直接使用下面的命令来达成生成引擎，这是我认为4.0的这个版本最大的升级点了！！！

sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25

我一般缩小宽度和深度来达到剪裁模型，加快运行速度的作用。这一api优化让我的调试十分的方便。
v5.0的tensorrtx库更新适配了yolov5的p6模型，如果你用的是p6模型训练，可以使用下面的命令。

sudo ./yolov5 -s ../yolov5s.wts yolov5s.engine s6
sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c6 0.17 0.25

3.python调用

编辑tensorrtx/yolov5_trt.py文件

注释掉import torch和import torchvision，因为在nano上安装这两个库是有些麻烦的特别是torchvision。
当然你也可以去试试安装这两个库，网上的一些教程我有试过是可以成功，就是麻烦。

INPUT_W和INPUT_H根据你自己之前设置的来，我build时是写608，所以这里也是608。
CONF_THRESH 与 IOU_THRESHOLD 可根据最后的显示效果返回来修改他，暂且不调整。

详见：https://blog.csdn.net/weixin_45569617/article/details/108145046