最近正在梳理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，添加输出信息，其部分代码如下：
// 向Network添加输入信息
ITensor* addInput(const char* name, DataType type, Dims dimensions);
// 向Network加入Conv层
IConvolutionLayer* addConvolutionNd(ITensor& input, int32_t nbOutputMaps,
Dims kernelSize, Weights kernelWeights, Weights biasWeights);
// 向Network加入Pooling层
IPoolingLayer* addPoolingNd(ITensor& input, PoolingType type, Dims windowSize);
// 向Network加入自定义层Plugin
IPluginV2Layer* addPluginV2(ITensor* const* inputs, int32_t nbInputs,
IPluginV2& plugin)；
// 向Network加入输出信息
void markOutput(ITensor& tensor);
这里的Layer，对应ONNX中不同类型的OP，不同类型Layer所包含的信息也不相同。我们先看下所有层的基类ILayer。
https://github.com/NVIDIA/TensorRT/blob/release/8.0/include/NvInfer.h#L478%EF%BC%89%EF%BC%8C%E4%B8%BB%E8%A6%81%E6%98%AF%E5%8A%9F%E8%83%BD%E6%98%AF%E8%AE%BE%E7%BD%AE%E8%AF%A5%E5%B1%82%E7%9A%84%E8%BE%93%E5%85%A5%E8%BE%93%E5%87%BA%E4%BF%A1%E6%81%AF%E3%80%81%E6%95%B0%E6%8D%AE%E7%B2%BE%E5%BA%A6%E4%BF%A1%E6%81%AF%E3%80%82%E4%B8%BB%E8%A6%81%E6%8E%A5%E5%8F%A3%E5%A6%82%E4%B8%8B
该类主要是功能是设置该层的输入输出信息、数据精度信息，主要代码如下：
// 设置层的输入信息
void setInput(int32_t index, ITensor& tensor);
// 获取层的输入Tensor
ITensor* getInput(int32_t index);
// 设置层的输出Tensor的数据类型
void setOutputType(int32_t index, DataType dataType);
// 获取层的输入Tensor
ITensor* getOutput(int32_t index);
// 设置层的精度类型
void setPrecision(DataType dataType);
看到这里的接口，有些同学可能会有疑惑，为什么只有setInput接口，没有setOutput接口？因为每一个层都会在内部产生输出Tensor，比如Conv层会把结果保存到一个Tensor中，不需要我们在外部设置，但是我们可以通过获取到getOutput接口输出Tensor信息。
我们看下IConvolutionLayer的定义。
https://github.com/NVIDIA/TensorRT/blob/release/8.0/include/NvInfer.h#L954%EF%BC%89%E7%9A%84%E5%AE%9A%E4%B9%89%EF%BC%8C%E4%B8%BB%E8%A6%81%E6%8E%A5%E5%8F%A3%E6%9C%89%E8%AE%BE%E7%BD%AELayer%E7%9A%84%E8%BE%93%E5%85%A5Tensor%E3%80%81%E8%BE%93%E5%87%BATensor%EF%BC%8C%E5%B7%B2%E7%BB%8F%E5%8D%B7%E7%A7%AF%E8%BF%90%E8%A1%8C%E6%98%AF%E5%B1%9E%E6%80%A7%E5%92%8Cweight
该类主要接口有设置Layer的输入Tensor、输出Tensor，以及卷积运算的属性和weight等信息，部分代码如下。
// 设置卷积运算Kernel的weight
void setKernelWeights(Weights weights);
// 设置卷积运算Kernel的bias
void setBiasWeights(Weights weights);
// 设置卷积运算的分组数量
void setNbGroups(int32_t nbGroups);
// 设置卷积运算的Padding信息
void setPrePadding(Dims padding);
void
setPostPadding(Dims padding);
void setPaddingMode(PaddingMode paddingMode);
void setPaddingNd(Dims padding)
// 设置卷积运算Kernel的尺寸
void setKernelSizeNd(Dims kernelSize);
// 设置卷积运算的Stride
void setStrideNd(Dims stride);
// 设置卷积运算的Dilation
void setDilationNd(Dims dilation);
TensorRT解析ONNX流程
在了解了ONNX和TensorRT的基本信息后，我们看下TensorRT的ONNX Parser解析ONNX模型过程。代码入口：
https://github.com/onnx/onnx-tensorrt/blob/4225037958191705a20684019a7df694c81ec39b/ModelImporter.cpp%23L505
第一，解析ONNX模型的输入信息，然后调用TensorRT接口添加输入信息，代码实现也比较简单，如下：
Status importInputs(ImporterContext* ctx, ::ONNX_NAMESPACE::GraphProto const& graph,
string_map* tensors)
{
// The weights come from the Initializer list in onnx graph
// Initializers are not really network inputs, so they need to be excluded.
std::unordered_set<std::string> initializers{};
for (const ::ONNX_NAMESPACE::TensorProto& initializer : graph.initializer())
{
initializers.emplace(initializer.name());
}
for (const ::ONNX_NAMESPACE::ValueInfoProto& input : graph.input())
{
TensorOrWeights tensor;
if (!initializers.count(input.name()))
{
nvinfer1::ITensor* tensor_ptr;
CHECK(importInput(ctx, input, &tensor_ptr));
tensor = tensor_ptr;
}
ctx->registerTensor(std::move(tensor), input.name());
}
return Status::success();
}
第二，对onnx模型的算子进行拓扑排序，按照拓扑序列，解析ONNX的算子，然后调用TensorRT对应的接口，在TensorRT内添加对应的layer，代码链接如下。
https://github.com/onnx/onnx-tensorrt/blob/4225037958191705a20684019a7df694c81ec39b/ModelImporter.cpp%23L96
这里说是的对应的layer，该对应的意思是，TensorRT的Layer不一定和ONNX的OP同名，但是在描述模型的计算图内有相同的表达意义。比如ONNX的BatchNorm对会转换成TensorRT中的ScaleLayer。这种映射关系，有时候是一对一，有时候是一对多，有时候是N对M，要根据不同IR中算子的含义、颗粒度等信息来具体情况具体分析。
这里以Conv算子为例，梳理一下添加单个算子的流程。
-
在计算图里查找Conv算子的输入。ONNX和TensorRT的Network，会对Tensor设置一个名字，因此可以根据名字对查找到对应的Tensor。
// Assemble node inputs. These may come from outside the subgraph.
std::vector nodeInputs;
std::ostringstream ssInputs{};
ssInputs <" [" <"] inputs: ";
for (const auto& inputName : node.input())
{
// Empty input names indicate optional inputs which have not been supplied.
if (inputName.empty())
{
nodeInputs.emplace_back(nullptr);
ssInputs <"[optional input, not set], ";
}
else
{
LOG_VERBOSE("Searching for input: " < ASSERT( (ctx->tensors().count(inputName)) && "Node input was not registered.", ErrorCode::kINVALID_GRAPH);
nodeInputs.push_back(ctx->tensors().at(inputName));
ssInputs <"[" <" -> " <"[" <"]" <<"], ";
}
}
LOG_VERBOSE(ssInputs.str());
-
根具算子的名称，找到对应算子的添加函数，然后执行该函数。
// Dispatch to appropriate converter.
const NodeImporter* importFunc{nullptr};
if (opImporters.count(node.op_type()))
{
importFunc = &opImporters.at(node.op_type());
}
else
{
LOG_INFO("No importer registered for op: " <". Attempting to import as plugin.");
importFunc = &opImporters.at("FallbackPluginImporter");
}
std::vector outputs;
try
{
GET_VALUE((*importFunc)(ctx, node, nodeInputs), &outputs);
}
catch (const std::exception& e)
{
return MAKE_ERROR(makeErrorExplanation(e, nodeName), ErrorCode::kINVALID_NODE);
}
if (ctx->hasError())
{
return MAKE_ERROR(makeErrorExplanation(ctx, nodeName), ErrorCode::kINVALID_NODE);
}
-
对于Conv算子，会执行对应的Conv添加函数，对应的代码链接：
https://github.com/onnx/onnx-tensorrt/blob/4225037958191705a20684019a7df694c81ec39b/builtin_op_importers.cpp%23L604
nvinfer1::ITensor* tensorPtr = &convertToTensor(inputs.at(0), ctx);
-
获取Conv算子的weight，对于算子的input[1]，代码链接
auto kernelWeights = inputs.at(1).weights();
-
如果Conv算子有bias，获取Conv算子的bias，对应算子的input[2]
-
获取Conv的属性，包括stride、padding、dilation、grroup等属性
-
在TensorRT Network中添加Conv Layer
nvinfer1::IConvolutionLayer* layer
= ctx->network()->addConvolutionNd(*tensorPtr, noutput, kernelSize,
kernelWeights, bias_weights);
-
设置Conv Layer相关属性，包括stride、padding、dilation、group等属性
-
获取Conv Layer的输出Tensor，然后返回该Tensor的地址
第三，解析onnx模型的输出信息，在trt没添加对应的输出信息。
// Mark outputs defined in the ONNX model (unless tensors are user-requested)
for (::ONNX_NAMESPACE::ValueInfoProto const& output : graph.output())
{
ASSERT((_importer_ctx.tensors().count(output.name())) && "The output tensor was not registered.",
ErrorCode::kINVALID_GRAPH);
nvinfer1::ITensor* output_tensor_ptr
= &convertToTensor(_importer_ctx.tensors().at(output.name()), &_importer_ctx);
LOG_VERBOSE("Marking " setName(output.name().c_str());
if (output_tensor_ptr->isNetworkInput())
{
// HACK WAR for TRT not allowing input == output
// TODO: Does this break things by changing the name of the input tensor?
output_tensor_ptr->setName(("__" + output.name()).c_str());
output_tensor_ptr = &identity(&_importer_ctx, output_tensor_ptr).tensor();
ASSERT(output_tensor_ptr && "Failed to add an Identity layer.", ErrorCode::kUNSUPPORTED_NODE);
output_tensor_ptr->setName(output.name().c_str());
}
nvinfer1::ITensor** user_output = _importer_ctx.getUserOutput(output.name().c_str());
if (!user_output)
{
_importer_ctx.network()->markOutput(*output_tensor_ptr);
nvinfer1::DataType output_trt_dtype;
ASSERT(convertDtype(output.type().tensor_type().elem_type(), &output_trt_dtype) && "Failed to convert ONNX date type to TensorRT data type.", ErrorCode::kUNSUPPORTED_NODE);
// For INT32 data type, output type must match tensor type
ASSERT( (output_tensor_ptr->getType() != nvinfer1::DataType::kINT32
|| output_trt_dtype == nvinfer1::DataType::kINT32) && "For INT32 tensors, the output type must also be INT32.",
ErrorCode::kUNSUPPORTED_NODE);
// Note: Without this, output type is always float32
output_tensor_ptr->setType(output_trt_dtype);
}
}
