編寫 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：允許開發人員覆寫轉換器登錄中的現有轉換器
轉換器只需要金鑰。
函式主體負責取得網路的目前狀態，並新增下一個子圖，以使用 TensorRT 運算執行修飾器中指定的運算。該函式會以原生 PyTorch 運算的方式提供引數，並針對凍結張量屬性的 numpy 陣列或 TensorRT ITensor 新增案例，這些 TensorRT ITensor 是先前節點的輸出張量，對應於圖中中間運算的邊緣/輸出張量。若要判斷預期的類型以及轉換器的傳回類型，請查看要轉換的運算的定義。在 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中的目標金鑰。例如：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 ITensor 的聯集。如果轉換器中僅支援類型子集，您也可以新增 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 運算，這些運算將呼叫其對應的轉換器。
@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 lowering 運算。
在大多數情況下，這比實作轉換器容易得多。因此，在可能的情況下，這應該是首先嘗試的方法。
