yolov5 model to tensorRT

Before that, verify that CUDA, CUDNN, TensorRT, and OpenCV are all installed correctly.
CUDA, CUDNN, TensorRT installation and verification reference: https://blog.csdn.net/long630576366/article/details/125039294?spm=1001.2014.3001.5501

OpenCV installation reference: https://blog.csdn.net/long630576366/article/details/124980126?spm=1001.2014.3001.5501

1. Code and model download

tensorrtx only supports the v5 version of yolov5, so the version should correspond. And the yolov5 code version must be the same as the yolov5s.pt model version, otherwise there will be errors, and the v5 version is used here.

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  #运行生成容器

2. Model conversion

2.1 Modify CMakeLists.txt

In the /tensorrtx/yolov5/ directory: gedit CMakeLists.txt
Modify the path of calling tensorRT to the path of your own installation.
insert image description here

2.1 Compile and convert

If you train your own model, modify the number of categories in 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

2.3 Possible problems:

1. Failed to generate engine:
(1) what(): driver error :
tensorRT version is incorrect, or the NVIDIA graphics card driver does not match the graphics card model.
(2) Segmentation fault :
Modify USE_FP16 to USE_FP32 in yolo5.cpp to solve.

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