极市平台 | 实践教程｜TensorRT中对ONNX模型解析过程

本文来源公众号“极市平台”，仅用于学术分享，侵权删，干货满满。

原文链接：实践教程｜TensorRT中对ONNX模型解析过程

作者丨文亚伟@知乎（已授权）

来源丨https://zhuanlan.zhihu.com/p/603338185

编辑丨极市平台

极市导读

本文主要介绍了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，添加输出信息，其部分代码如下：

// 向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。

TensorRT/include/NvInfer.h at release/8.0 · NVIDIA/TensorRT (github.com)

该类主要是功能是设置该层的输入输出信息、数据精度信息，主要代码如下：

// 设置层的输入信息
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

TensorRT/include/NvInfer.h at release/8.0 · NVIDIA/TensorRT (github.com)

该类主要接口有设置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<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计算图的主要流程。如果文中有纰漏之处，欢迎批评指正，谢谢！

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。