手把手教程 | 如何用TensorRT部署YOLOv7

作者 | 一天到晚潜水的鱼  编辑 | 汽车人

点击下方卡片，关注“自动驾驶之心”公众号

ADAS巨卷干货，即可获取

点击进入→自动驾驶之心【模型部署】技术交流群

后台回复【模型部署工程】获取基于TensorRT的分类、检测任务的部署源码！

前言

YOLOv7已经出来一段时间了，之前一直没有时间去看。最近抽空学习了一下，不禁感慨一下，现在的YOLO系列目标检测算法真是卷成麻花了，各种算法层出不穷，让人目不暇接。最重要的是，我感觉自己学不动了

之前YOLOX和YOLOv6出来后，我分别写了文章介绍如何使用TensorRT的C++接口部署这两个模型，并且把代码放在了gitee上。

本文将介绍如何用TensorRT的C++接口部署YOLOv7。

实现过程

「1. 下载PyTorch模型」

YOLOv7的PyTorch模型可以从官方GitHub仓库的Releases中下载，目前的版本号为v0.1。下面的操作是下载yolov7.pt模型:

wget -c https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt

「2. 转换为ONNX模型」

下载好PyTorch格式的模型文件后，可以用官方提供的脚本转换为ONNX模型：

python export.py --weights ./yolov7.pt --grid --simplify --img-size 640 640

官方的脚本还提供把NMS操作集成到ONNX模型中的功能，我个人不太想这么做，感兴趣的读者可以试一下。

上面的操作导出的ONNX模型已经把后处理操作集成到了模型中，下图红框中的操作就是把模型三个检测头输出的内容经过处理后再汇总到一起作为输出，这样做的好处是可以避免自己去写复杂的后处理代码。

「3. 基于TensorRT部署ONNX模型」

如何使用TensorRT的C++接口部署ONNX模型可以参考我之前给【自动驾驶之心】公众号写的这篇文章：手把手教学！TensorRT部署实战：YOLOv5的ONNX模型部署。

这篇文章以YOLOv5为例，通过大量的代码一步步讲解如何使用TensorRT框架部署ONNX模型，主要目的是希望读者能够通过这篇文章学习到TensorRT模型部署的基本流程，比如如何准备输入数据、如何调用API用模型做推理、如何解析模型的输出结果。相信看完这篇文章的读者 已经基本掌握使用TensorRT部署ONNX模型的基本方法，并且能做到举一反三自己部署其他模型了，比如YOLOX、YOLOv6、YOLOv7等等。

我就来说一下重点吧，「重点就是这篇文章的代码可以直接拿来部署YOLOv7！」

既然这样，本文就不再重复叙述了。如果还没有看这篇文章的话，那就赶紧点击上面的链接查看吧。

结果

用yolov7.onnx模型测试的几个结果如下:

下面是跑视频的结果：

在我的GeForce GTX 1050 Ti显卡（只支持FP32）上测试各个模型的耗时如下表所示：

模型	输入尺寸	耗时
yolov7-tiny.onnx	640x640	8 ms
yolov7.onnx	640x640	66 ms
yolov7x.onnx	640x640	113 ms
yolov7-w6.onnx	640x640	56 ms
yolov7-e6.onnx	640x640	83 ms
yolov7-d6.onnx	640x640	108 ms
yolov7-e6e.onnx	640x640	126ms

换到装了GeForce RTX 3070显卡（支持FP32和FP16）的服务器上去试一下：

模型	输入尺寸	模型精度	耗时
yolov7-tiny.onnx	640x640	FP32	2 ms
yolov7.onnx	640x640	FP32	13 ms
yolov7x.onnx	640x640	FP32	20 ms
yolov7-w6.onnx	640x640	FP32	12 ms
yolov7-e6.onnx	640x640	FP32	17 ms
yolov7-d6.onnx	640x640	FP32	21 ms
yolov7-e6e.onnx	640x640	FP32	25ms
yolov7-tiny.onnx	640x640	FP16	1 ms
yolov7.onnx	640x640	FP16	6 ms
yolov7x.onnx	640x640	FP16	9 ms
yolov7-w6.onnx	640x640	FP16	6 ms
yolov7-e6.onnx	640x640	FP16	7 ms
yolov7-d6.onnx	640x640	FP16	9 ms
yolov7-e6e.onnx	640x640	FP16	10ms

果然在高级显卡上就是跑得快呀，等我有钱了一定要去换个好点的显卡！

总结

没啥好总结的，对于这种优秀的开源工作，我只能用两个字来评价：真香！

【自动驾驶之心】全栈技术交流群

自动驾驶之心是首个自动驾驶开发者社区，聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、BEV感知、多传感器融合、SLAM、光流估计、深度估计、轨迹预测、高精地图、规划控制、模型部署落地、自动驾驶仿真测试、硬件配置、AI求职交流等方向；

添加汽车人助理微信邀请入群

备注：学校/公司+方向+昵称

自动驾驶之心【知识星球】

想要了解更多自动驾驶感知（分类、检测、分割、关键点、车道线、3D目标检测、多传感器融合、目标跟踪、光流估计、轨迹预测）、自动驾驶定位建图（SLAM、高精地图）、自动驾驶规划控制、领域技术方案、AI模型部署落地实战、行业动态、岗位发布，欢迎扫描下方二维码，加入自动驾驶之心知识星球（三天内无条件退款），日常分享论文+代码，这里汇聚行业和学术界大佬，前沿技术方向尽在掌握中，期待交流！