yolov5系列文章目录

前言

部署时，为了提高推理效率，通常不直接使用训练生成的pt模型，而是把pt模型转换成可以通过tensorrt加速推理的engin文件。

本文以yolov5s v5.0版本为例，给同学们做演示。
下载源码，可以使用git，当然也可以下载压缩包解压，本文使用git方法：

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

下载模型yolov5s.pt模型：
https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
这里也可以使用自己训练好的pt模型文件，本文我将采用官方模型做转换。

把{tensorrtx}/yolov5/目录下gen_wts.py 复制到{ultralytics}/yolov5目录下：

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

cd到{ultralytics}/yolov5目录下:

cd {
    
    ultralytics}/yolov5

运行：

python gen_wts.py yolov5s.pt

将会生成一个yolov5s.wts文件在这里插入图片描述
若使用自己的模型，只需把上述代码换成：python gen_wts.py xxx.pt

cd {
    
    tensorrtx}/yolov5/

如果采用的是自己训练的模型，则需要更新 yololayer.h 中的 CLASS_NUM 在这里插入图片描述
如需要调整batch_size，需要在yolov5.cpp中进行设置（默认值为1）：

mkdir build
cd build
cp {
    
    ultralytics}/yolov5/yolov5s.wts {
    
    tensorrtx}/yolov5/build

cmake ..

在这里插入图片描述

make

在这里插入图片描述

sudo ./yolov5 -s yolov5s.wts yolov5s.engine s

在这里插入图片描述
至此，build目录下就会生成libmyplugins.so与yolov5s.engine文件：

sudo ./yolov5 -d yolov5s.engine ../images

这里自己修改测试图片images的路径
在这里插入图片描述
此时，便会生成inference后的图片：

在这里插入图片描述

采用trt模型推理时间相比于pt模型肯定是减少的，后续有机会测一下。
参考文档：https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5

如果阅读本文对你有用，欢迎一键三连呀！！！
2021年8月20日09:23:03
在这里插入图片描述