TensorRT来加速YOLO v5推理与检测
文章目录
- TensorRT来加速YOLO v5推理与检测
- TensorRT简介
- 一、TensorRT安装
- 1. 电脑基础环境
- 2. 查看 cuda 版本
- 3. 安装TensorRT
- 4. 验证TensorRT
- 二、YOLO v5模型导出函数
- 1.onnx 模型导出(def export_onnx()函数)
- 2.openvino模型导出(def export_openvino()函数)
- 3.coreml模型导出(def export_coreml()函数)
- 4.TensorRT模型导出(def export_engine()函数)
- 三、YOLO v5使用TensorRT进行加速检测
- 1. 模型导出
- 2. 图片检测
- 3. 使用yolov5s.pt 和 yolov5s.engine(tensorrt)进行检测
- 3.1 使用yolov5s.pt, 推理时间: 5.49s
- 3.2 使用 yolov5s.engine(tensorrt), 推理时间: 4.20s
- 四、总结
TensorRT简介
TensorRT是一个高性能的深度学习推理(Inference)优化器,可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架,将TensorRT和NVIDIA的GPU结合起来,能在几乎所有的框架中进行快速和高效的部署推理。
TensorRT 是一个C++库,从 TensorRT 3 开始提供C++ API和Python API,主要用来针对 NVIDIA GPU进行 高性能推理(Inference)加速。
基本特性:
用于高效实现已训练好的深度学习模型的推理过程的SDK
内含推理优化器和运行时环境
使DL模型能以更高吞吐量和更低的延迟运行
有C++和python的API,完全等价可以混用
TensorRT主要做了以下几点来加快推理速度:
算子融合
量化
动态张量显存
内核自动调整
多流执行
一、TensorRT安装
1. 电脑基础环境
- 操作系统: win10
- 显卡: 4060TI 16GB
- pytorch 版本:1.13.0 使用 print(torch._ _version_ _) 查看 pytorch 版本
- python: 3.8
2. 查看 cuda 版本
- 激活 conda 中的一个环境, 并输入 nvcc -V 进行查看, 我的版本如下图所示
3. 安装TensorRT
需要到nvidia官网去下载对应cuda版本的TensorRT安装包,我这里下载对应版本压缩包如下,下载解压即可。
下载好将压缩包进行解压,激活 conda 环境, 然后 cd 进入下面几个文件夹,安装 whl 文件。
- 将 解压后的TensorRT-8.5.3.1\include 中头文件拷贝到 安装的cuda目录中 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\include中
- 将 解压后的TensorRT-8.5.3.1\lib 中所有的lib文件拷贝到 安装的cuda目录中 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\lib\x64中
- 将 解压后的TensorRT-8.5.3.1\lib 中所有的dll文件拷贝到 安装的cuda目录中 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin中
- 将 :\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin 路径添加到用户的 Path 环境变量中
在上面操作添加以后,重启电脑,此时环境变量生效。
4. 验证TensorRT
激活conda 环境,依次输入下面语句。
python
import tensorrt
print(tensorrt.__version__)
并运行下载解压后安装包里的一个实例的 sample.py 文件, 来检查 tensorrt 是否能正常使用。
直接在激活的环境下, cd 进入,比如我的文件目录, 直接cd
C:\Users\xiaoalei\Desktop\TensorRT-8.5.3.1\samples\python\network_api_pytorch_mnist
- 主要看后面的一部分地址
-输入 python sample.py 运行
- 出现上面情况就是表示成功,(运行报错一般是安装包没有安装,直接 pip install 就行,我在安装提示只需安装 pycuda,直接 pip install pycuda)
二、YOLO v5模型导出函数
而在YOLOv5 6.0版本更新后,新增了一个export.py文件,支持大部分框架模型的导出,包括TensorRT。这里介绍部分模型导出函数
1.onnx 模型导出(def export_onnx()函数)
2.openvino模型导出(def export_openvino()函数)
3.coreml模型导出(def export_coreml()函数)
4.TensorRT模型导出(def export_engine()函数)
三、YOLO v5使用TensorRT进行加速检测
1. 模型导出
直接导出模型:(我这里对自己训练好的模型和权重文件进行导出)
- 这里的--half表示半精度模型,使用半精度可以加快推理速度,但会损失一定精度,直接导出可以不加
python export.py --weights runs/train/exp/weights/best.pt --data linhuo.yaml --include engine --device 0 --half
2. 图片检测
验证图片检测速度, 需要对 detect.py 进行修改,最后的程序入口进行修改
修改后如下:
if __name__ == '__main__':
start_time = time.time()
opt = parse_opt()
main(opt)
end_time = time.time()
print("程序花费时间{}秒".format(end_time - start_time))
3. 使用yolov5s.pt 和 yolov5s.engine(tensorrt)进行检测
命令行语句:
3.1 使用yolov5s.pt, 推理时间: 5.49s
python detect.py --weights runs/train/exp/weights/best.pt
3.2 使用 yolov5s.engine(tensorrt), 推理时间: 4.20s
python detect.py --weights runs/train/exp/weights/best.engine
四、总结
TensorRT能加速YOLO v5推理与检测。
后续会补充视频推理检测