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模型测试或者你自己修改也可以】


目录
相关文章
|
8月前
|
存储 计算机视觉 异构计算
使用python&C++对bubbliiiing的yolo系列进行opencv.dnn进行推理部署
使用python&C++对bubbliiiing的yolo系列进行opencv.dnn进行推理部署
102 0
|
8月前
|
数据可视化 Python
Python用 PyMC3 贝叶斯推理案例研究:抛硬币和保险索赔发生结果可视化
Python用 PyMC3 贝叶斯推理案例研究:抛硬币和保险索赔发生结果可视化
|
人工智能 知识图谱 Python
python实现知识推理,图可达,pydatalog,kanren,sympy
python实现知识推理,图可达,pydatalog,kanren,sympy
178 0
|
机器学习/深度学习 Python
【Python机器学习】实验06 贝叶斯推理3
【Python机器学习】实验06 贝叶斯推理
81 0
|
机器学习/深度学习 数据可视化 API
Paddle Lite是什么,快速上手Python推理,pdmodel使用
“Paddle Lite 是一组工具,可帮助开发者在移动设备、嵌入式设备和 loT 设备上运行模型,以便实现设备端机器学习。”
453 0
|
1月前
|
人工智能 数据可视化 数据挖掘
探索Python编程:从基础到高级
在这篇文章中,我们将一起深入探索Python编程的世界。无论你是初学者还是有经验的程序员,都可以从中获得新的知识和技能。我们将从Python的基础语法开始,然后逐步过渡到更复杂的主题,如面向对象编程、异常处理和模块使用。最后,我们将通过一些实际的代码示例,来展示如何应用这些知识解决实际问题。让我们一起开启Python编程的旅程吧!
|
30天前
|
存储 数据采集 人工智能
Python编程入门:从零基础到实战应用
本文是一篇面向初学者的Python编程教程,旨在帮助读者从零开始学习Python编程语言。文章首先介绍了Python的基本概念和特点,然后通过一个简单的例子展示了如何编写Python代码。接下来,文章详细介绍了Python的数据类型、变量、运算符、控制结构、函数等基本语法知识。最后,文章通过一个实战项目——制作一个简单的计算器程序,帮助读者巩固所学知识并提高编程技能。
|
19天前
|
Unix Linux 程序员
[oeasy]python053_学编程为什么从hello_world_开始
视频介绍了“Hello World”程序的由来及其在编程中的重要性。从贝尔实验室诞生的Unix系统和C语言说起,讲述了“Hello World”作为经典示例的起源和流传过程。文章还探讨了C语言对其他编程语言的影响,以及它在系统编程中的地位。最后总结了“Hello World”、print、小括号和双引号等编程概念的来源。
103 80
|
2月前
|
存储 索引 Python
Python编程数据结构的深入理解
深入理解 Python 中的数据结构是提高编程能力的重要途径。通过合理选择和使用数据结构,可以提高程序的效率和质量
152 59
|
8天前
|
Python
[oeasy]python055_python编程_容易出现的问题_函数名的重新赋值_print_int
本文介绍了Python编程中容易出现的问题,特别是函数名、类名和模块名的重新赋值。通过具体示例展示了将内建函数(如`print`、`int`、`max`)或模块名(如`os`)重新赋值为其他类型后,会导致原有功能失效。例如,将`print`赋值为整数后,无法再用其输出内容;将`int`赋值为整数后,无法再进行类型转换。重新赋值后,这些名称失去了原有的功能,可能导致程序错误。总结指出,已有的函数名、类名和模块名不适合覆盖赋新值,否则会失去原有功能。如果需要使用类似的变量名,建议采用其他命名方式以避免冲突。
29 14