TensorRT中对ONNX模型解析过程

最近正在梳理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加入自定义层PluginIPluginV2Layer* 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);// 获取层的输入TensorITensor* getInput(int32_t index);// 设置层的输出Tensor的数据类型void setOutputType(int32_t index, DataType dataType);// 获取层的输入TensorITensor* 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的weightvoid setKernelWeights(Weights weights);// 设置卷积运算Kernel的biasvoid 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);// 设置卷积运算的Stridevoid setStrideNd(Dims stride);// 设置卷积运算的Dilationvoid 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<TensorOrWeights>* 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<TensorOrWeights> nodeInputs;        std::ostringstream ssInputs{};        ssInputs << nodeName << " [" << node.op_type() << "] 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: " << inputName);                ASSERT( (ctx->tensors().count(inputName)) && "Node input was not registered.", ErrorCode::kINVALID_GRAPH);                nodeInputs.push_back(ctx->tensors().at(inputName));                ssInputs << "[" << inputName << " -> " << nodeInputs.back().shape() << "[" << nodeInputs.back().getType() << "]" <<"], ";            }        }        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: " << node.op_type() << ". Attempting to import as plugin.");            importFunc = &opImporters.at("FallbackPluginImporter");        }        std::vector<TensorOrWeights> 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

• 在Conv算子添加函数内，大致流程如下：

• 获取输入Tensor，对于算子的input[0]

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 " << output_tensor_ptr->getName() << " as output: " << output.name());        output_tensor_ptr->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);        }    }

总结

本文主要介绍了ONNX和TensorRT的IR信息，并且梳理了从ONNX转换成TensorRT计算图的主要流程。如果文中有纰漏之处，欢迎批评指正，谢谢！