AutoTuneCombination

class tensorrt.plugin.AutoTuneCombination(io_types: str | None = None, layouts: str | None = None, tactics: Iterable[int] | None = None)

构建插件 I/O 支持的类型/格式组合集。

还可以声明每个此类类型/格式组合的任何自定义策略。策略只是另一种计算同一 I/O 类型/格式组合的插件输出的方法（例如，如果有多个内核可用）。

参数:

  • io_types (str, optional) –

    类型组合的字符串表示形式。

    有效格式为“type0,type1,…,type#io”，其中“type”的形式为“TYPE0[sep]TYPE1[sep]…”。

    TYPE 是 trt.DataType 的有效字符串表示形式。这些包括 trt.float32 的“FP32”，trt.float16 的“FP16”。其他数据类型的字符串表示形式与其在 trt.DataType 枚举中的名称相同。

    [sep] 是有效分隔符，可以是“|”或“*”。给定的 io_types 中只能出现其中一个分隔符。

    (1). “|”表示依赖组合：一个 I/O 的类型对另一个 I/O 的类型的依赖。例如，“FP32|FP16,FP32|FP16”表示 IO 只能同时为 FP32 或同时为 FP16。

    (2). “*”表示独立组合。例如，“FP32*FP16,FP32|FP16,FP32|FP16”表示第一个输入独立地为 FP32 或 FP16，而与其余 IO 无关。

  • layouts (str, optional) –

    格式组合的字符串表示形式。

    有效格式为“format0,format1,…,format#io”，其中“format”的形式为“FORMAT0[sep]FORMAT1[sep]…”。

    FORMAT 是 trt.TensorFormat 的有效字符串表示形式。这些是 trt.TensorFormat 的枚举值的字符串版本。例如，trt.TensorFormat.LINEAR 的“LINEAR”。

    [sep] 是有效分隔符，可以是“|”或“*”。规则与 io_types 的规则相同。

  • tactics (Iterable[int], optional) – 此类型/格式组合的自定义策略。每个自定义策略必须是正整数。默认为默认策略 (0)。

对于具有 3 个 I/O 的插件，I/O 索引 0 和 1 依赖性地为 FP32/FP16，索引 2 独立性地为 FP32/FP16。
1@trtp.autotune("my::plugin")
2def autotune(inp0: trtp.TensorDesc, inp1: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
3    # The following would result in the following type combinations:
4    # [FP32, FP32, FP32], [FP16, FP16, FP32], [FP32, FP32, FP16], [FP16, FP16, FP16]
5    return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16, FP32|FP16", "LINEAR", [1, 2])]
对于具有 2 个 I/O 的插件，输入/输出支持 FP32 的 LINEAR 或 HWC 格式以及 FP16 的 LINEAR 格式。
1@trtp.autotune("my::plugin")
2def autotune(inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
3    # Even though (FP16, HWC) is not a valid combination (see next example), TRT should intelligently reject those
4    # and pass the following combinations to the impl function:
5    # [{FP32, FP32}, {LINEAR, LINEAR}], [{FP32, FP32}, {HWC, LINEAR}], [{FP16, FP32}, {LINEAR, LINEAR}]
6    return [trtp.AutoTuneCombination("FP32*FP16, FP32", "LINEAR*HWC, LINEAR", [1, 2])]
对于具有 2 个 I/O 的插件，输入/输出支持 FP32 的 LINEAR 或 HWC 格式以及 FP16 的 LINEAR 格式（第二种方法）。
1@trtp.autotune("my::plugin")
2def autotune(inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
3    # We can use two AutoTuneCombination objects to avoid communicating illegal combinations
4    return [trtp.AutoTuneCombination("FP32*FP16, FP32", "LINEAR, LINEAR", [1, 2]), trtp.AutoTuneCombination("FP32, FP32", "HWC, LINEAR", [1, 2])]
pos(pos: Iterable[int], io_types: str, layouts: str = 'LINEAR') → None

为指定的 I/O 索引集指定 I/O 类型和格式。

参数:

  • pos (Iterable[int]) – I/O 索引。输入索引为 [0, 1, …, #inputs - 1]，输出索引为 [#inputs, #inputs + 1, …, #inputs + #outputs - 1]。

  • io_types (str) – 这些 I/O 索引的数据类型。

  • layouts (str, optional) – 这些 I/O 索引的张量格式。默认为“LINEAR”。

Raises:

ValueError – 如果已指定任何这些 I/O 索引的类型或布局。

对于具有 3 个 I/O 的插件，I/O 索引 0 和 1 依赖性地为 FP32/FP16，索引 2 独立性地为 FP32/FP16。
1@trtp.autotune("my::plugin")
2def autotune(inp0: trtp.TensorDesc, inp1: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
3    c = trtp.AutoTuneCombination()
4    c.pos([0, 1], "FP32|FP16", "LINEAR")
5    c.pos(2, "FP32*FP16") # Omitting format is the same as declaring it to be LINEAR.
6    c.tactics([1, 2])
7    return [c]
tactics(tactics: Iterable[int]) → None

为此类型/格式组合指定自定义策略

参数:

tactics (Iterable[int]) – 自定义策略。这些必须是正整数。