onnx转换tensorrt的两种方法

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档文章目录前言一、tensorrt是什么？二、使用步骤1.引入库2.读入数据总结前言提示：这里可以添加本文要记录的大概内容：例如：随着人工智能的不断发展，机器学习这门技术也越来越重要，很多人都开启了学习机器学习，本文就介绍了机器学习的基础内容。提示：以下是本篇文章正文内容，下面案例可供参考一、tensorrt是什么？示例：pandas

weixin_42476942

25088人浏览 · 2021-02-08 15:39:23

weixin_42476942 · 2021-02-08 15:39:23 发布

文章目录

前言
一、tensorrt是什么？
二、方法
- 方法一.配置环境，onnx-tensorrt 6的安装
- 方法二.直接使用python脚本，转换onnx至trt

前言

本文主要讲解onnx转换至trt的过程

提示：以下是本篇文章正文内容，下面案例可供参考

一、tensorrt是什么？

tensorRT 是推理优化器，能对训练好的模型进行优化。可以理解为只有前向传播的深度学习框架，这个框架可以将 Caffe，TensorFlow的网络模型解析，然后与tensorRT中对应的层进行一一映射，把其他框架的模型统一全部转换到tensorRT中，然后在tensorRT中可以针对NVIDIA自家GPU实施优化策略，并进行部署加速。当你的网络训练完之后，可以将训练模型文件直接丢进tensorRT中，而不再需要依赖深度学习框架（Caffe，TensorFlow等)，

二、方法

先把onnx转化为TensorRT的trt文件，然后让c++环境下的TensorRT直接加载trt文件，从而构建engine，本文主要讲解onnx转换至trt的过程。

方法一.配置环境，onnx-tensorrt 6的安装

基础环境：pytorch 1.1 ubuntu16.04 tensorrt 6.0 cuda 10.0 x86系统

1.1下载onnx-tensorrt
https://github.com/onnx/onnx-tensorrt/tree/6.0
在这里插入图片描述
1.2下载后onnx-tensorrt/third_party/ onnx中内容是空的，需要下载与tensorrt6.0匹配的onnx1.4
https://github.com/onnx/onnx/releases

1.3 将下载后的内容，解压放置到onnx-tensorrt/third_party/onnx/

1.4 在onnx-tensorrt-release-6.0文件下的CMakeLists.txt加入
link_directories(/home/rex/data/projects/TensorRT-6.0.1.5/lib)
include_directories(/usr/local/cuda/include)
include_directories(/home/rex/data/projects/TensorRT-6.0.1.5/include)

1.5在onnx-tensorrt-release-6.0文件下新建build文件夹

1.6export PATH="$PATH:/home/rex/data/projects/TensorRT-6.0.1.5"

1.7cmake …
编译成功如下：
在这里插入图片描述 1.8make -j8

1.9 sudo make install
至此编译结束
1.10 在onnx-tensorrt-release-6.0文件下，运行在onnx2trt my_model.onnx -o my_engine.trt
即可生成转换好的trt
上述过程中，可能遇到的问题：

问题一：未在build中声明export PATH="$PATH:/home/rex/data/projects/TensorRT-6.0.1.5"

问题二：cmake error at CmakeLists.txt:121 (add_subdirectory)
The source directory
/home/xxx/Downloads/onnx-tensorrt-master/third_party/onnx
does not contain a Cmakelists.txt file
步骤1.2-1.3未进行或者未成功。
在这里插入图片描述
问题三：版本不匹配

方法二.直接使用python脚本，转换onnx至trt

适合用于tensorrt 6版本
若转换失败，出现节点问题，需要可视化onnx网络模型定位节点问题

import os
import tensorrt as trt
import sys

TRT_LOGGER = trt.Logger()
model_path = 'vgg16.onnx'
engine_file_path = "vgg16.trt"
# EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)  # batchsize=1
EXPLICIT_BATCH = []

# with trt.Builder(TRT_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) \
#        as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network,
                                                                                                            TRT_LOGGER) as parser:
    builder.max_workspace_size = 1 << 28
    builder.max_batch_size = 1
    print(network)
    if not os.path.exists(model_path):
        print('ONNX file {} not found.'.format(model_path))
        exit(0)
    print('Loading ONNX file from path {}...'.format(model_path))
    with open(model_path, 'rb') as model:
        print('Beginning ONNX file parsing')
        if not parser.parse(model.read()):
            print('ERROR: Failed to parse the ONNX file.')
            for error in range(parser.num_errors):
                print('parser.get_error(error)', parser.get_error(error))

    last_layer = network.get_layer(network.num_layers - 1)
    network.mark_output(last_layer.get_output(0))

    network.get_input(0).shape = [1, 3, 640, 160]
    print('Completed parsing of ONNX file')
    engine = builder.build_cuda_engine(network)
    with open(engine_file_path, "wb") as f:
        f.write(engine.serialize())
        print('save  trt success!!')