PluginV2

向 TensorRT 添加具有自定义功能的节点。

另请参阅

使用自定义层扩展 TensorRT

属性

插件特定属性。

输入

插件特定输入。

输出

插件特定输出。

数据类型

插件特定数据类型。

形状信息

插件特定形状。

体积限制

插件特定限制。

DLA 支持

插件特定限制。

示例

插件

def init_gelu_plugin(bias):
    for plugin_creator in trt.get_plugin_registry().plugin_creator_list:
        if plugin_creator.name == "CustomGeluPluginDynamic":
            gelu_data_type = trt.PluginField("type_id", np.array([0], dtype=np.int32), type=trt.PluginFieldType.INT8)
            gelu_bias = trt.PluginField("bias", bias, type=trt.PluginFieldType.FLOAT32)
            field_collection = trt.PluginFieldCollection([gelu_data_type, gelu_bias])
            return plugin_creator.create_plugin(name="CustomGeluPluginDynamic", field_collection=field_collection)

    raise Exception('CustomGeluPluginDynamic Plugin not found')

bias = np.array([[[1, -0.5, 0, 0, 0.7]]], dtype=np.float32)

input_layer = network.add_input(name="input_layer", dtype=trt.float32, shape=(1, 1, 5))
lrelu = network.add_plugin_v2(inputs=[input_layer], plugin=init_gelu_plugin(bias))
network.mark_output(lrelu.get_output(0))

inputs[input_layer.name] = np.array(
    [[
        [3.0, -4.3, 22.8, -2.97, 143.2],
    ]]
)
outputs[lrelu.get_output(0).name] = lrelu.get_output(0).shape

input_with_bias = inputs[input_layer.name] + bias
expected_output = 0.5 * input_with_bias * (1 + np.tanh(np.sqrt(2 / math.pi) * (input_with_bias + 0.044715 * np.power(input_with_bias, 3))))

expected[lrelu.get_output(0).name] = expected_output

C++ API

有关 C++ IPluginV2Layer 算子的更多信息，请参阅 C++ IPluginV2Layer 文档。

Python API

有关 Python IPluginV2Layer 算子的更多信息，请参阅 Python IPluginV2Layer 文档。