YOLOv5:TensorRT加速YOLOv5模型推理

前言

前提条件

相关介绍

  • Python是一种跨平台的计算机程序设计语言。是一个高层次的结合了解释性、编译性、互动性和面向对象的脚本语言。最初被设计用于编写自动化脚本(shell),随着版本的不断更新和语言新功能的添加,越多被用于独立的、大型项目的开发。
  • PyTorch 是一个深度学习框架,封装好了很多网络和深度学习相关的工具方便我们调用,而不用我们一个个去单独写了。它分为 CPU 和 GPU 版本,其他框架还有 TensorFlow、Caffe 等。PyTorch 是由 Facebook 人工智能研究院(FAIR)基于 Torch 推出的,它是一个基于 Python 的可续计算包,提供两个高级功能:1、具有强大的 GPU 加速的张量计算(如 NumPy);2、构建深度神经网络时的自动微分机制。
  • YOLOv5是一种单阶段目标检测算法,该算法在YOLOv4的基础上添加了一些新的改进思路,使其速度与精度都得到了极大的性能提升。它是一个在COCO数据集上预训练的物体检测架构和模型系列,代表了Ultralytics对未来视觉AI方法的开源研究,其中包含了经过数千小时的研究和开发而形成的经验教训和最佳实践。

TensorRT加速YOLOv5模型推理

YOLOv5项目官方源地址

git clone https://github.com/ultralytics/yolov5.git

将训练好的YOLOv5模型权重转换成TensorRT引擎

在这里插入图片描述

python export.py --imgsz 1024 --weights ./runs/train/exp/weights/best.pt --include engine --device 0

在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

YOLOv5 best.pt推理测试

python3 detect.py --imgsz 1024 --data data/my.yaml --source ./imgs --weights runs/train/exp/weights/best.pt 

在这里插入图片描述

TensorRT Engine推理测试

python3 detect.py --imgsz 1024 --data data/my.yaml --source ./imgs --weights runs/train/exp/weights/best.engine

在这里插入图片描述

小结

  • 可能是因为RTX 3070性能比较高,所以速度差别不大。
  • 从结果上看,所需时间,从0.009s(0.008s)提升到了0.006s(0.006s)。

参考

[1] https://github.com/ultralytics/yolov5

猜你喜欢

转载自blog.csdn.net/FriendshipTang/article/details/131023963