YOLOv4 tensorrt推理 python版【附代码】

简介: 笔记

学了几天的tensorRT,又经过了几天的努力终于实现了YOLOv4 tensorRT推理,这篇文章将把这些成果开源出来,供大家免费使用。


YOLOv4代码我采用的是b站up主Bubbliiiing,相信大家应该都比较熟悉这位大佬。


关于trt的推理部分我是参考了官方YOLOV5 6.2版本。

剪枝以后的模型也可以用tensorrt进行推理【理论是可行的,虽然我还没试】


环境说明


windows10


cuda10.2


cudnn8.2.1


pytorch1.7


tensorrt8.2.5.1


python 3.7


显卡:NVIDIA 1650 4G(比较拉跨)


注:linux下我还没有试,可能有些代码需要改,而且trt的版本也会受影响


注意python版本的在下载完tensorrt后的文件夹中有个python文件夹,然后用pip安装【注意:TRT应该是8.2X以上可以安装python般的,其他的是C++】


其他tensorrt的学习资料可以看我tensorrt专栏。


torch2onnx


改功能的实现在torch2onnx.py中。


你需要修改的地方有:


simplity:是否开启simplity模式


output_path: # 输出onnx路径


num_classes: # 类的数量


ckpt_path: # torch权重路径


input_shape: # 输入网络图像尺寸大小


python torch2onnx.py


转换过程如下:

graph(%images : Float(1:1108992, 3:369664, 608:608, 608:1, requires_grad=0, device=cpu),

     %yolo_head3.1.weight : Float(255:256, 256:1, 1:1, 1:1, requires_grad=1, device=cpu),

     %yolo_head3.1.bias : Float(255:1, requires_grad=1, device=cpu),

     %yolo_head2.1.weight : Float(255:512, 512:1, ...

```


可以通过Netron工具对onnx可视化


onnx2engine


改功能的实现在onnx2trt.py中。


你需要修改的地方有:


onnx_path: # onnx文件路径


engine_path: # 输出engine路径


python onnx2trt.py

生成engine的时间比较长,在我电脑上大概用了10分钟左右,如果用trtexec也可以生成engine,这个构建时间会快很多。


推理功能在predict.py中。


参数说明:

--weights: # 权重路径

--img: # 开启图像预测

--video: # 开启视频预测

--video_path: # 视频路径

--fps: # FPS测试

--onnx: # 开启onnx推理

--engine: # 开启trt预测

--input_shape: # 输入大小,默认608 * 608

--conf: # 置信度阈值

--nms: # NMS阈值


torch推理


图像推理:

python predict.py --weights model_data/yolo4_weights.pth --img --conf 0.5 --nms 0.4


视频推理:

python predict.py --weights model_data/yolo4_weights.pth --video --video_path 0


FPS测试:

python predict.py --weights model_data/yolo4_weights.pth --fps

onnx推理


图像推理:


python predict.py --weights model_data/yolov4.onnx --img --conf 0.5 --nms 0.4

视频推理:


python predict.py --weights model_data/yolov4.onnx --video --video_path 0

FPS测试:


python predict.py --weights model_data/yolov4.onnx --fps


engine推理


图像推理:


python predict.py --weights model_data/yolov4.engine --img --conf 0.5 --nms 0.4

视频推理:


python predict.py --weights model_data/yolov4.engine --video --video_path 0

FPS测试:


python predict.py --weights model_data/yolov4.engine --fps


检测结果:

1.jpeg


代码:


git clone https://github.com/YINYIPENG-EN/YOLOv4_tensorRT_pytorch.git



说明:本项目暂时未添加动态输入


所遇问题:


如果报下面的错误:


一般有这么几种情况:


1.cuda10.2需要打两个补丁(官网就有)


2.如果打了补丁还报错,一种是可能遇到了不支持的算子,一种是内存不够。


我trt8.2x版本中发现如果有包含全连接层的onnx不能转成engine,即便用trrexec强行转也会报这种错误,但如果我把全连接层用卷积层代替就没问题,我不知道更高版本的会不会解决这个问题。

[10/21/2022-18:32:29] [TRT] [E] 2:[ltWrapper.cpp::nvinfer1::rt::CublasLtWrapper::setupHeuristic::334] Error C
ode 2: Internal Error (Assertion cublasStatus == CUBLAS_STATUS_SUCCESS failed. )

刚开始在用trt推理的时候发现准确率极差,但onnx和torch都很正常,于是我将预测代码中的图像处理由PIL度图像方式改为了opencv后有非常大的改善,但出现慢屏的框,于是通过分析onnx和engine推理的ouputs发现engine将yolo 的三个head顺序颠倒了,这个问题不知道为什么,重新调整顺序后就可以正常检测了【不过在测试FPS和map部分我没修改图像处理代码,因为FPS就只是算个model推理时间而已不用管他推理结果好坏,而测试map部分还是建议大家用训练torch模型测试或者你自己修改也可以】


目录
相关文章
|
7月前
|
存储 计算机视觉 异构计算
使用python&C++对bubbliiiing的yolo系列进行opencv.dnn进行推理部署
使用python&C++对bubbliiiing的yolo系列进行opencv.dnn进行推理部署
92 0
|
7月前
|
数据可视化 Python
Python用 PyMC3 贝叶斯推理案例研究:抛硬币和保险索赔发生结果可视化
Python用 PyMC3 贝叶斯推理案例研究:抛硬币和保险索赔发生结果可视化
|
人工智能 知识图谱 Python
python实现知识推理,图可达,pydatalog,kanren,sympy
python实现知识推理,图可达,pydatalog,kanren,sympy
165 0
|
机器学习/深度学习 Python
【Python机器学习】实验06 贝叶斯推理3
【Python机器学习】实验06 贝叶斯推理
75 0
|
机器学习/深度学习 数据可视化 API
Paddle Lite是什么,快速上手Python推理,pdmodel使用
“Paddle Lite 是一组工具,可帮助开发者在移动设备、嵌入式设备和 loT 设备上运行模型,以便实现设备端机器学习。”
445 0
|
17天前
|
存储 数据挖掘 开发者
Python编程入门:从零到英雄
在这篇文章中,我们将一起踏上Python编程的奇幻之旅。无论你是编程新手,还是希望拓展技能的开发者,本教程都将为你提供一条清晰的道路,引导你从基础语法走向实际应用。通过精心设计的代码示例和练习,你将学会如何用Python解决实际问题,并准备好迎接更复杂的编程挑战。让我们一起探索这个强大的语言,开启你的编程生涯吧!
|
23天前
|
机器学习/深度学习 人工智能 TensorFlow
人工智能浪潮下的自我修养:从Python编程入门到深度学习实践
【10月更文挑战第39天】本文旨在为初学者提供一条清晰的道路,从Python基础语法的掌握到深度学习领域的探索。我们将通过简明扼要的语言和实际代码示例,引导读者逐步构建起对人工智能技术的理解和应用能力。文章不仅涵盖Python编程的基础,还将深入探讨深度学习的核心概念、工具和实战技巧,帮助读者在AI的浪潮中找到自己的位置。
|
23天前
|
机器学习/深度学习 数据挖掘 Python
Python编程入门——从零开始构建你的第一个程序
【10月更文挑战第39天】本文将带你走进Python的世界,通过简单易懂的语言和实际的代码示例,让你快速掌握Python的基础语法。无论你是编程新手还是想学习新语言的老手,这篇文章都能为你提供有价值的信息。我们将从变量、数据类型、控制结构等基本概念入手,逐步过渡到函数、模块等高级特性,最后通过一个综合示例来巩固所学知识。让我们一起开启Python编程之旅吧!
|
23天前
|
存储 Python
Python编程入门:打造你的第一个程序
【10月更文挑战第39天】在数字时代的浪潮中,掌握编程技能如同掌握了一门新时代的语言。本文将引导你步入Python编程的奇妙世界,从零基础出发,一步步构建你的第一个程序。我们将探索编程的基本概念,通过简单示例理解变量、数据类型和控制结构,最终实现一个简单的猜数字游戏。这不仅是一段代码的旅程,更是逻辑思维和问题解决能力的锻炼之旅。准备好了吗?让我们开始吧!
|
10天前
|
Python
Python编程入门:从零开始的代码旅程
本文是一篇针对Python编程初学者的入门指南,将介绍Python的基本语法、数据类型、控制结构以及函数等概念。文章旨在帮助读者快速掌握Python编程的基础知识,并能够编写简单的Python程序。通过本文的学习,读者将能够理解Python代码的基本结构和逻辑,为进一步深入学习打下坚实的基础。