Onnx-tensorrt详解之nvonnxparser库

Onnx-tensorrt工程提供了所有onnx builtin layer向trtmodel转换的parser代码。

Onnx-tensorrt工程链接： https://github.com/onnx/onnx-tensorrt

一、nvonnxparser库概览

nvonnxparser库的核心代码文件如下：

set(IMPORTER_SOURCES
  NvOnnxParser.cpp
  ModelImporter.cpp
  builtin_op_importers.cpp
  onnx2trt_utils.cpp
  ShapedWeights.cpp
  OnnxAttrs.cpp
)  //见CMakeLists.txt文件

最终，这些代码被编译成动态链接库nvonnxparser.so和静态链接库nvonnxparser_static.a 。

add_library(nvonnxparser SHARED ${IMPORTER_SOURCES})
target_link_libraries(nvonnxparser PUBLIC onnx_proto nvonnxparser_plugin ${PROTOBUF_LIBRARY}
add_library(nvonnxparser_static STATIC ${IMPORTER_SOURCES})
target_link_libraries(nvonnxparser_static PUBLIC onnx_proto nvonnxparser_plugin ${PROTOBUF_LIBRARY} ${CUDNN_LIBRARY} ${TENSORRT_LIBRARY})

==========================================================

以上代码文件的逻辑关系由外而内：
1. NvOnnxParser

是最外层接口，tensorrt的代码auto parser = nvonnxparser::createParser(*network, gLogger)中createParser的函数就定义在此文件中；

2. ModelImporter  

nvonnxparser::createParser函数内部代码return new onnx2trt::ModelImporter(network, logger)，调用了位于ModelImporter中的函数。

ModelImporter中定义了：
parser->parse(args);   //A method called parse() from the object of type iParser is called to read the model file and populate the TensorRT network; The specific args depend on what format parser is used.  For args are information about the pretrained onnxFile.
parser->parseFromFile(onnx_filename, ILogger::Severity::kWARNING);

3. builtin_op_importers

ModelImporter调用builtin_op_importers中定义的函数getBuiltinOpImporterMap()；builtin_op_importers是parse的核心代码。
 

4. onnx2trt_utils.cpp，ShapedWeights.cpp，OnnxAttrs.cpp是builtin_op_importers中用到的一些函数。

 

二、parse核心代码builtin_op_importers

builtin_op_importers，即builtin_opterators_importers。

• onnxmodel到trtmodel的parse代码。从onnxmodel的input出发，最后，输出trtmodel的输出tensor_ptr；
• onnx支持的builtin operators包括Conv, Argmax, Unsample,Relu等，具体可以参考operators.md文件；
• 文件中根据onnx层的类型名调用相应的DEFINE_BUILTIN_OP_IMPORTER(Conv), DEFINE_BUILTIN_OP_IMPORTER(Argmax),  DEFINE_BUILTIN_OP_IMPORTER(Unsample), DEFINE_BUILTIN_OP_IMPORTER(Relu)等，从而完成对应层的onnx2trtmodel的parser。

主要完成以下三项工作：
1、将onnx输入数据转化为trt要求的数据格式
2、建立trt层
3、计算trt输出结果

以DEFINE_BUILTIN_OP_IMPORTER(Conv) 为例：
//************将onnx输入数据转化为trt要求的数据格式*******************
nvinfer1::ITensor* tensor_ptr = &inputs.at(0).tensor();
auto kernel_weights = inputs.at(1).weights();   //onnxmodel的输入格式 inputs=['x','W'],转化为trtmodel输出的数据格式
int noutput = kernel_weights.shape.d[0];

//*************************建立trt层*********************
nvinfer1::IConvolutionLayer* layer = ctx->network()->addConvolution(*tensor_ptr, noutput, kernel_size, kernel_weights, bias_weights);  //此时，onnx的layer已经转化为trtmodel的layer,   ctx（context的简写）就是trtmodel的network。  trt官方文档给出的添加convolution层的例子：IConvolutionLayer* conv1 = network->addConvolution(*scale_1->getOutput(0), 20, DimsHW{5, 5}, mWeightMap["conv1filter"], mWeightMap["conv1bias"]);

//***********************计算trt输出结果********************
tensor_ptr = layer->getOutput(0); //利用trtmodel计算输出输出tensor，并作为输出返回
return {{tensor_ptr}}; //返回输出tensor  Y

PS:  builtin_op_importers代码说明：

DEFINE_BUILTIN_OP_IMPORTER() 函数是builtin_op_importer的核心内容；

另外，builtin_op_importer文件中前面定义了DEFINE_BUILTIN_OP_IMPORTER()共用的一些ultils函数；

argMinMaxHelper函数是DEFINE_BUILTIN_OP_IMPORTER(ArgMax)和DEFINE_BUILTIN_OP_IMPORTER(ArgMin)共用的。

外部接口函数如下：

string_map& getBuiltinOpImporterMap() {
  static string_map builtin_op_importers;
  return builtin_op_importers;
}

namespace {

bool registerBuiltinOpImporter(std::string op,
                               NodeImporter const& importer) {
  bool inserted = getBuiltinOpImporterMap().insert({op, importer}).second;
  assert(inserted);
  return inserted;
}

三、自定义层

如果需要自定义层(非onnx中builtin的层)，则需要在此文件中写自定义的parse方法，同样完成以上说明的三项工作。

1. 自定义input数据的转化;

2. 通过addPlugin接口将自定义层写入网络中。以SplitPlugin为例，使用ctx->addPlugin(new SplitPlugin(axis, output_lengths), {&inputs.at(0).tensor()}); 

3. 定义output的计算方式。