前言
上一篇博客给大家介绍了LabVIEW开放神经网络交互工具包【ONNX】,今天我们就一起来看一下如何使用LabVIEW开放神经网络交互工具包实现TensorRT加速YOLOv5。
以下是YOLOv5的相关笔记总结,希望对大家有所帮助。
内容 | 地址链接 |
---|---|
【YOLOv5】LabVIEW+OpenVINO让你的YOLOv5在CPU上飞起来 | |
【YOLOv5】LabVIEW OpenCV dnn快速实现实时物体识别(Object Detection) |
一、TensorRT简介
TensorRT是一个高性能的深度学习推理(Inference)优化器,可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架,将TensorRT和NVIDIA的GPU结合起来,能在几乎所有的框架中进行快速和高效的部署推理。主要用来针对 NVIDIA GPU进行 高性能推理(Inference)加速。
通常我们做项目,在部署过程中想要加速,无非就那么几种办法,如果我们的设备是CPU,那么可以用openvion,如果我们希望能够使用GPU,那么就可以尝试TensorRT了。那么为什么要选择TensorRT呢?因为我们目前主要使用的还是Nvidia的计算设备,TensorRT本身就是Nvidia自家的东西,那么在Nvidia端的话肯定要用Nvidia亲儿子了。
不过因为TensorRT的入门门槛略微有些高,直接劝退了想要入坑的玩家。其中一部分原因是官方文档比较杂乱;另一部分原因就是TensorRT比较底层,需要一点点C++和硬件方面的知识,学习难度会更高一点。我们做的开放神经网络交互工具包GPU版本,在GPU上做推理时,ONNXRuntime可采用CUDA作为后端进行加速,要更快速可以切换到TensorRT,虽然和纯TensorRT推理速度比还有些差距,但也十分快了。如此可以大大降低开发难度,能够更快更好的进行推理。。
二、准备工作
按照
三、YOLOv5模型的获取
为方便使用,博主已经将yolov5模型转化为onnx格式,可在百度网盘下载 链接:
1.下载源码
将Ultralytics开源的YOLOv5代码Clone或下载到本地,可以直接点击Download ZIP进行下载,
下载地址:
2.安装模块
解压刚刚下载的zip文件,然后安装yolov5需要的模块,记住cmd的工作路径要在yolov5文件夹下:
打开cmd切换路径到yolov5文件夹下,并输入如下指令,安装yolov5需要的模块
pip install -r requirements.txt
3.下载预训练模型
打开cmd,进入python环境,使用如下指令下载预训练模型:
1 import torch 2 ? 3 # Model 4 model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # or yolov5n - yolov5x6, custom
?
成功下载后如下图所示:
4.转换为onnx模型
在yolov5之前的yolov3和yolov4的官方代码都是基于darknet框架实现的,因此opencv的dnn模块做目标检测时,读取的是.cfg和.weight文件,非常方便。但是yolov5的官方代码是基于pytorch框架实现的。需要先把pytorch的训练模型.pt文件转换到.onnx文件,然后才能载入到opencv的dnn模块里。
将.pt文件转化为.onnx文件,主要是参考了nihate大佬的博客:https://blog.csdn.net/nihate/article/details/112731327
将export.py做如下修改,将def export_onnx()中的第二个try注释掉,即如下部分注释:
1 ''' 2 try: 3 check_requirements(('onnx',)) 4 import onnx 5 ? 6 LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') 7 f = file.with_suffix('.onnx') 8 print(f) 9 ? 10 torch.onnx.export( 11 model, 12 im, 13 f, 14 verbose=False, 15 opset_version=opset, 16 training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.eva l, 17 do_constant_folding=not train, 18 input_names=['images'], 19 output_names=['output'], 20 dynamic_axes={ 21 'images': { 22 0: 'batch', 23 2: 'height', 24 3: 'width'}, # shape(1,3,640,640) 25 'output': { 26 0: 'batch', 27 1: 'anchors'} # shape(1,25200,85) 28 } if dynamic else None) 29 ? 30 # Checks 31 model_onnx = onnx.load(f) # load onnx model 32 onnx.checker.check_model(model_onnx) # check onnx model 33 ? 34 # Metadata 35 d = {'stride': int(max(model.stride)), 'names': model.names} 36 for k, v in d.items(): 37 meta = model_onnx.metadata_props.add() 38 meta.key, meta.value = k, str(v) 39 onnx.save(model_onnx, f)'''
并新增一个函数def my_export_onnx():
1 def my_export_onnx(model, im, file, opset, train, dynamic, si