yolov5转tensorrt模型

最新推荐文章于 2026-05-02 09:14:09 发布

原创最新推荐文章于 2026-05-02 09:14:09 发布 · 6.8k 阅读

本内容遵循CC 4.0 BY-SA版权协议

收录于

深度学习

该文章详细介绍了如何在Windows 10环境下，使用CPU或GPU搭建Yolov5，并将预训练模型yolov5s转换为TensorRT引擎，以实现推理速度的提升。步骤包括生成.wts文件、编译和运行，最后展示了模型转换后的效果，适合需要优化模型推理速度的读者参考。

Python3.8

Conda

Python

Python 是一种高级、解释型、通用的编程语言，以其简洁易读的语法而闻名，适用于广泛的应用，包括Web开发、数据分析、人工智能和自动化脚本

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
在这里插入图片描述

您可能感兴趣的与本文相关的镜像

Python3.8

Conda

Python

Python 是一种高级、解释型、通用的编程语言，以其简洁易读的语法而闻名，适用于广泛的应用，包括Web开发、数据分析、人工智能和自动化脚本

标签

#python #tensorrt #深度学习