前言
各位朋友,上一篇,我们介绍了tensorRT的一些基本概念,今天我们一起来探究一下如何基于LabVIEW使用TensorRT实现图像分类,本文所用到的工具包的安装与下载方法可参考博文:https://blog.csdn.net/virobotics/article/details/129304465
三、使用LabVIEW TensorRT实现图像分类完整步骤
1、模型获取及转换为onnx
- 安装pytorch和torchvision
- 获取torchvision中的模型:mobile(我们获取预训练好的模型):
original_model = models.mobilenet_v2(pretrained=True)
- 转onnx
def get_pytorch_onnx_model(original_model):
# define the directory for further converted model save
onnx_model_path = dirname
# define the name of further converted model
onnx_model_name = "mobilenet.onnx"
# create directory for further converted model
os.makedirs(onnx_model_path, exist_ok=True)
# get full path to the converted model
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
# generate model input
generated_input = Variable(
torch.randn(1, 3, 224, 224)
)
# model export into ONNX format
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"],
output_names=["output"],
opset_version=11
)
return full_model_path
完整获取及模型转换python代码如下:
import os
import torch
import torch.onnx
from torch.autograd import Variable
from torchvision import models
dirname, filename = os.path.split(os.path.abspath(__file__))
print(dirname)
def get_pytorch_onnx_model(original_model):
# define the directory for further converted model save
onnx_model_path = dirname
# define the name of further converted model
onnx_model_name = "mobilenet.onnx"
# create directory for further converted model
os.makedirs(onnx_model_path, exist_ok=True)
# get full path to the converted model
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
# generate model input
generated_input = Variable(
torch.randn(1, 3, 224, 224)
)
# model export into ONNX format
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"],
output_names=["output"],
opset_version=11
)
return full_model_path
def main():
# initialize PyTorch MobileNetV2
original_model = models.mobilenet_v2(pretrained=True)
# get the path to the converted into ONNX PyTorch model
full_model_path = get_pytorch_onnx_model(original_model)
print("PyTorch MobileNet model was successfully converted: ", full_model_path)
if __name__ == "__main__":
main()
注:该模型所用数据集为:ImageNet数据集
2、onnx转化为engine(save engine_Simple.vi)
使用onnx_to_engine.vi,将该vi拖拽至前面板空白区域,创建并输入onnx的路径以及engine的路径,type即精度,可选择FP32或FP16,肉眼观看精度无大差别。(一般FP16模型比FP32速度快一倍)
转换的完整程序如下:
点击运行,等待1~3分钟,模型转换成功,可在刚刚设定的路径中找到我们转化好的mobilenet.engine.
Q:为什么要转换模型,不直接调用ONNX?
A:tensorRT内部加载ONNX后其实是做了一个转换模型的工作,该过程时间长、占用内存巨大。因此不推荐每次初始化都加载ONNX模型,而是加载engine。
3、LabVIEW TensorRT调用 mobilenet实现图像分类(call_tensorRT.vi)
1、加载mobilenet.engine2、根据转换前的onnx模型设置输入输出缓存 3、将预处理后的输入tensor及其数据大小(即1 3 224 * 224)送至run函数中。(可使用输入一维数组的方式,亦可使用输入数据指针地址的方式,后者稍快) 4、根据模型输出的数据量,首先初始化一个相同数据量的数组(1维、2维、3维皆可,根据用户后处理方便可自行选择),使用Get_Result获取相同数据量和形状的结果。
默认第一个输出的index为0,如有多个输出,第n个输出的index为n-1 5、完整图像分类源码 6、分类结果及速度,如下图所示,使用tensorRT FP16,分类一帧大概需要0.5毫秒左右,速度极快 7、不同工具包实现图像分类速度比较,对比可见对于分类,纯TenosRT推理速度是onnx使用TensoRT推理速度的2倍
四、使用INetworkDefinition编辑模型
1、tensorRT高级API
1、INetworkDefinition(INetwork文件夹):可新建一个空白网络或加载一个Onnx网络,获取其所有的层(算子,ILayer)、输入和输出张量(ITensor),用户也可以执行以下操作:添加算子、输入输出、将某一个张量设置为输出、保存成engine文件等 2、ILayer:INetworkDefinition中的任意一层,实质为一个算子。tensorRT中共有41种Ilayer子类,如右图。使用属性节点可获取、编辑每一种子类(如IConvilutionLayer)的相关信息
使用INetwork下的多态VI addLayer可在网络中添加指定层
3、ITensor:INetworDefinition中的变量,一般为某一个ILayer的输入或输出。可通过以下四个函数获取某一个ITensor:
INetworkDefinition.getInput(index)
INetworkDefinition.getOutput(index)
ILayer.getInput(index)
ILayer.getOutput(index)
在INetworkDefinition中,使用markOutput将某个ITensor设置为网络的输出,使用unMarkOutput将某个原本是输出的ITensor设置为非输出。
2、给pytorch分类模型加一层softMax
softMax理论知识
假设有一个数组V,V i 表示V中的第i个元素,那么这个元素的softmax值为:
在pytorch分类模型的输出结果中,j=1000,i为0~999
为什么softmax是这种形式?首先,我们知道概率有两个性质:
1)预测的概率为非负数;
2)各种预测结果概率之和等于1。
softmax就是将在负无穷到正无穷上的预测结果按照这两步转换为概率的。
使用LabVIEW给分类模型加一层softMax(save engine_add_softMax.vi)
- 加载mobilenet.onnx模型;
- 将原来的输出tensor设置为非输出;
- 使用addSoftMax,的输出tensor后面加一层,返回该层句柄;
- 使用属性节点的Axes设置为2。因为原本的输出shape为(1,1000),因此要在1000这个维度的基础上添加softMax,因此Axes=2;
- 设置新加一层的输出为整个模型的输出,保存模型。
调用并推理加了softMax的分类模型engine(call_tensorRT_softMax.vi)
五、完整项目源码下载
大家可关注微信公众号: VIRobotics,回复关键字:TensorRT实现mobilenet图像分类 获取本次分享内容的完整项目源码及模型。
TensorRT实现mobilenet图像分类
附加说明
操作系统:Windows10
python:3.6及以上
LabVIEW:2018及以上 64位版本
视觉工具包:techforce_lib_opencv_cpu-1.0.0.98.vip
LabVIEW TensorRT工具包:virobotics_lib_tensorrt-1.0.0.22.vip
总结
以上就是今天要给大家分享的内容。如果有问题可以在评论区里讨论,提问前请先点赞支持一下博主哦,如您想要探讨更多关于LabVIEW与人工智能技术,欢迎加入我们的技术交流群:705637299。进群请备注:LabVIEW机器视觉
**如果文章对你有帮助,欢迎✌关注、👍点赞、✌收藏