前言

各位朋友，今天我们一起来探究一下如何基于LabVIEW使用TensorRT实现图像分类需要做的一些准备工作，本文所用到的工具包的安装与下载方法可参考博文：https://blog.csdn.net/virobotics/article/details/129304465

一、TensorRT简介

TensorRT是一个高性能的深度学习推理（Inference）优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。主要用来针对 NVIDIA GPU进行高性能推理（Inference）加速。

TensorRT与其他推理引擎速度对比：

注：测试电脑cpu为i7-11800H，Intel集显为i7-11650G7，独显为笔记本RTX 3070（包括预处理和后处理）

二、TensorRT工具包常用函数介绍

TensorRT工具包所在位置：

TrtSession：推理函数系列，可直接推理engine(trt)模型和onnx模型

TensorRT工具包常用函数介绍：

三、模型获取及转换为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()

附加说明

操作系统：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

总结

以上就是今天要给大家分享的内容。下一篇我们将给大家讲解图像分类使用tensorRT工具包在LabVIEW中的部署。如果有问题可以在评论区里讨论，如您想要探讨更多关于LabVIEW与人工智能技术，欢迎加入我们的技术交流群：705637299。进群请备注：LabVIEW机器视觉

**如果文章对你有帮助，欢迎✌关注、👍点赞、✌收藏