yolov5模型 转 tensorRT_yolov5转tensorrt

在这之前，先验证 CUDA 、 CUDNN 、 TensorRT、 OpenCV是否都正确安装。
CUDA 、 CUDNN 、 TensorRT 安装及验证参考：https://blog.csdn.net/long630576366/article/details/125039294?spm=1001.2014.3001.5501

OpenCV 安装参考：https://blog.csdn.net/long630576366/article/details/124980126?spm=1001.2014.3001.5501

1、代码、模型下载

tensorrtx 只支持到yolov5 的 v5 版本，所以版本要对应。而且yolov5 代码版本和 yolov5s.pt 模型版本要一致，否则肯定会出错，这里用 v5 版本。

git clone -b v5.0 https://github.com/ultralytics/yolov5.git  #拷贝代码
git clone https://github.com/wang-xinyu/tensorrtx.git  #拷贝代码  已经手动下载了的不用管
wget https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5  #将tensorrtx对应的py文件，复制到YOLO训练模型的项目中去
cd {ultralytics}/yolov5  #进入YOLO项目
python gen_wts.py -w yolov5s.pt -o yolov5s.wts  #运行生成容器
1
2
3
4
5
6

2、模型转换

2.1 修改 CMakeLists.txt

在 /tensorrtx/yolov5/ 目录下: gedit CMakeLists.txt
修改调用 tensorRT 的路径为自己安装的路径。

2.1 编译、转换

如果训练自己的模型， yololayer.h 中修改类别个数。

cd {tensorrtx}/yolov5/  #进入tensorrt
// update CLASS_NUM in yololayer.h if your model is trained on custom dataset
mkdir build  #建立新文件夹
cd build  # 进入建立的文件夹
cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/build  # 将前面转换的wts复制过来
cmake ..  # 编译
make   # 转换
sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw]  // serialize model to plan file  #运行转换的文件生成对应的engine容器  sudo ./yolov5 -s yolov5s.wts yolov5s.engine s  后面的s对应预训练模型的yolov5s.pt 所以前期要确定你的是哪个模型，对应的是v5的那个版本，这个版本一定要对应，否则转换会出错
sudo ./yolov5 -d [.engine] [image folder]  // deserialize and run inference, the images in [image folder] will be processed.
// For example yolov5s
sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
sudo ./yolov5 -d yolov5s.engine ../samples
// For example Custom model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml
sudo ./yolov5 -s yolov5_custom.wts yolov5.engine c 0.17 0.25
sudo ./yolov5 -d yolov5.engine ../samples
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

2.3 可能遇到的问题：

1、生成engine失败：
（1）what(): driver error:
tensorRT 版本不对，或者NVIDIA 显卡驱动与显卡型号不匹配。
（2）segmentation fault:
在yolo5.cpp 中将USE_FP16 修改成USE_FP32 解决。

//#define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
#define USE_FP32  // set USE_INT8 or USE_FP16 or USE_FP32
#define DEVICE 0  // GPU id
1
2
3