使用 TensorRT Python 插件编写自定义算子¶

TensorRT (TRT) 提供了广泛的内置算子，可以满足大多数用例。但是，您可能希望出于多种原因引入自定义算子，包括

支持 TRT 开箱即用不支持的全新算子
为已支持的算子添加行为略有不同的算子

TRT 通过插件提供对自定义算子的支持。本指南展示了如何实现和包装定义插件行为的 Python 函数，以便可以将它们作为自定义算子添加到网络中。

插件的组成¶

首先，插件定义需要两个函数，并使用 tensorrt.plugin 模块提供的装饰器进行包装

tensorrt.plugin.register()：返回输出张量的形状和类型特征。函数签名还定义了输入张量和插件运行所需的任何属性。
tensorrt.plugin.impl()：执行插件计算

可选地，如果插件能够支持其 I/O 的多种数据类型/张量布局组合，或者能够支持多个后端（例如内核），它还可以利用 TensorRT 的自动调优功能，使其能够在目标平台上选择插件的最佳性能配置。要启用自动调优，只需定义一个用 tensorrt.plugin.autotune() 包装的函数。

让我们考虑几个示例来理解这些函数的具体细节。

示例：循环填充插件¶

循环填充对于循环卷积等操作很有用。下图表示原始图像（红色）如何循环填充一次（绿色）和两次（蓝色）

alt text

插件应具有以下特性

输入：\(N\) 维输入
属性：\(m\) 维参数 pads，其中 \(m\) 为偶数且 \(m/2 \le N\)。pads 表示在输入张量的最后 \(m/2\) 个维度中的每个维度之前和之后应用的填充量。
输出：填充后的张量，其数据类型与输入相同。输出的第 i 个倒数维度 \(d^{\text{out}}_{N-i-1}\) 可以用相应的输入维度 \(d^{\text{in}}_{N-i-1}\) 表示为

\[ d^{\text{out}}_{N-i-1} = d^{\text{in}}_{N-i-1} + \text{pads}_{2i} + \text{pads}_{2i + 1} \]

让我们在 tensorrt.plugin.register 函数中捕获此信息

import tensorrt.plugin as trtp
import numpy.typing as npt

@trtp.register("example::circ_pad_plugin")
def circ_pad_plugin_desc(
    inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32]
) -> trtp.TensorDesc:
    ndim = inp0.ndim
    out_desc = inp0.like()

    for i in range(np.size(pads) // 2):
        out_desc.shape_expr[ndim - i - 1] += int(
            pads[i * 2] + pads[i * 2 + 1]
        )

    return out_desc

参数 example::circ_pad_plugin 定义了插件的命名空间（“example”）和名称（“circ_pad_plugin”）。使用 tensorrt.plugin.TensorDesc 注释的输入参数表示输入张量；所有其他参数都被解释为插件属性。支持的属性类型包括

int、float、str、bool、bytes。不支持这些类型的列表/元组。
以下类型的 1-D Numpy 数组：int8、int16、int32、int64、float16、float32、float64、bool。这些必须使用 numpy.typing.NDArray[dtype] 注释，其中 dtype 是预期的 Numpy 数据类型。

输出签名是描述输出的 trt.plugin.TensorDesc。要构造输出张量描述符，我们从 inp0.like() 开始，它返回一个与 inp0 具有相同形状和类型特征的张量描述符。由于输出形状实际上与输入不同，因此可以通过 tensorrt.plugin.TensorDesc.shape_expr 访问输出形状的符号表达式，并将预期形状写入其中。

现在让我们定义计算函数，并使用 tensorrt.plugin.impl() 进行装饰。为简单起见，让我们利用 PyTorch 的 torch.nn.functional.pad 来进行计算。

import tensorrt.plugin as trtp

@trtp.impl("example::circ_pad_plugin")
def circ_pad_plugin_impl(
    inp0: trtp.Tensor,
    pads: npt.NDArray[np.int32],
    outputs: Tuple[trtp.Tensor],
    stream: int
) -> None:
    inp_t = torch.as_tensor(inp0, device="cuda")
    out_t = torch.as_tensor(outputs[0], device="cuda")

    out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
    out_t.copy_(out)

请注意，装饰后的函数接收每个输入和输出的 tensorrt.plugin.Tensor。与 TensorDesc 不同，Tensor 引用底层数据缓冲区，可以直接通过 tensorrt.plugin.Tensor.data_ptr() 访问。当使用 Torch 和 OpenAI Triton 内核时，更容易使用 torch.as_tensor() 零拷贝构造与 tensorrt.plugin.Tensor 对应的 torch.Tensor。

选择最佳性能插件配置：自动调优¶

假设插件能够支持 FP32 和 FP16 I/O，以及线性张量布局（张量格式）。如果性能是关键，并且不确定 FP32 或 FP16 在目标平台上哪个性能更好，我们可以定义一个用 tensorrt.plugin.autotune() 装饰的函数。

@trtp.autotune("example::circ_pad_plugin")
def circ_pad_plugin_autotune(
    inp0: trtp.TensorDesc,
    pads: npt.NDArray[np.int32],
    outputs: Tuple[trtp.TensorDesc],
) -> List[trtp.AutoTuneCombination]:
    return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR")]

装饰后的函数必须返回 tensorrt.plugin.AutoTuneCombination 的列表。在本例中，我们定义一个组合 AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR")；这表明输入和输出必须都是 FP32 或 FP16，并且每个都具有线性格式。当定义自动调优函数时，在引擎构建期间，TRT 可能会针对每种类型/格式组合使用随机输入执行插件（调用其 impl 函数）。

将插件添加到 TensorRT 网络¶

现在我们已经定义了插件，就可以将其添加到 TensorRT 网络中。根据您的工作流程，您可能需要通过 TRT Python API 将插件直接添加到 TRT tensorrt.INetworkDefinition，或者您可能希望在加载包含自定义算子的 ONNX 图时通过 TRT ONNX 解析器添加插件。

使用 TRT Python API 添加插件¶

tensorrt.INetworkDefinition.add_plugin() API 可用于将插件添加到网络定义（tensorrt.INetworkDefinition 的实例）

input_tensor = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
plugin_layer = network.add_plugin(trt.plugin.op.example.circ_pad_plugin(input_tensor, pads = pads))

请注意，已注册的插件可在 tensorrt.plugin.op 下其注册的命名空间和名称中找到。

加载带有自定义算子的 ONNX 模型¶

可以加载带有自定义 op 节点的 ONNX 模型，您需要通过 TRT 插件运行该模型。为了使 TRT ONNX 解析器能够正确识别您的插件映射到感兴趣的 ONNX 节点，请确保

ONNX 节点的 op 属性与您的插件名称完全相同。
节点包含一个名为 plugin_namespace 的字符串属性，其中包含您的插件的命名空间。

例如，如果使用 ONNX Graphsurgeon，则可以按如下方式构造自定义 op 节点

import onnx_graphsurgeon as gs

var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
var_y = gs.Variable(name="y", dtype=np.float32)

circ_pad_node = gs.Node(
    name="circ_pad_plugin",
    op="circ_pad_plugin",
    inputs=[var_x],
    outputs=[var_y],
    attrs={"pads": pads, "plugin_namespace": "example"},
)

高级用法¶

示例：具有数据相关输出形状的算子 - Non-zero¶

Non-zero 是一种操作，用于查找输入张量的非零元素的索引。因此，它具有数据相关的输出形状 (DDS)，典型的形状计算无法通过输入形状完成。

为了处理 DDS，每个数据相关输出维度的范围必须用大小张量表示，大小张量是一个标量，它以输入形状的形式向 TRT 传达该维度的上限和自动调优值。TRT 引擎构建可能会针对自动调优值进行优化，但该维度的范围可能会在运行时扩展到上限。

在本示例中，我们考虑一个 2D 输入张量 inp0；输出将是一个 \(N \times 2\) 张量（一组 \(N\) 个 2D 索引），其中 \(N\) 是非零索引的数量。最多，所有元素都可能为非零，因此上限可以表示为 upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]。

平均而言，我们可以预期一半的输入将填充为零，因此可以使用该值作为自动调优值来构造大小张量

st = trt.plugin.size_tensor(opt = upper_bound // 2, upper_bound = upper_bound)

现在我们准备好构造输出形状。st.expr() 返回大小张量的形状表达式，因此可以按 trt.plugin.from_shape_expr((st.expr(), 2), dtype=trt.int32) 的形式构造输出形状的张量描述符。TRT 要求任何大小张量也必须作为插件的输出。将它们放在一起，我们得到以下结果

import tensorrt.plugin as trtp

@trtp.register("example::non_zero_plugin")
def non_zero_plugin_desc(
    inp0: trtp.TensorDesc,
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
    upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]
    st = trtp.size_tensor(upper_bound // 2, upper_bound)
    return trtp.from_shape_expr((st.expr(), 2), dtype=trt.int32), st

示例：使用 I/O 别名实现原地自定义操作¶

原地计算可以通过 TRT 插件通过别名 I/O 完成。即，需要原地修改的输入可以由输入-输出对表示，其中输出别名为输入。例如，考虑一个简单的逐元素加法插件。如果需要原地加法，可以按如下方式实现

import tensorrt.plugin as trtp

@trtp.register("sample::elemwise_add_plugin_")
def add_plugin_desc_(inp0: trtp.TensorDesc) -> trtp.TensorDesc:
    return inp0.aliased()

请注意，inp0.aliased() 生成一个别名为 inp0 的输出 TensorDesc。

示例：具有多个后端的插件：使用自定义策略¶

可能存在多种内核或库可用于执行插件计算（对于相同的 IO 数据类型/张量布局组合），但无法预先确定哪个后端在目标平台上最快。这种备用后端称为插件的策略。为了让 TRT 找出最快的策略，可以使用自动调优功能。

假设我们有一个 OpenAI Triton 内核 circ_pad_kernel，它可以执行上述示例中的循环填充操作。我们可以要求 TRT 在 OpenAI Triton 内核和 torch.nn.functional.pad 之间选择最快的，如下所示

import tensorrt.plugin as trtp
from enum import IntEnum

class Tactic(IntEnum):
    TORCH = 1
    TRITON = 2

@trt.plugin.autotune("sample::circ_pad_plugin")
def circ_pad_plugin_autotune(inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
    c = trtp.AutoTuneCombination()
    c.pos([0, 1], "FP32|FP16")
    c.tactics([int(Tactic.TORCH), int(Tactic.TRITON)])
    return [c]

请注意，我们在此处使用另一种构造 trt.plugin.AutoTuneCombination 的方法 – 即，通过 pos(...) 填充类型/格式信息，并通过 tactics(...) 指定策略。

现在，impl 函数可以接受一个额外的 int 参数 tactic 以使用适当的后端

@trtp.impl("example::circ_pad_plugin")
def circ_pad_plugin_impl(
    inp0: trtp.Tensor,
    pads: npt.NDArray[np.int32],
    outputs: Tuple[trtp.Tensor],
    stream: int,
    tactic: int
) -> None:
    inp_t = torch.as_tensor(inp0, device="cuda")
    out_t = torch.as_tensor(outputs[0], device="cuda")

    if tactic == Tactic.TORCH:
        out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
        out_t.copy_(out)
    elif tactic == Tactic.TRITON:
        N = inp0.ndim
        all_pads = np.zeros((N * 2,), dtype=np.int32)
        out_dims = trtp.Shape(tuple(inp0.shape))

        block_size = 256
        num_blocks = tuple(
            [int((np.prod(out_dims) + block_size - 1) // block_size)]
        )

        circ_pad_kernel[num_blocks](inp_t, ..., BLOCK_SIZE=block_size)

FAQ¶

插件创建器、插件注册和插件库发生了什么变化？

如果您之前处理过基于类的 TRT 插件，您可能会有这个问题。对于基于装饰器的插件，只要定义了 tensorrt.plugin.register 函数，tensorrt.plugin 模块就会自动处理注册。

为了模拟插件库的体验，您可以在单独的 Python 模块上的公共命名空间下定义插件，然后在需要加载该插件“库”时加载该模块。

如何处理插件属性的序列化/反序列化？

tensorrt.plugin 模块会自动将 tensorrt.plugin.impl 函数签名中包含的任何插件属性序列化到 TRT 引擎中。加载该引擎时，这些属性将被反序列化并传递回 tensorrt.plugin.impl 函数。

注意

如果插件计算仅需要 tensorrt.plugin.register 中包含的插件属性的子集，则仅将这些属性包含在 tensorrt.plugin.impl 函数签名中。这应避免不必要的数据序列化，从而生成更精简的 TRT 引擎。

局限性¶

在 Python 中定义的 TRT 插件会导致 TRT 引擎在执行时依赖于具有这些插件定义的 Python 环境的可用性。要构建独立于 Python 的 TRT 引擎，建议使用 TRT C++ 插件接口。
本指南介绍了通过 tensorrt.plugin 模块实现的插件，该模块支持插件使用的一些最常见用例。对于更高级的用例，例如传递形状输入，建议将插件定义为 tensorrt.IPluginV3 接口的直接实现。

其他资源¶

使用自定义层扩展 TensorRT