点击下方卡片，关注「集智书童」公众号
作者丨文亚伟@知乎 来源丨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。
该类主要是功能是设置该层的输入输出信息、数据精度信息，主要代码如下：
// 设置层的输入信息
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的定义。
该类主要接口有设置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模型过程。代码入口：
第一，解析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，代码链接如下。
这里说是的对应的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添加函数，对应的代码链接：
- 在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计算图的主要流程。如果文中有纰漏之处，欢迎批评指正，谢谢！
[
[
CVPR 2023｜SQR：对于训练DETR-family目标检测的探索和思考](https://mp.weixin.qq.com/s?__biz=MzU5OTA2Mjk5Mw==&mid=2247504428&idx=3&sn=8cacd143a074d113c8eac8776acccebc&chksm=feb83292c9cfbb844c7388474f29a1734a47029eff9dcb087c5c0db4649b3a818c2ddff3287d&scene=21#wechat_redirect)
[
扫码加入👉「集智书童」交流群
（备注： 方向+学校/公司+昵称 ）
想要了解更多：
前沿AI视觉感知全栈知识👉「分类、检测、分割、关键点、车道线检测、3D视觉（分割、检测）、多模态、目标跟踪、NerF」
行业技术方案 👉「AI安防、AI医疗、AI自动驾驶」
AI模型部署落地实战 👉「CUDA、TensorRT、NCNN、OpenVINO、MNN、ONNXRuntime以及地平线框架」
欢迎扫描上方二维码，加入「 集智书童-知识星球 」，日常分享论文、学习笔记、问题解决方案、部署方案以及全栈式答疑，期待交流！
免责声明
凡本公众号注明“来源：XXX（非集智书童）”的作品，均转载自其它媒体，版权归原作者所有，如有侵权请联系我们删除，谢谢。
点击下方“ 阅读原文 ”，
了解更多AI学习路上的 「武功秘籍」