本文介绍 onnx模型tensorrt推理实践

onnx模型tensorrt推理实践

This article was original written by Jin Tian, welcome re-post, first come with https://jinfagang.github.io . but please keep this copyright info, thanks, any question could be asked via wechat: jintianiloveu

现在很多场景模型部署都是用tensorrt,因为显卡速度比cpu快很多,加上一些场景还真的不在乎成本,因此为了使得系统更加稳定可靠,用tensorrt就显得很不错了,典型的场景比如自动驾驶。

那么我觉得最重要的是,如何了解tensorrt,以及如果将onnx模型用tensorrt进行推理。

tensorrt 安装

tensorrt的安装就比较简单了,现在常用的是tensorrt5.1. 你可以下载cuda对应的版本,下载之后,主要进行这么几个步骤:

  1. 添加lib 到path里面:

    export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/home/xx/tensorrt5.1/lib:$LD_LIBRARY_PATH
    

    假设你的tensorrt解压在home目录。这样也是为了后续方便。

  2. 编译samples

    tensorrt自带了几个sample,可以cd进去make,然后在上一层的bin里面可以看到可执行文件。

  3. 安装onnx-tensorrt 打通onnx和tensorrt的桥梁

    直接clone然后编译:

    git clone https://github.com/onnx/onnx-tensorrt.git --recursive
    cmake .. -DTENSORRT_ROOT=~/tensorrt5.1
    

    然后就可以得到一个可执行的二进制文件:onnx2trt. 这样的话就可以将onnx转到trt上。

那么问题来了,我为什么要安装在home目录呢?等一下include的时候cmake又无法找到文件。。。那就直接在cmake里面指定吧,总感觉直接丢到系统目录不是一件好事。

onnx模型转trt

那么问题来了,如何将onnx转到tensorrt呢?onnx有一个onnx_tensorrt的转换工具,编译之后即可转换。比如我们要将mobilenetv2的onnx模型转到trt,那么:

onnx2trt mobilenetv2-1.0.onnx -o mobilenetv2-1.0.trt

需要注意的是,上面的trt实际上就是engine了已经。但是问题是这个trt,怎么欧通过tensorrt来做inference。

使用tensorrt推理trt模型

既然转到trt模型,那么用tensorrt如何推理呢。实际上tensorrt可以直接推理onnx模型,但是本质上还是要在代码中将模型转换为trt再进行推理的。

假设我们刚才通过onnx2trt工具,已经得到了一个engine模型,也就是说mobilenetv2_engine.trt. 我们用python代码来尝试inference一下:

import tensorrt as trt

G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogServerity.INFO)
engine = trt.utils.load_engine(G_LOGGER, 'mobilenetv2_engine.trt')

# create an execution context
context = engine.create_execution_context()

# allocate inputs and outputs
import pycuda.driver as cuda
import pycuda.autoinit

d_input = cuda.mem_alloc(input_shapes)
d_output = cuda.mem_alloc(output_shapes)
bindings = [int(d_input), int(d_output)]

# feed input data and execute inference
stream = cuda.Stream()
cuda.memcpy_htod_async(d_input, data, stream)
context.enqueue(batch_size, bindings, stream.handle, None)
cuda.memcpy_dtoh_async(output, d_output, stream)
stream.synchronize()

# or do something like this
cuda.memcpy_htod(d_input, data)
context.execute(batch_size, bindings)
cuda.memcpy_dtoh(output, d_output)

此时,我们应该已经将trt的engine 加载进来了。

python的例子是有了,但是C++的没有。

上面的api是老的tensorrt的api,新版的已经有很大的不一样了。具体看下一篇报道。

参考文献或网址

  1. https://faldict.github.io/faldict/tensorrt-summary/