TensorRT Developer Guide - chapter4

Chapter 4. EXTENDING TENSORRT WITH CUSTOM LAYERS

TensorRT支持多种类型的层，并且其功能不断扩展。但是，可能存在支持的层不满足模型的特定需求的情况。在这种情况下，用户可以通过使用C ++和Python API的IPluginV2类实现自定义层来扩展TensorRT功能。自定义层（通常称为插件）由应用程序实现和实例化，它们的生命周期必须跨越TensorRT引擎中的使用范围。

4.1 使用C ++ API添加自定义层

通过扩展IPluginV2和IPluginCreator类来实现自定义层。

IPluginV2：

IPluginV2是您应该为插件实现的基类。它包括版本控制支持，并帮助启用支持除NCHW和单精度之外的其他数据格式的自定义图层。

IPluginCreator：

IPluginCreator是自定义图层的创建者类，用户可以使用它来获取插件名称、版本和插件字段参数。它还提供了在网络构建阶段创建插件对象的方法，并在推理期间对其进行反序列化。

在以前版本的TensorRT中，您为自定义图层实现了IPluginExt。虽然仍然支持此API，但我们强烈建议您转移到IPluginV2以便能够使用所有新的插件功能。

TensorRT还提供了通过调用REGISTER_TENSORRT_PLUGIN（pluginCreator）来注册插件的功能，将插件创建者静态注册到插件注册表。在运行时，可以使用extern函数getPluginRegistry（）查询插件注册表。插件注册表存储指向所有已注册的插件创建器的指针，可用于根据插件名称和版本查找特定的插件创建器。TensorRT库包含可以加载到您的应用程序中的插件。所有这些插件的版本都设置为1.这些插件的名称是：
‣ RPROI_TRT
‣ Normalize_TRT
‣ PriorBox_TRT
‣ GridAnchor_TRT
‣ NMS_TRT
‣ LReLU_TRT
‣ Reorg_TRT
‣ Region_TRT
‣ Clip_TRT

要在应用程序中使用TensorRT注册的插件，必须加载libnvinfer_plugin.so库并且必须注册所有插件。这可以通过在应用程序代码中调用initLibNvInferPlugins（void * logger，const char * libNamespace）（）来完成。

如果您有自己的插件库，则可以包含一个类似的入口点，以便在注册表中的唯一命名空间下注册所有插件。这可确保在不同插件库的构建时间内没有插件名称冲突。

有关这些插件的更多信息，请参阅NvInferPlugin.h文件参考。

使用Plugin Creator，可以调用IPluginCreator :: createPlugin（）函数，该函数返回IPluginV2类型的插件对象。可以使用addPluginV2（）创建并添加层到TensorRT网络，然后将层绑定到给定的插件。该方法还返回一个指向层（IPluginV2Layer类型）的指针，该指针可用于访问层或插件本身（通过getPlugin（））。

例如，要将名称为pluginName、版本为pluginVersion的插件层添加到网络中，可以使用以下指令：

//Use the extern function getPluginRegistry to access the global TensorRT Plugin Registry
auto creator = getPluginRegistry()->getPluginCreator(pluginName, pluginVersion);
const PluginFieldCollection* pluginFC = creator->getFieldNames();

//populate the field parameters (say layerFields) for the plugin layer
PluginFieldCollection *pluginData = parseAndFillFields(pluginFC, layerFields);

//create the plugin object using the layerName and the plugin meta data
IPluginV2 *pluginObj = creator->createPlugin(layerName, pluginData);

//add the plugin to the TensorRT network using the network API
auto layer = network.addPluginV2(&inputs[0], int(inputs.size()), pluginObj);

… (build rest of the network and serialize engine)
pluginObj->destroy() // Destroy the plugin object
… (destroy network, engine, builder)
… (free allocated pluginData)

pluginData应该在传递给createPlugin之前在堆上分配PluginField条目。

上面的createPlugin方法将在堆上创建一个新的插件对象并返回指向它的指针。确保销毁pluginObj，如上所示，以避免内存泄漏。

在序列化期间，TensorRT引擎将在内部存储所有IPluginV2类型插件的插件类型、插件版本和命名空间（如果存在）。在反序列化期间，TensorRT引擎会查找此信息，以便从插件注册表中找到插件创建器。这使TensorRT引擎能够在内部调用IPluginCreator :: deserializePlugin（）方法。反序列化期间创建的插件对象将由TensorRT引擎通过调用IPluginV2 :: destroy（）方法在内部销毁。

在以前的TensorRT版本中，您必须实现nvinfer1 :: IPluginFactory类以在反序列化期间调用createPlugin方法。对于使用TensorRT注册并使用addPluginV2添加的插件，不再需要这样做。

4.1.1 示例1：使用C ++为Caffe添加自定义层

要在C ++中添加自定义层，请实现IPluginExt类。对于基于Caffe的网络，如果使用TensorRT Caffe Parser，您还将实现nvcaffeparser1 :: IPluginFactoryExt（对于IPluginExt类型的插件）和nvinfer1 :: IPluginFactory类。有关更多信息，请参阅4.3节从框架导入模型时使用自定义图层。

以下示例代码添加了一个名为FooPlugin的新插件：

class FooPlugin : public IPluginExt
{
...implement all class methods for your plugin
};
class MyPluginFactory : public nvinfer1::IPluginFactory, public
nvcaffeparser1::IPluginFactoryExt
{
...implement all factory methods for your plugin
};

如果您使用的是在IPluginV2类型的TensorRT插件注册表中注册的插件，那么您不需要实现nvinfer1 :: IPluginFactory类。但是，您确实需要实现nvcaffeparser1 :: IPluginFactoryV2和IPluginCreator类并注册它们。

class FooPlugin : public IPluginV2
{
...implement all class methods for your plugin
};
class FooPluginFactory : public nvcaffeparser1::IPluginFactoryV2
{
virtual nvinfer1::IPluginV2* createPlugin(...)
{
...create and return plugin object of type FooPlugin
}
bool isPlugin(const char* name)
{
...check if layer name corresponds to plugin
}
}
class FooPluginCreator : public IPluginCreator
{
...implement all creator methods here
};
REGISTER_TENSORRT_PLUGIN(FooPluginCreator);

以下示例说明了如何使用C ++为Caffe网络添加自定义插件层：

samplePlugin有一个用户实现的插件；
sampleFasterRCNN使用在TensorRT插件注册表中注册的插件。

4.1.2 示例2：使用C++添加一个UFF不支持的自定义层

要使用TensorRT运行TensorFlow网络，必须先将其转换为UFF格式。

以下步骤在C ++中为TensorFlow网络添加自定义插件层：

1.实现IPluginV2和IPluginCreator类，如4.1.1所示；
2.将TensorFlow操作映射到插件操作。你可以使用graphurgeon。例如，请参阅以下代码段以将TensorFlow Relu6操作映射到插件：

import graphsurgeon as gs
my_relu6 = gs.create_plugin_node(name=”MyRelu6”, op=”Clip_TRT”, clipMin=0.0,
clipMax=6.0)
Namespace_plugin_map = { “tf_relu6” : my_relu6 }
def preprocess(dynamic_graph):
dynamic_graph.collapse_namespaces(namespace_plugin_map)

在上面的代码中，tf_relu6是TensorFlow图中Relu6节点的名称。它将tf_relu6节点映射到具有操作“Clip_TRT”的自定义插件节点，该操作是要使用的插件的名称。将上面的代码保存到名为config.py的文件中。如果插件层需要参数，则应将它们作为参数传递给gs.create_plugin_node。在这种情况下，clipMin和clipMax是剪辑插件所期望的参数。

3.使用preprocess -p标志设置调用UFF转换器：

convert-to-uff frozen_inference_graph.pb -p config.py -t

这将生成一个UFF文件，其中TensorFlow操作被TensorRT插件节点替换。

4.进行预处理并且使用TensorRT UFF解析器将其转换为UFF文件。更多细节请参考4.3节。sampleUffSSD示例说明了如何使用C ++添加UFF不支持的自定义层。有关如何预处理图形的演示，请参阅示例文件夹中的config.py.

4.2 使用Python API添加自定义层

虽然C ++ API是实现自定义图层的首选语言;但是由于可以轻松访问CUDA和cuDNN等库，因此您还可以在Python应用程序中使用自定义层。

您可以使用C ++ API创建自定义层，在Python中使用pybind11打包层，然后将插件加载到Python应用程序中。有关更多信息，请参阅3.2节在Python中创建网络定义。

相同的自定义层实现可以用于C ++和Python。有关更多信息，请参阅位于/ usr / src / tensorrt / samples / fc_plugin_caffe_mnist /目录中的fc_plugin_caffe_mnist Python示例。

4.2.1 示例1：使用Python将自定义层添加到TensorRT网络

可以使用插件节点将自定义层添加到Python中的任何TensorRT网络中。Python API有一个名为add_plugin_v2的函数，它允许您将插件节点添加到网络中。以下示例说明了这一点。它创建了一个简单的TensorRT网络，并通过查找TensorRT插件注册表添加了一个Leaky ReLU插件节点。

import tensorrt as trt
import numpy as np

TRT_LOGGER = trt.Logger()
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list

def get_trt_plugin(plugin_name):
	plugin = None
	for plugin_creator in PLUGIN_CREATORS:
		if plugin_creator.name == plugin_name:
			lrelu_slope_field = trt.PluginField("neg_slope", np.array([0.1],dtype=np.float32), trt.PluginFieldType.FLOAT32)
			field_collection =trt.PluginFieldCollection([lrelu_slope_field])
			plugin = plugin_creator.create_plugin(name=plugin_name,field_collection=field_collection)
	return plugin

def main():
	with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
		builder.max_workspace_size = 2**20
		input_layer = network.add_input(name="input_layer", dtype=trt.float32,shape=(1, 1))
		lrelu = network.add_plugin_v2(inputs=[input_layer],plugin=get_trt_plugin("LReLU_TRT"))
		lrelu.get_output(0).name = "outputs"
		network.mark_output(lrelu.get_output(0))

4.2.2 示例2：使用Python添加一个UFF不支持的自定义层

TensorFlow网络可以转换为UFF格式，并使用Python接口与TensorRT一起运行。为此，我们使用了graphurgeon API。如果您正在编写自己的插件，则需要通过实现IPluginExt和IPluginCreator类在C ++中实现它，如4.1.1节示例1：使用C ++为Caffe添加自定义层所示。

以下步骤说明如何使用UFF Parser使用在TensorRT插件注册表中注册的插件节点运行自定义层。

1.通过调用trt.init_libnvinfer_plugins（TRT_LOGGER，’’）注册TensorRT插件（或者在已注册自己的插件的地方加载.so文件）。
2.准备网络并检查TensorFlow输出：

tf_sess = tf.InteractiveSession()
tf_input = tf.placeholder(tf.float32, name="placeholder")
tf_lrelu = tf.nn.leaky_relu(tf_input, alpha=lrelu_alpha, name="tf_lrelu")
tf_result = tf_sess.run(tf_lrelu, feed_dict={tf_input: lrelu_args})
tf_sess.close()

3.准备命名空间映射。将操作LReLU_TRT映射到TensorRT附带的Leaky ReLU插件。

trt_lrelu = gs.create_plugin_node(name="trt_lrelu", op="LReLU_TRT",negSlope=lrelu_alpha)
namespace_plugin_map = {"tf_lrelu": trt_lrelu}

4.使用graphurgeon转换TensorFlow图并保存到UFF：

dynamic_graph = gs.DynamicGraph(tf_lrelu.graph)
dynamic_graph.collapse_namespaces(namespace_plugin_map)

5.运行UFF解析器并将结果与​​TensorFlow进行比较：

uff_model = uff.from_tensorflow(dynamic_graph.as_graph_def(), ["trt_lrelu"],output_filename=model_path,text=True)
parser = trt.UffParser()
parser.register_input("placeholder", [lrelu_args.size])
parser.register_output("trt_lrelu")
parser.parse(model_path, trt_network)

有关更多信息，请参阅uff_custom_plugin示例。

4.3 从框架导入模型时使用自定义图层

TensorRT解析器使用层操作字段来识别网络中的特定层是否是TensorRT支持的操作。

Tensorflow:
与先前版本的TensorRT相比，本版本对TensorFlow中的自定义层如何使用TensorRT UFF解析器进行了更改。对于TensorFlow模型，使用UFF转换器将图转换为UFF文件。在此过程中，如果网络包含插件层，则还需要将这些层的操作字段映射到TensorRT中相应的已注册插件。这些插件既可以是TensorRT附带的插件，也可以是您编写的自定义插件。网络中的插件字段名称也应该与插件期望的字段匹配。这可以使用graphurgeon来完成，如8.5节使用Graph Surgeon API预处理TensorFlow图中所解释的那样，并且通过使用带有UFF转换器的配置文件在sampleUffSSD中进行了演示。

UFF Parser将为每个不受支持的操作查找插件注册表。如果它找到与任何已注册的插件名称匹配，则解析器将从输入网络解析插件字段参数并使用它们创建插件对象。然后将此对象添加到网络中。在以前的TensorRT版本中，您必须实现nvuffparser :: IPluginFactoryExt并手动将插件参数传递给createPlugin（…）函数。尽管仍然可以使用此流程，但不再需要添加Plugin API。有关更多信息，请参阅：
‣ IPluginV2 and IPluginCreator in the C++ API
‣ IPluginV2 and IPluginCreator in the Python API

caffe：
对于Caffe模型，请使用nvcaffeparser1 :: IPluginFactoryV2类。解析器的setPluginFactoryV2方法在解析器中设置工厂以启用自定义层。在解析模型描述时，对于每个层，解析器调用isPluginV2来检查工厂是否对应于自定义层;如果是这样，解析器实例化插件，使用层的名称、模型权重和相应参数调用createPlugin（以便工厂可以实例化相应的插件）。如果单个工厂与不同的图层名称相关联，则单个工厂可以支持的插件数量没有限制。

对于Caffe解析器，如果使用setPluginFactoryV2和IPluginFactoryV2，则在反序列化期间创建的插件对象将由引擎通过调用IPluginExt :: destroy（）在内部销毁。您只负责销毁在网络创建步骤中创建的插件对象，如4.1.1节使用C ++ API添加自定义图层中所示。

samplePlugin示例说明了如何扩展nvcaffeparser1 :: IPluginFactoryExt以使用自定义图层，而sampleUffSSD使用UFF Parser来使用自定义图层。

有关使用TensorRT的自定义图层的Python用法，请参阅Caffe网络的fc_plugin_caffe_mnist示例，以及UFF网络的uff_custom_plugin和uff_ssd示例。

4.3.1 示例1：将自定义层添加到TensorFlow模型

要使用TensorRT运行TensorFlow网络，必须先将其转换为UFF格式。在转换过程中，可以使用graphurgeon实用程序将自定义图层标记为插件节点。

然后，UFF转换器将处理后的图形转换为UFF格式，然后由UFF分析器运行。然后，UFF Parser将插件节点添加到TensorRT网络。

有关使用C ++ API的详细信息，请参阅示例2：使用C ++添加UFF中不支持的自定义层。
有关使用Python API的详细信息，请参阅示例2：使用Python添加UFF中不支持的自定义层。此外，uff_ssd Python示例演示了Python中用于使用TensorRT运行TensorFlow目标检测网络的端到端工作流。

4.4 插件API

所有新插件都应该实现IPluginV2和IPluginCreator类。此外，新插件还应该调用REGISTER_TENSORRT_PLUGIN（…）宏来向TensorRT插件注册表注册插件，或者创建一个等同于initLibNvInferPlugins（）的init函数。

4.4.1 将插件从TensorRT 5.0.0 RC迁移到TensorRT 5.0.x

已添加IPluginV2插件类。 IPluginExt类已恢复为与4.0 GA中的版本兼容。

4.5 自定义图层的最佳实践

用户自定义层：
要将自定义图层实现创建为TensorRT插件，您需要为插件实现IPluginV2类和IPluginCreator类。

有关这两个API类的详细信息，请参阅4.4节插件API说明。

对于Caffe网络，请参阅示例1：使用C ++为Caffe添加自定义层。
对于TensorFlow（UFF）网络，请参见示例2：使用C ++添加UFF中不支持的自定义层。

使用UFF注册API：
有关如何在C ++和Python中使用UFF的插件的示例，请参阅示例1：使用C ++为Caffe添加自定义层和示例2：使用Python添加UFF中不支持的自定义层。

自定义层调试问题：
必须释放插件中分配的内存以确保没有内存泄漏。如果在initialize（）函数中获取资源，则需要在terminate（）函数中释放它们。应该优先在插件类析构函数或destroy（）方法中释放所有其他内存分配。4.1节使用C ++ API添加自定义图层详细概述了这一点，并提供了使用插件时的最佳实践的一些注意事项。