Dynamo 转换器¶
Torch-TensorRT 中的 dynamo 转换器库位于 TensorRT/py/torch_tensorrt/dynamo/conversion
中。
步骤¶
操作集¶
dynamo 中的转换器由 aten_trace
生成,属于 aten_ops_converters
(FX 早期有 acc_ops_converters
、aten_ops_converters
或 nn_ops_converters
,具体取决于生成它的跟踪)。这些转换器使用 dynamo_tensorrt_converter
注册到 dynamo。装饰的函数具有参数 - network, target, args, kwargs, name
,这是所有运算符模式共有的。这些函数映射在 aten
转换器注册字典中(目前是 FX 和 dynamo 转换器的集合,FX 很快就会被弃用),键是函数目标名称。
aten_trace 由
torch_tensorrt.dynamo.trace(..)
针对导出路径生成,而torch_tensorrt.compile(ir=dynamo)
针对编译路径生成。导出路径使用
aten_tracer
,而编译中的备用跟踪由 AOT Autograd 库生成。这两个跟踪都将 Torch 运算符简化为一组简化的 Aten 运算符。
如上所述,如果您想添加新的转换器,其实现将包含在 TensorRT/py/torch_tensorrt/dynamo/conversion/impl
中。尽管在 TensorRT/py/torch_tensorrt/fx/impl
中的通用实现库中包含了相应的转换器实现,用于 FX 转换器,但本文档重点介绍了 dynamo 中 aten_ops
转换器的实现。
转换器实现¶
在本节中,我们将说明编写转换器需要实现的步骤。我们根据激活、运算符、降级传递实现或评估器对它们进行划分。每个步骤都将通过示例进行详细说明。
注册
转换器需要使用适当的 op 代码注册到
dynamo_tensorrt_converter
。
激活类型
示例:
leaky_relu
aten_ops_converters:Dynamo_converters
在
py/torch_tensorrt/dynamo/conversion/aten_ops_converters
中定义。您需要使用dynamo_tensorrt_converter
装饰器注册在跟踪中生成的 op 代码。在本例中,用于注册或转换器注册键的 op 代码是torch.ops.aten.leaky_relu.default
@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default) def aten_ops_leaky_relu( network: TRTNetwork, target: Target, args: Tuple[Argument, ...], kwargs: Dict[str, Argument], name: str, ) -> Union[TRTTensor, Sequence[TRTTensor]]: return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
tensorrt_converter
(用于 FX 注册)和dynamo_tensorrt_converter
是类似的装饰器函数,但存在一些差异。
两者都在注册表(Python 字典)中注册转换器 - 分别是
CONVERTERS
和DYNAMO_CONVERTERS
。这两个字典合并在一起,形成整体转换器注册表。字典使用
OpOverLoad
作为键,我们将在下面使用示例详细说明。两者都返回装饰的转换器实现。
CONVERTERS
直接注册装饰的converter_implementation
函数,而DYNAMO_CONVERTERS
具有额外的参数并注册ConverterSupport
对象。附加参数是:
key:实现转换器的节点目标(例如,torch.ops.aten.leaky_relu.Tensor)
enabled:是否应启用/缓存转换器。默认为 None,表示 capability_validator 函数始终为 true。这意味着默认情况下,所有“key”类型的节点都可以由该转换器支持。有关更多详细信息,请参阅
embedding
示例。capability_validator:一个函数,用于评估节点是否可以通过装饰的转换器进行转换。它默认为 None,表示 capability_validator 函数始终为 true。这意味着默认情况下,所有“key”类型的节点都可以由该转换器支持。有关更多详细信息,请参阅
embedding
示例。priority:转换器相对于具有相同目标的其他转换器的优先级级别。
ConverterSupport
是converter_implementation
和capability_validator
的组合。由
tensorrt_converter
和dynamo_tensorrt_converter
装饰的函数具有以下参数,这些参数由上面提到的跟踪函数自动生成。
network:以
call_module
或call_function
形式存在的节点,其目标是键。target:上面
call_module
或call_function
中的目标键。例如:torch.ops.aten_.leaky_relu.default
。请注意,torch.ops.aten._leaky_relu
是OpOverloadPacket
,而torch.ops.aten_.leaky_relu.default
是OpOverload
。args:传递给上面
call_module
或call_function
的参数。kwargs:传递给上面
call_module
或call_function
的关键字参数。name:包含目标名称的字符串。
作为编写新转换器的用户,您只需确保从生成的跟踪中提取适当的参数,并传递到实现库函数
activation.leaky_relu
中的实现函数(我们将在下面详细讨论)。操作类型
示例:
fmod
它遵循与上述转换器相同的步骤。在这种情况下,操作码为
torch.ops.aten.fmod.Scalar
或torch.ops.aten.fmod.Tensor
。因此,这两个操作码都在py/torch_tensorrt/dynamo/conversion/aten_ops_converters
中注册。请注意,torch.ops.aten.fmod
是OpOverLoadPacket
,而注册表以torch.ops.aten.fmod.Scalar
或torch.ops.aten.fmod.Tensor
为键,它是OpOverLoad
示例:
embedding
它遵循与上述转换器相同的步骤。在这种情况下,操作码为
torch.ops.aten.embedding.default
。有一些转换器需要考虑特殊情况。在这些情况下,应使用capability_validators
使用@dynamo_tensorrt_converter
注册转换器。我们将通过torch.ops.aten.embedding.default
说明这一点。它有参数 -scale_grad_by_freq
和sparse
,目前不受实现支持。在这种情况下,我们可以编写验证器embedding_param_validator
,它实现的是给定这些参数,转换器不受支持,并通过以下方式注册转换器:因此,如果存在新的转换器,其中某些特殊情况不受支持,则可以在
capability_validator
中指定它们。评估器类型
示例:
operator.getitem
评估器被归类为这样,因为它们不会对图进行任何修改。这在
py/torch_tensorrt/dynamo/conversion/op_evaluators.py
中实现,并使用相应的capbility_validator
。操作码为operator.getitem
。实现库
Dynamo 转换器将位于
py/torch_tensorrt/dynamo/conversion/impl
中
激活
示例:
leaky_relu
实现将放在
py/torch_tensorrt/dynamo/conversion/impl/activation.py
中。这是所有激活函数定义和实现的地方。def leaky_relu( network: TRTNetwork, target: Target, source_ir: Optional[SourceIR], name: str, input_val: TRTTensor, alpha: Optional[Any], ): #implementation实现函数具有以下参数。
network : 来自装饰函数注册的
network
target: 来自装饰函数注册的
target
source_ir: 枚举属性。
SourceIR
枚举定义在py/torch_tensorrt/dynamo/conversion/impl/converter_utils
中name: 来自装饰函数注册的
name
input_val: 从装饰函数注册的 args 或 kwargs 中提取的适当参数
alpha: 从装饰函数注册的 args 或 kwargs 中提取的适当参数。如果非空,它将设置创建的 TensorRT 激活层的 alpha 属性,例如:用于 leaky_relu、elu、hardtanh
beta: 从装饰函数注册的 args 或 kwargs 中提取的适当参数。如果非空,它将设置创建的 TensorRT 激活层的 beta 属性,例如:用于 hardtanh
dyn_range_fn: 一个可选函数,它接受 TensorRT 张量的动态范围,并返回输出动态范围
实现函数调用
py/torch_tensorrt/dynamo/conversion/impl/activation.py
中的convert_activation
函数。此函数将通过network.add_activation
添加适当的激活层。运算符
对于 Dynamo,实现将放在
py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py
中。这是所有逐元素函数定义和实现的地方。对于新的运算符,应确定它所属的类别。以下是一些示例
像
fmod
这样的逐元素运算符存在于py/torch_tensorrt/dynamo/conversion/impl/elementwise
中。py/torch_tensorrt/dynamo/conversion/impl/elementwise/base
包含逐元素运算符的基本函数。像
sqrt
这样的单目运算符将存在于py/torch_tensorrt/dynamo/conversion/impl/unary
中。py/torch_tensorrt/dynamo/conversion/impl/unary/base
包含单目运算符的基本函数。像
softmax
、layer_norm
、batch_norm
这样的归一化运算符将存在于py/torch_tensorrt/dynamo/conversion/impl/normalization
中。由于没有所有运算符共有的基本运算,因此没有基本文件。但是,如果所有归一化运算符都有公共函数,则可以选择实现基本文件像
slice
、select
、where
、embedding
这样的单个运算符将存在于py/torch_tensorrt/dynamo/conversion/impl/*.py
中。它们将具有与上述相同的 API 结构的单个运算符实现,但具有不同的单个参数请注意,上述运算符将具有要实现的公共函数,这些函数应放在
py/torch_tensorrt/dynamo/conversion/impl/converter_utils.py
中降低类型
有一些转换器可以分解为子运算,并且不需要单独的转换器注册。此类转换器可以通过
lowering passes
实现示例:
addmm
分解通过
py/torch_tensorrt/dynamo/backend/lowering/_decompositions.py
中的register_decomposition
注册。我们定义addmm_replacement
并用 torch 操作替换它,这些操作将调用其对应的转换器。@register_decomposition(torch.ops.aten.addmm, registry=DECOMPOSITIONS) 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 目录中存在一些预先存在的 Dynamo 分解,在这种情况下,应使用它们。在这种情况下,请在
py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py
中的torch_enabled_decompositions
中启用分解。类似地,您可以选择在torch_disabled_decompositions
中禁用任何分解。请注意,已经在降低中定义的那些将优先于 torch 降低操作。
测试¶
Dynamo 测试
Dynamo 测试存在于
tests/py/dynamo/lowering/test_decompositions.py
中,用于降低操作。上述转换器将很快移植到 Dynamo 测试比较
fx.symbolic_trace
和torch_tensorrt.dynamo.compile
的结果。测试
expected_op
和unexpected_op
。expected_op
:将操作降低到的操作。例如:对于addmm
,为mul
和add
请注意,指定
disable_passes= True
用于您不希望降低传递的情况(在测试转换器时应该是默认值)unexpected_op
:原始操作。例如:对于addmm
,为addmm
如果上述两个条件中的任何一个失败,测试应该失败