本文介绍 onnx-tensorrt实现添加自己的模型plugin

onnx-tensorrt实现添加自己的模型plugin

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

总结一下现在基本上深度学习模型部署的套路:

  • 一般都是把模型training然后export 到onnx;
  • 但是onnx如何推理呢,要么使用一些backend,或者使用onnx-tensorrt;

实际上使用tensorrt加速基本上是大势所驱啊。tensorrt目前支持的层不少,要是能再多一些,那么基本上大部分的层都可以采用tensorrt来推理了。但是有时候呢,有一些层是不支持的,此时要想尽一切办法来进行补充。举个例子,比如maskrcn模型,截止到目前,它的模型支持情况可以如下所示:

All ops used in MaskRCNN-FPN-Resnet50:
MaxPool
ConstantOfShape
ConstantOfShape not supported in onnx_tensorrt or TensorRT 5.1 now.
TopK
Exp
Gather
NonZero
NonZero not supported in onnx_tensorrt or TensorRT 5.1 now.
Reshape
ReduceMin
Greater
Greater not supported in onnx_tensorrt or TensorRT 5.1 now.
Transpose
Sqrt
Shape
Add
Split
Squeeze
Div
Gemm
NonMaxSuppression
Constant
Conv
Expand
Expand not supported in onnx_tensorrt or TensorRT 5.1 now.
Not
Not not supported in onnx_tensorrt or TensorRT 5.1 now.
Slice
Equal
Equal not supported in onnx_tensorrt or TensorRT 5.1 now.
Cast
Log
And
And not supported in onnx_tensorrt or TensorRT 5.1 now.
Sub
Resize
RoiAlign
Scatter
Scatter not supported in onnx_tensorrt or TensorRT 5.1 now.
Mul
Clip
Unsqueeze
Less
Less not supported in onnx_tensorrt or TensorRT 5.1 now.
Concat
Softmax
Floor
Relu
Sigmoid
ConvTranspose
Flatten

可以看到,只有一些看起来十分简单的层不支持,其他的,其实都支持。但是这些层看似简单,实现起来呢?我想应该不简单。比如说这么一些比较层,逻辑层等等。

我们将尝试在onnxtrt中加入Bi-linear Upsamling 层,也就是双线性差值上采样层,来探索一下在onnxtrt中添加一个层的难度。

## 寻找一个测试模型

要添加BilinearUpsampling层就得先知道这个层可以用哪个模型来测试,这个模型要满足两个条件:

  • 直接用onnx2trt无法转换的,因为它包含BilinearUpsampling层;
  • 其他的层都应该被支持。

这样通过我们的添加之后,我们的onnx-tensorrt就可以转换它,从而进一步的尝试使用TensorRT的推理引擎来load这个engine进行推理。