tensorrt.plugin.autotune

tensorrt.plugin.autotune(plugin_id: str) → Callable

包装一个函数，为已通过 trt.plugin.register 注册的插件定义自动调优逻辑。

自动调优是 TensorRT 在 IO 类型/格式组合以及插件声明为支持的任何自定义策略上执行插件的过程。延迟最低的（类型、格式、策略）组合用于在引擎构建后执行插件。

注意

自动调优函数是可选的。如果未指定，TensorRT 将假定插件仅支持在网络创建时指定的输入类型、通过 trt.plugin.register 指定的输出类型以及所有 I/O 的线性格式。

此 API 仅旨在用作装饰器。装饰函数不需要为输入参数或返回值提供类型提示；但是，指定的任何类型提示都将根据 trt.plugin.register 签名进行验证以确保一致性。

该函数的模式如下

(inp0: TensorDesc, inp1: TensorDesc, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, outputs: Tuple[TensorDesc]) -> List[AutoTuneCombination]

• 首先传递输入张量，每个张量都由 TensorDesc 描述。

• 接下来声明插件属性。并非所有包含在 trt.plugin.register 中的属性都必须在此处指定 – 它们可以是子集。

• 该函数应返回 AutoTuneCombination 的列表。

参数:

plugin_id – 插件的 ID，格式为“{namespace}::{name}”，必须与 trt.plugin.register 期间使用的 ID 匹配

一个 elementwise add 插件，它同时支持 FP32 和 FP16 线性 I/O，并且希望针对 2 个自定义策略进行调优。

import tensorrt.plugin as trtp

@trtp.register("my::add_plugin")
def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
    return inp0.like()

@trtp.autotune("my::add_plugin")
def add_plugin_autotune(inp0: trtp.TensorDesc, block_size: int, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:

    return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR", [1, 2])]

与上面的示例相同，但使用 AutoTuneCombination 的逐索引构造

import tensorrt.plugin as trtp

@trtp.register("my::add_plugin")
def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
    return inp0.like()

@trtp.autotune("my::add_plugin")
def add_plugin_autotune(inp0: trtp.TensorDesc, block_size: int, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
    c = trtp.AutoTuneCombination()
    c.pos(0, "FP32|FP16", "LINEAR")
    c.pos(1, "FP32|FP16") # index 1 is the output. Omitting format is the same as declaring it to be LINEAR.
    c.tactics([1, 2])
    return [c]

另请参阅

• AutoTuneCombination