这里主要介绍了ONNX和TensorRT的IR信息,并且梳理了从ONNX转换成TensorRT计算图的主要流程。
最近正在梳理TensorRT的ONNX Parser源码,该Parser的核心功能是将模型ONNX IR转换成TensorRT IR。
ONNX基础
首先,我们来看一下ONNX模型格式的基础知识,大家可以参考以下文章,在此不太赘述。
一图看懂ONNX模型格式3:https://zhuanlan.zhihu.com/p/425232454
ONNX学习笔记:https://zhuanlan.zhihu.com/p/346511883
TensorRT IR基础
其次,我们看一下TensorRT中构建IR的接口。在TensorRT中,没有使用Protobuffer定义IR,但是提供了相关接口,帮助用户自己定义IR。描述IR信息的类叫做INetworkDefinition,代码链接如下。
https://github.com/NVIDIA/TensorRT/blob/release/8.0/include/NvInfer.h%23L5417
该类提供的功能有,添加输入信息,添加layer,添加输出信息,其部分代码如下:
这里的Layer,对应ONNX中不同类型的OP,不同类型Layer所包含的信息也不相同。我们先看下所有层的基类ILayer。
该类主要是功能是设置该层的输入输出信息、数据精度信息,主要代码如下:
看到这里的接口,有些同学可能会有疑惑,为什么只有setInput接口,没有setOutput接口?因为每一个层都会在内部产生输出Tensor,比如Conv层会把结果保存到一个Tensor中,不需要我们在外部设置,但是我们可以通过获取到getOutput接口输出Tensor信息。
我们看下IConvolutionLayer的定义。
该类主要接口有设置Layer的输入Tensor、输出Tensor,以及卷积运算的属性和weight等信息,部分代码如下。
TensorRT解析ONNX流程
在了解了ONNX和TensorRT的基本信息后,我们看下TensorRT的ONNX Parser解析ONNX模型过程。代码入口:
第一,解析ONNX模型的输入信息,然后调用TensorRT接口添加输入信息,代码实现也比较简单,如下:
第二,对onnx模型的算子进行拓扑排序,按照拓扑序列,解析ONNX的算子,然后调用TensorRT对应的接口,在TensorRT内添加对应的layer,代码链接如下。
这里说是的对应的layer,该对应的意思是,TensorRT的Layer不一定和ONNX的OP同名,但是在描述模型的计算图内有相同的表达意义。比如ONNX的BatchNorm对会转换成TensorRT中的ScaleLayer。这种映射关系,有时候是一对一,有时候是一对多,有时候是N对M,要根据不同IR中算子的含义、颗粒度等信息来具体情况具体分析。
这里以Conv算子为例,梳理一下添加单个算子的流程。
- 在计算图里查找Conv算子的输入。ONNX和TensorRT的Network,会对Tensor设置一个名字,因此可以根据名字对查找到对应的Tensor。
- 根具算子的名称,找到对应算子的添加函数,然后执行该函数。
- 对于Conv算子,会执行对应的Conv添加函数,对应的代码链接:
- 在Conv算子添加函数内,大致流程如下:
- 获取输入Tensor,对于算子的input[0]
- 获取Conv算子的weight,对于算子的input[1],代码链接
- 如果Conv算子有bias,获取Conv算子的bias,对应算子的input[2]
- 获取Conv的属性,包括stride、padding、dilation、grroup等属性
- 在TensorRT Network中添加Conv Layer
- 设置Conv Layer相关属性,包括stride、padding、dilation、group等属性
- 获取Conv Layer的输出Tensor,然后返回该Tensor的地址
第三,解析onnx模型的输出信息,在trt没添加对应的输出信息。
总结
本文主要介绍了ONNX和TensorRT的IR信息,并且梳理了从ONNX转换成TensorRT计算图的主要流程。