Yolov5m.pt转tensorRT

准备自己训练好的yolov5m.pt pytorch模型
在yolov5网址下选择正确的yolov5版本，我用的是yolov5-v5.0。

依照官网上的步骤在同一根目录下git下载yolov5与tensorrt-yolov5，并将自己的yolov5m.py文件放在yolov5文件夹下；
然后从tensorrt-yolov5中复制gen_wts.py文件到yolov5中，再进入到yolov5文件夹中。

cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5

cd {ultralytics}/yolov5

python gen_wts.py -w yolov5m.pt -o yolov5s.wts

然后进入tensorrt-yolov5文件夹下

cd {tensorrtx}/yolov5/

// update CLASS_NUM in yololayer.h if your model is trained on custom dataset （这一步至关重要，如果自己训练的yolov5检测模型只检测人，必须在yololayer.h中修改CLASS_NUM为1，如果还有其他检测类别，则CLASS_NUM为检测类别总数）

mkdir build

cd build

cp {ultralytics}/yolov5/yolov5m.wts {tensorrtx}/yolov5/build   
// 复制生成的yolov5m.wts到tensorrt-yolov5

修改CMakeLists.txt中的tensorrt路径（该路径为自己ubuntu系统上的tensorrt路径）：

# tensorrt
include_directories(/home/XXX/TensorRT-7.2.2.3/include)
link_directories(/home/XXX/TensorRT-7.2.2.3/lib)

然后在build文件夹下开始编译

cmake ..

make

sudo ./yolov5 -s yolov5m.wts yolov5m.engine m  
// serialize model to plan file

sudo ./yolov5 -d yolov5m.engine ../samples  
// deserialize and run inference, the images in [image folder] will be processed.

然后修改yolov5_trt.py中的.engine路径

在tensorrt-yolov5/yolov5文件夹下运行yolov5_trt.py：