YOLOv4 tensorrt推理 python版【附代码】

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

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

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

有关YOLOv4的libtorch部署可以参考我另一篇文章：YOLOV4 libtorch

有关YOLOv4剪枝参考我另一篇文章： YOLOV4剪枝

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

---

环境说明

windows10

cuda10.2

cudnn8.2.1

pytorch1.7

tensorrt8.2.5.1

python 3.7

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

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

---

  tensorrt安装可以看我另一篇文章：tensorrt安装与yolov5部署C++

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

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

---

目录

torch2onnx

torch推理

onnx推理

engine推理

检测结果：

代码：

权重：

所遇问题：

---

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

检测结果：

---

代码：

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

权重：

链接：https://pan.baidu.com/s/1mlBMuiKqH3glZPqFGunXmw 
提取码：yypn

项目新增加了语音报警功能【看v2.0版本即可】

v2.0下载：

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

项目新增：目标跟踪(利用iou进行跟踪)，可选用diou。看v3.0版本

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

 运行(可结合其他功能使用)：

python predict.py --weights [weight path] --video --video_path 0 --engine --track --track_class

 

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

所遇问题：

如果报下面的错误：

一般有这么几种情况：

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