编写 Dynamo 转换器¶
Torch-TensorRT 中的 Dynamo 转换器库位于 TensorRT/py/torch_tensorrt/dynamo/conversion。
转换器实现¶
注册¶
转换器是一个用 torch_tensorrt.dynamo.dynamo_tensorrt_converter 装饰的函数，该函数遵循函数签名
@torch_tensorrt.dynamo.conversion.dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
def leaky_relu_converter(
ctx: torch_tensorrt.dynamo.conversion.ConversionCtx,
target: Target,
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[tensorrt.ITensor, Sequence[tensorrt.ITensor]]:
装饰器接受多个参数
key: 实现该转换器的节点目标（例如，torch.ops.aten.leaky_relu.default）
enabled: 该转换器是否应被启用为转换器注册表中可用的转换器
capability_validator: 一个 lambda 函数，它可以接受一个torch.fx.Node并确定转换器是否能够正确处理该节点。如果验证器返回False，子图划分器将确保该节点在编译图的 PyTorch 中运行。
priority: 允许开发人员覆盖转换器注册表中现有的转换器
转换器所需的一切就是 key。
函数体负责获取网络的当前状态，并添加下一个子图，使用 TensorRT 操作来执行装饰器中指定的 op。函数将接收到与原生 PyTorch op 相同的参数，外加用于冻结 Tensor 属性的 NumPy 数组或 TensorRT ITensors（它们是前一个节点的输出 Tensor，对应于图中中间操作的边/输出 Tensor）。要确定预期类型以及转换器的返回类型，请查看要转换的 op 的定义。对于 aten 操作，此文件将是事实来源：https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml 由于开发人员可能编写的许多转换器是低级运算符的组合，因此不需要用原始 TensorRT 实现转换器，torch_tensorrt.dynamo.conversion.impl 子包包含了许多可以链接起来创建 TensorRT 子图的操作的实现。
ctx: 编译器的当前状态。转换器主要将操作 ctx.net，它是正在构建的tensorrt.INetworkDefinition。其他元数据，包括用户提供的设置，也可以在此结构中获得。
target:call_module或call_function上方的目标 key。例如：torch.ops.aten_.leaky_relu.default。请注意，torch.ops.aten._leaky_relu是OpOverloadPacket，而torch.ops.aten_.leaky_relu.default是OpOverload。
args: 传递给特定节点的参数（由torch_tensorrt.dynamo.conversion.TRTInterpreter收集）。这些参数与 kwargs 一起用于构建代表 INetworkDefinition 中当前节点的特定 TensorRT 子图。
kwargs: 传递给特定节点的参数（由torch_tensorrt.dynamo.conversion.TRTInterpreter收集）。
name: 包含目标名称的字符串
该函数预计将返回 tensorrt.ITensor 或一些 tensorrt.ITensor 的集合，供 torch_tensorrt.dynamo.conversion.TRTInterpreter 使用，与要转换的操作的输出签名相匹配
能力验证¶
有些转换器有需要考虑的特殊情况。在这种情况下，应该使用 capability_validators 来使用 @dynamo_tensorrt_converter 注册转换器。我们通过 torch.ops.aten.embedding.default 来说明这一点。它有参数 - scale_grad_by_freq 和 sparse，这些参数目前在实现中不支持。在这种情况下，我们可以编写验证器 embedding_param_validator，该验证器实现给定这些参数时转换器不支持，并通过以下方式注册转换器：
类型契约¶
该函数预计遵循签名建立的类型契约。这包括接受有效的 PyTorch 类型 + 用于常量张量的 NumPy 数组以及 TensorRT ITensors 的联合。如果转换器只支持部分类型，您还可以添加 torch_tensorrt.dynamo.conversion.converter_utils.enforce_tensor_types，它允许您指定一个字典，将输入位置与这些输入可以接受的类型进行映射。在可能的情况下，装饰器将尝试将输入转换为匹配这些类型，优先使用提供的顺序。int 键将引用 args 中的位置参数。str 键将引用 kwargs 中的关键字参数。
示例：Convolution¶
默认的卷积转换器同时使用能力验证器和类型强制，以防止在不支持的情况下运行。能力验证器在划分期间运行，以确定特定的卷积节点是否可以转换为 TensorRT 或需要在 PyTorch 中运行。在这里，验证器确保卷积的维度不超过 3D。类型强制器将在调用转换器之前自动转换输入，将输入转换为转换器支持的类型，从而限制作者必须处理的用例数量。
@dynamo_tensorrt_converter(
torch.ops.aten.convolution.default, capability_validator=lambda conv_node: conv_node.args[7] in ([0], [0, 0], [0, 0, 0])
) # type: ignore[misc]
@enforce_tensor_types(
{
0: (TRTTensor,),
1: (np.ndarray, torch.Tensor, TRTTensor),
2: (np.ndarray, torch.Tensor, TRTTensor),
}
) # type: ignore[misc]
def aten_ops_convolution(
ctx: ConversionContext,
target: Target,
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
评估器¶
某些操作不会产生 TensorRT 子图作为副作用。这些被称为评估器。
示例：
operator.getitem评估器之所以被归类为这样，是因为它们不对图进行任何修改。这在
py/torch_tensorrt/dynamo/conversion/op_evaluators.py中实现，并具有相应的capbility_validator。操作码是operator.getitem。
运算符分解¶
有些转换器可以分解为 PyTorch 中的子操作，无需单独注册转换器。此类转换器可以通过分解来实现
示例：addmm¶
分解通过 register_torch_trt_decomposition 装饰器进行注册。我们定义 addmm_replacement 并用 torch ops 替换它，然后将调用相应的转换器。
@torch_tensorrt.dynamo.lowering.register_torch_trt_decomposition(torch.ops.aten.addmm)
def addmm_replacement(
input_: torch.Tensor, mat1: torch.Tensor, mat2: torch.Tensor, *, beta=1, alpha=1
) -> torch.Tensor:
return torch.add(
torch.mul(input_, beta), torch.mul(torch.matmul(mat1, mat2), alpha)
)
您可以通过编辑 torch_tensorrt.dynamo.lowering.torch_enabled_decompositions 和 torch_tensorrt.dynamo.lowering.torch_disabled_decompositions 来修改运行的分解。
注意：
torch_tensorrt.dynamo.lowering.torch_enabled_decompositions和torch_tensorrt.dynamo.lowering.torch_disabled_decompositions必须是互斥集，并且torch_tensorrt.dynamo.lowering中已定义的分解将优先于 torch 降低 ops。
很多时候，这比实现转换器要容易得多。因此，在可能的情况下，这应该是首先尝试的方法。
