TensorRT教程15：使用Python API 添加自定义插件层

使用Python API 添加自定义插件层

示例：使用 python 向网络中添加自定义层

import tensorrt as trt
import numpy as np

#step1:创建logger：日志记录器
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():
    #创建builder,network
    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插件层
        lrelu = network.add_plugin_v2(inputs=[input_layer], plugin=get_trt_plugin("LReLU_TRT"))
        #设置输出层name
        lrelu.get_output(0).name = "outputs"
        #网络标记输出层
        network.mark_output(lrelu.get_output(0))

示例：使用 Python 添加 UFF 中不支持的自定义层

官方示例：

https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html#uff_custom_plugin

import tensorrt as trt
import numpy as np

#step1:创建logger：日志记录器
TRT_LOGGER = trt.Logger()
#初始化注册插件层
trt.init_libnvinfer_plugins(TRT_LOGGER, '')

#准备网络并检查 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()

#准备命名空间映射
trt_lrelu = gs.create_plugin_node(name="trt_lrelu", op="LReLU_TRT", negSlope=lrelu_alpha)
namespace_plugin_map = { "tf_lrelu": trt_lrelu }

#转换 TensorFlow 使用 GraphSurgeon 图并保存到 UFF
dynamic_graph = gs.DynamicGraph(tf_lrelu.graph)
dynamic_graph.collapse_namespaces(namespace_plugin_map)

#运行 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)

示例：使用 Python API 实现BatchNorm算子

批量归一化层：BatchNorm

通过TRT支持的scale层来实现

#---------------BatchNorm层---------------
#获取训练后的BN相关参数
g0 = params['batchnorm0_gamma'].asnumpy().reshape(-1)
m0 = extra_params['batchnorm0_moving_mean'].asnumpy().reshape(-1)
v0 = extra_params['batchnorm0_moving_var'].asnumpy().reshape(-1)

#注意一定要设置为1e-3，如果设置为1e-4等更小的，输出的结果就和tf不完全一样，这是tf特色
scale0 = g0 / np.sqrt(v0 + 1e-3)
shift0 = -m0 / np.sqrt(v0 + 1e-3) * g0 + params['batchnorm0_beta'].asnumpy().reshape(-1)
power0 = np.ones(len(g0), dtype=np.float32)
bag += [scale0 shift0 power0]
#添加BN层
batchNormLayer= network.add_scale(layer.get_output(0), trt.ScaleMode.CHANNEL,shift0, scale0, power0)