使用自定義轉換器過載 Torch-TensorRT 轉換器

如果由於某種原因您想要更改特定 PyTorch 操作到 TensorRT 的轉換行為，可以透過編寫自定義轉換器並重載 Torch-TensorRT 的來實現。這可能是因為您想使用自定義核而不是 TensorRT 的核，或者因為您想在 TensorRT 中使用與 Torch-TensorRT 通常使用的實現不同的層實現。

在本教程中，我們將演示如何使用一個自定義轉換器來過載 Torch-TensorRT 將 torch.nn.functional.gelu 操作轉換為 TensorRT 的過程，該自定義轉換器使用了 GeLU 層的不同實現。

import logging
import sys

import torch
import torch_tensorrt

GeLU 在 PyTorch 中有兩種模式，一種使用 erf 函式，另一種使用 tanh 近似。TensorRT 本身支援這兩種實現作為啟用層，但假設我們只想在 tanh 模式下在 TensorRT 中使用 GeLU 的自定義實現。

class GeLU(torch.nn.Module):
    def __init__(self, mode="tanh"):
        super().__init__()
        self.mode = mode

    def forward(self, x):
        return torch.nn.functional.gelu(x, approximate=self.mode)


my_mod = GeLU(mode="tanh")
ex_input = torch.randn(2, 5).to("cuda")

作為基線，我們可以使用標準的 Torch-TensorRT GeLU 轉換器（在 tanh 近似模式下）來處理我們的模組。

my_standard_gelu = torch_tensorrt.compile(
    my_mod, arg_inputs=(ex_input,), min_block_size=1
)
print(my_standard_gelu.graph)
print(my_standard_gelu(ex_input))

編寫自定義轉換器

轉換器是函式，它們獲取 PyTorch 圖中特定 PyTorch 操作例項，並將其轉換為正在構建的 TensorRT 圖中的等效 TensorRT 操作集。它們使用 @torch_tensorrt.dynamo.conversion.dynamo_tensorrt_converter 裝飾器註冊到 Torch-TensorRT。在程式碼層面，轉換器接收當前的轉換狀態（ConversionCtx）、圖中要轉換的下一個運算子以及該節點的引數，並返回該操作的佔位符輸出，同時副作用是將必要的 TensorRT 層插入到 TensorRT 網路中。

from typing import Dict, Sequence, Tuple, Union

import tensorrt as trt
from torch.fx.node import Argument, Node, Target
from torch_tensorrt.dynamo import CompilationSettings
from torch_tensorrt.dynamo.conversion import ConversionContext

轉換器元資料

@torch_tensorrt.dynamo.conversion.dynamo_tensorrt_converter(
    # The PyTorch operation to convert, when this operation is encountered, this converter will be called
    torch.ops.aten.gelu.default,
    # Validators are functions that determine that given a specific node, if it can be converted by the converter
    capability_validator=lambda node, settings: (
        "approximate" in node.kwargs and node.kwargs["approximate"] == "tanh"
    ),
    # Can this converter be used in cases where the input shapes are dynamic
    supports_dynamic_shapes=True,
    # Set the priority of the converter to supersede the default one
    priority=torch_tensorrt.dynamo.conversion.ConverterPriority.HIGH,
    # Whether the converter requires a dynamic output allocator to run (e.g. data dependent ops)
    requires_output_allocator=True,
)

對於定義轉換器的裝飾器，有一個必需引數和幾個可選引數。所有轉換器都需要一個目標運算子，其理念是當圖中存在 torch.ops.aten.gelu.default 的例項時，將呼叫此轉換器。

在目標運算子之後，您可以提供額外的元資料，這些元資料定義了轉換器的能力以及該轉換器相對於針對該目標的其他可能轉換器的優先順序

定義轉換器能力的主要工具是 capability_validator 引數，它是一個 lambda 函式，接受圖中的特定節點以及使用者編譯設定，並返回一個布林值，指示該轉換器是否可用於該節點。這個驗證器函式在圖分割槽階段之前針對轉換器目標操作的每個例項執行。在此階段沒有透過驗證器的轉換器的節點將在執行時在 PyTorch 中執行。這在您只想在特定情況下使用自定義轉換器時非常有用，就像在我們只希望在 approximate == "tanh" 時使用我們的轉換器的情況一樣。

與驗證器不同的是 supports_dynamic_shapes 引數，它是一個布林值，指示在輸入形狀為動態的情況下是否可以使用該轉換器。如果此引數設定為 False，在使用者提供的輸入為動態的情況下，此轉換器將被停用。如果沒有支援動態形狀的替代方案，該操作將在 PyTorch 中執行。

最後是 priority 引數，它是一個列舉，來自 torch_tensorrt.dynamo.conversion.ConverterPriority 類，定義了轉換器的優先順序。兩個選項是 HIGH 和 STANDARD。註冊為 STANDARD 的轉換器將被新增到給定操作的轉換器列表末尾，而註冊為 HIGH 的轉換器將被新增到列表開頭。候選轉換器將按此優先順序順序評估其適用性，並使用第一個透過驗證器的轉換器。

轉換器實現

轉換器函式本身接受以下引數：當前轉換上下文、目標運算子、目標運算子的引數、目標運算子的關鍵字引數以及目標運算子的名稱。引數可以是任何 Python 原語、torch.Tensor、np.Arrays 或 ITensor 物件。轉換器函式應主要以 TensorRT ITensor 的形式返回目標運算子的輸出。這些輸入和輸出應與目標 PyTorch 運算子的 schema 對應，其 schema 可在此處找到 https://pytorch.com.tw/docs/main/torch.compiler_ir.html。

由於 Torch-TensorRT 涵蓋了核心 ATen 操作集，它已經將許多常見的低階操作抽象為輔助函式，可用於構建 TensorRT 網路。這使開發人員能夠避免直接建立 TensorRT 層的樣板程式碼，轉而專注於轉換的高階邏輯。輔助函式位於 torch_tensorrt.dynamo.conversion.impl 模組中，設計上可與原始 TensorRT 實現組合和互操作。在本例中，我們將使用 impl 中的 Torch-TensorRT mul、add 和 tanh 函式來實現我們的替代 GeLU 層。

def aten_ops_gelu(
    ctx: ConversionContext,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
    # The schema for torch.ops.aten.gelu.default is gelu(Tensor self, *, str approximate=’none’) -> Tensor

    from torch_tensorrt.dynamo import SourceIR
    from torch_tensorrt.dynamo.conversion import impl

    # Cheap way to allow layer names to be unqiue
    op_count = 0

    def get_op_count():
        nonlocal op_count
        op_count += 1
        return op_count

    mul = lambda x, y: impl.elementwise.mul(
        ctx,
        target,
        name=f"mul_{get_op_count()}",
        source_ir=SourceIR.ATEN,
        lhs_val=x,
        rhs_val=y,
    )
    add = lambda x, y: impl.elementwise.add(
        ctx,
        target,
        name=f"add_{get_op_count()}",
        source_ir=SourceIR.ATEN,
        lhs_val=x,
        rhs_val=y,
    )
    tanh = lambda x: impl.activation.tanh(
        ctx, target, name=f"tanh_{get_op_count()}", source_ir=SourceIR.ATEN, input_val=x
    )

    # So we know that our custom converter is being run instead of the standard one
    print("\n\n---------------------------")
    print("Using custom GeLU converter")
    print("---------------------------\n\n")

    x_7 = mul(args[0], 0.5)
    x_8 = mul(args[0], 0.79788456080000003)
    x_9 = mul(args[0], 0.044714999999999998)
    x_10 = mul(x_9, args[0])
    x_11 = add(x_10, 1.0)
    x_12 = mul(x_8, x_11)
    x_13 = tanh(x_12)
    x_14 = add(x_13, 1.0)
    x_15 = mul(x_7, x_14)

    return x_15

使用我們的自定義轉換器

現在我們可以重新編譯並看到我們的自定義轉換器被呼叫以將 GeLU 轉換為 TensorRT。

my_custom_gelu = torch_tensorrt.compile(
    my_mod, arg_inputs=(ex_input,), min_block_size=1
)

print(my_custom_gelu.graph)
print(my_custom_gelu(ex_input))

我們可以驗證我們的實現與 tanh 近似的 TensorRT 實現相匹配。

print(
    f"tanh approximations are close: {torch.allclose(my_standard_gelu(ex_input), my_custom_gelu(ex_input))}"
)

最後，我們想驗證，如果 approximate 引數未設定為 tanh，則不使用我們的自定義轉換器。

my_mod_erf = GeLU(mode="none")
my_gelu_erf = torch_tensorrt.compile(
    my_mod_erf, arg_inputs=(ex_input,), min_block_size=1
)

注意，我們沒有看到來自自定義轉換器的列印語句，這表明它沒有被使用。然而，檢視圖，我們仍然可以看到建立了一個 TensorRT 引擎來執行 GeLU 操作。在這種情況下，我們的自定義轉換器的驗證器返回 False，因此轉換系統移動到列表中的下一個轉換器，即標準 GeLU 轉換器，並使用它來轉換操作。

print(my_gelu_erf.graph)
print(my_gelu_erf(ex_input))