TensorFlow
模型:要求输入二进制.pb
模型- 如果只有 checkpoint,可以使用
python/ckpt_to_pb.py
将模型转化为.pb
模型。 - 如果当前模型需要满足具有多输出时,输出层不能被其他任何层作为输入层调用。例如,模型网络为
Input->MatMul->BiasAdd
,可以生成以MatMul
作为输出的.pb
文件并利用 Forward 推理,不能同时将MatMul
和BiasAdd
作为输出层。 - 如果需要更改模型的输出层,可以使用
pythpb_change_output.py
进行快速更改,即生成新模型,后续再使用 Forward 进行构建。
- 如果只有 checkpoint,可以使用
PyTorch
模型:要求输入Torch Jit
模型,对模型有一定限制,详见 PyTorch 使用说明 。Keras
模型:要求输入.h5
模型。ONNX
模型:要求输入.onnx
模型。
通过 TensorFlow
/ PyTorch
/ Keras
/ ONNX
模型构建推理引擎的主要流程基本一致,细节不尽相同,每一种模型拥有各自的 Builder
(TfBuilder
/ TorchBuilder
/ KerasBuilder
/ OnnxBuilder
),调用各模型定制的 Parser
对原模型进行分析打包,生成 TrtNetworkDesc
中间结构和 std::vector<TrtLayerDesc>
数据,并交由 TrtForward
类的 Build
方法进行 INetworkDefinition
的 TensorRT
网络构建和 Engine
生成 (OnnxBuilder
将会省略部分步骤)。具体构流程可参考下方的流程图。
- 原始网络模型 -> 中间结构描述创建器 (流程图可参考此处)
- 当用户导入原始网络模型时,我们会创建一个
Parser
类的对象parser
,用于将当前模型转换成网络层描述。该对象中的描述创建器管理器会对每种网络层创建器进行注册,每个创建器继承自各自命名空间下的ILayerDescCreator
类,并重写基类虚函数Check
和Create
。Check
函数用于检查当前结点是否符合创建器需求;Create
函数根据当前结点信息调用对应的创建器,返回该结点的中间形式TrtLayerDesc
对象,同时将该结点的(原模型)输入结点保存到node_inputs
中。需要注意的是,部分创建器是基于结点(node)的,如clamp_creator
,activation_creator
等,他们是最低层的创建器;还有一部分创建器是基于 module 的,如bert_creator
、lrn_creator
等,他们的Check
函数会检查当前结点及之前的多个结点。一般来说,基于 module 的创建器在后续构建网络时优化程度更高,应当优先被检查,即在创建Parser
类的对象时被优先注册。 parser
会先调用Graph::Load
函数将原始网络模型转换成Graph
类的对象graph_
,随后调用Graph::ExtractGraphInfos
函数提取输入结点和输出结点信息,并分别保存到inputs_
和outputs_
对象中。- 对于所有的输入结点,我们调用
Parser::CreateInputDescs
函数,根据batch_size
创建InputDescs
,同时我们还会维护一张哈希表created_desc_map_
,保存输入结点和网络层描述的映射关系;对于其他所有结点网络层描述的创建,由于我们在上一步保存了所有输出结点信息,这里采用了自底向上的搜索方式,即从输出结点开始,直到搜索到输入结点为止。这一递归过程通过Parser::ParseOperaion
函数完成,在递归的过程中,我们不断地保存当前结点的输入结点信息,并创建对应的网络层描述,最终我们可以得到一个能够描述整个模型的对象network_
。具体递归流程可参考下方的伪代码。
- 当用户导入原始网络模型时,我们会创建一个
- 中间结构描述创建器 ->
TensorRT
网络层创建器 (流程图可参考此处)- 当
parser
完成对原始网络模型的解析后,我们还需要创建一个TrtNetworkCreator
类的对象creator
,用于将原始网络模型的网络层描述network_
转换成TensorRT
推理引擎。这一过程与上一小节中的创建器类似,都是先注册对应网络层的创建器,再进行后续的查找。由于创建器已经指定了对应格式的TrtLayerDesc
,因此创建器不再需要Check
函数,可以直接通过层级名字获取。由于层级描述已经统一,这里的实现过程较上一节中的创建器更加简单,用户无需额外关注。
- 当
- 网络层插件
- 对于一些特殊层,我们使用了插件形式对这些层进行了支持。开发新的插件一般需要满足以下条件之一:
- 使用
TensorRT
自带的Layer
不能支持或拼接表示出需要的层; - 使用
TensorRT
自带的Layer
实现效率较低; - 使用
TensorRT
自带的Layer
进行拼接表示的流程过长,既不利于提速,又不利于后续开发维护。
- 使用
- 在我们的测试中,
Gather
层用于Embedding
时的表现较差,另外Padding
层支持的模式有限,基于类似的原因开发或加入了多款插件,详见目录source\trt_engine\trt_network_crt\plugins
。
- 对于一些特殊层,我们使用了插件形式对这些层进行了支持。开发新的插件一般需要满足以下条件之一:
- 与
TensorFlow
模型相关- 对于图片输入的 CV 网络,
TF
默认的输入格式为NHWC
,因此构建网络时,我们将四维的输入做了一个转置,使其输入格式变为NCHW
。这个改动对于其他的四维输入会有不正确的影响,需要在未来修复。 - 对于
IteratorGetNext
这样单个结点中有多个输入的情况,会用专门的负数input_reference_
标记其位置。当其被引用后恢复为正数,这样能够保证后续查找UnusedInputs
能正常工作。同时,在调用ParseOperaion
函数时也要对其进行专门处理,以保证获得正确的输入对应。 TensorFlow
网络中的weights
利用TF
的TF_TryEvaluateConstant
API 进行获取(调用链:tf::Output::GetConstantTensor->TF_TryEvaluateConstant
)。对于具有较大参数的TF
网络,请使用ckpt_to_pb.py
进行模型转换,脚本会自动将较大的参数单独提取成文件,避免读入此类模型时出现内存相关的问题。
- 对于图片输入的 CV 网络,
- 与
PyTorch
模型相关PyTorch
的 C++ API 提供了FuseLinear
的优化,可以帮助将线性层合并,在分析网络时首先使用该优化可以降低分析的复杂程度。在更高的PyTorch
版本(>=1.7)中还有Inline
和FoldFloorDivide
等其他优化,也可以使用。- 在进行网络分析前需要调用
GraphUtils::EvalAll
函数,该函数依次将每个结点作为输出,然后运行网络,以此得到每个结点的输出值。这一操作的主要目的是获取PyTorch
各个常量层的数据值。对于大型网络,这一步骤需要等待一定的时间。 - 在调用完
GraphUtils::EvalAll
之后还需要调用GraphUtils::RemoveRedundantNodes
函数来移除部分冗余结点,这部分结点非常影响网络分析,详见代码内的注释。
- 与
ONNX
模型相关OnnxBuilder
在解析 ONNX 模型时,直接调用了NvOnnxParser.h
中提供的nvonnxparser::createParser
和nvonnxparser::IParser::parseFromFile
接口,省略了上述通用模式中的第一个步骤和第二个步骤中生成网络定义的环节,因为这些接口可以直接返回原模型的网络定义nvinfer1::INetworkDefinition
。OnnxBuilder
在获得原模型的网络定义后,将直接调用nvinfer1::builder::buildEngine
生成 TRT 引擎。