快捷方式

Dynamo 转换器

Torch-TensorRT 中的 dynamo 转换器库位于 TensorRT/py/torch_tensorrt/dynamo/conversion 中。

步骤

操作集

dynamo 中的转换器由 aten_trace 生成,属于 aten_ops_converters(FX 早期有 acc_ops_convertersaten_ops_convertersnn_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 是类似的装饰器函数,但存在一些差异。

      1. 两者都在注册表(Python 字典)中注册转换器 - 分别是 CONVERTERSDYNAMO_CONVERTERS。这两个字典合并在一起,形成整体转换器注册表。

      2. 字典使用 OpOverLoad 作为键,我们将在下面使用示例详细说明。

      3. 两者都返回装饰的转换器实现。

      4. CONVERTERS 直接注册装饰的 converter_implementation 函数,而 DYNAMO_CONVERTERS 具有额外的参数并注册 ConverterSupport 对象。

      5. 附加参数是:

        1. key:实现转换器的节点目标(例如,torch.ops.aten.leaky_relu.Tensor)

        2. enabled:是否应启用/缓存转换器。默认为 None,表示 capability_validator 函数始终为 true。这意味着默认情况下,所有“key”类型的节点都可以由该转换器支持。有关更多详细信息,请参阅 embedding 示例。

        3. capability_validator:一个函数,用于评估节点是否可以通过装饰的转换器进行转换。它默认为 None,表示 capability_validator 函数始终为 true。这意味着默认情况下,所有“key”类型的节点都可以由该转换器支持。有关更多详细信息,请参阅 embedding 示例。

        4. priority:转换器相对于具有相同目标的其他转换器的优先级级别。

      6. ConverterSupportconverter_implementationcapability_validator 的组合。

      tensorrt_converterdynamo_tensorrt_converter 装饰的函数具有以下参数,这些参数由上面提到的跟踪函数自动生成。

      1. network:以 call_modulecall_function 形式存在的节点,其目标是键。

      2. target:上面 call_modulecall_function 中的目标键。例如:torch.ops.aten_.leaky_relu.default。请注意,torch.ops.aten._leaky_reluOpOverloadPacket,而 torch.ops.aten_.leaky_relu.defaultOpOverload

      3. args:传递给上面 call_modulecall_function 的参数。

      4. kwargs:传递给上面 call_modulecall_function 的关键字参数。

      5. name:包含目标名称的字符串。

      作为编写新转换器的用户,您只需确保从生成的跟踪中提取适当的参数,并传递到实现库函数 activation.leaky_relu 中的实现函数(我们将在下面详细讨论)。

    • 操作类型

      示例:fmod

      它遵循与上述转换器相同的步骤。在这种情况下,操作码为 torch.ops.aten.fmod.Scalartorch.ops.aten.fmod.Tensor。因此,这两个操作码都在 py/torch_tensorrt/dynamo/conversion/aten_ops_converters 中注册。请注意,torch.ops.aten.fmodOpOverLoadPacket,而注册表以 torch.ops.aten.fmod.Scalartorch.ops.aten.fmod.Tensor 为键,它是 OpOverLoad

      示例:embedding

      它遵循与上述转换器相同的步骤。在这种情况下,操作码为 torch.ops.aten.embedding.default。有一些转换器需要考虑特殊情况。在这些情况下,应使用 capability_validators 使用 @dynamo_tensorrt_converter 注册转换器。我们将通过 torch.ops.aten.embedding.default 说明这一点。它有参数 - scale_grad_by_freqsparse,目前不受实现支持。在这种情况下,我们可以编写验证器 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
      

      实现函数具有以下参数。

      1. network : 来自装饰函数注册的 network

      2. target: 来自装饰函数注册的 target

      3. source_ir: 枚举属性。 SourceIR 枚举定义在 py/torch_tensorrt/dynamo/conversion/impl/converter_utils

      4. name: 来自装饰函数注册的 name

      5. input_val: 从装饰函数注册的 args 或 kwargs 中提取的适当参数

      6. alpha: 从装饰函数注册的 args 或 kwargs 中提取的适当参数。如果非空,它将设置创建的 TensorRT 激活层的 alpha 属性,例如:用于 leaky_relu、elu、hardtanh

      7. beta: 从装饰函数注册的 args 或 kwargs 中提取的适当参数。如果非空,它将设置创建的 TensorRT 激活层的 beta 属性,例如:用于 hardtanh

      8. 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 中。这是所有逐元素函数定义和实现的地方。对于新的运算符,应确定它所属的类别。以下是一些示例

      1. fmod 这样的逐元素运算符存在于 py/torch_tensorrt/dynamo/conversion/impl/elementwise 中。 py/torch_tensorrt/dynamo/conversion/impl/elementwise/base 包含逐元素运算符的基本函数。

      2. sqrt 这样的单目运算符将存在于 py/torch_tensorrt/dynamo/conversion/impl/unary 中。 py/torch_tensorrt/dynamo/conversion/impl/unary/base 包含单目运算符的基本函数。

      3. softmaxlayer_normbatch_norm 这样的归一化运算符将存在于 py/torch_tensorrt/dynamo/conversion/impl/normalization 中。由于没有所有运算符共有的基本运算,因此没有基本文件。但是,如果所有归一化运算符都有公共函数,则可以选择实现基本文件

      4. sliceselectwhereembedding 这样的单个运算符将存在于 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 测试

    1. 比较 fx.symbolic_tracetorch_tensorrt.dynamo.compile 的结果。

    2. 测试 expected_opunexpected_op

      1. expected_op:将操作降低到的操作。例如:对于 addmm,为 muladd

      2. 请注意,指定 disable_passes= True 用于您不希望降低传递的情况(在测试转换器时应该是默认值)

      3. unexpected_op:原始操作。例如:对于 addmm,为 addmm

如果上述两个条件中的任何一个失败,测试应该失败

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源