編寫 Dynamo Converter¶
Torch-TensorRT 中的 dynamo converter 庫位於 TensorRT/py/torch_tensorrt/dynamo/conversion。
Converter 實現¶
註冊¶
Converter 是一個使用 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: 實現 converter 的節點目標（例如，torch.ops.aten.leaky_relu.default）
enabled: 此 converter 是否應被啟用，以便在 converter 登錄檔中可用
capability_validator: 一個 lambda 函式，它接受一個torch.fx.Node並判斷此 converter 是否能正確處理此節點。如果驗證器返回False，子圖分割槽器將確保此節點在編譯圖中於 PyTorch 中執行。
priority: 允許開發者覆蓋 converter 登錄檔中現有的 converter
Converter 所需的全部內容只有 key。
函式體負責獲取網路的當前狀態，並新增下一個子圖以使用 TensorRT 操作執行裝飾器中指定的操作。函式提供的引數與原生 PyTorch 操作提供的引數相同，但增加了用於凍結張量屬性的 numpy 陣列或 TensorRT ITensor（它們是前一個節點的輸出張量，對應於圖中的中間操作的邊/輸出張量）。要確定預期的型別以及 converter 的返回型別，請檢視要轉換的操作的定義。對於 aten 操作，此檔案是事實來源：https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml 由於開發者可能編寫的許多 converter 是較低級別操作的組合，因此無需在原始 TensorRT 中實現 converter，torch_tensorrt.dynamo.conversion.impl 子包包含許多可以鏈式組合以建立 TensorRT 子圖的操作實現。
ctx: 編譯器的當前狀態。Converter 主要操作 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 用於構建一個特定的 TensorRT 子圖，表示 INetworkDefinition 中的當前節點。
kwargs: 傳遞給特定節點的引數（由torch_tensorrt.dynamo.conversion.TRTInterpreter收集）。
name: 包含目標名稱的字串
函式預期返回 tensorrt.ITensor 或一些 tensorrt.ITensor 的集合，供 torch_tensorrt.dynamo.conversion.TRTInterpreter 使用，匹配被轉換操作的輸出簽名
能力驗證¶
有些 converter 需要考慮特殊情況。在這種情況下，應該使用 capability_validators 透過 @dynamo_tensorrt_converter 註冊 converter。我們透過 torch.ops.aten.embedding.default 來闡述這一點。它有兩個引數 - scale_grad_by_freq 和 sparse，這些引數目前不支援。在這種情況下，我們可以編寫 validator embedding_param_validator，它實現的是，如果給定這些引數，則不支援此 converter，並透過以下方式註冊 converter
型別契約¶
函式應遵循簽名建立的型別契約。這包括接受有效的 PyTorch 型別與常量張量的 numpy 陣列以及 TensorRT ITensor 的並集。如果 converter 中僅支援部分型別，您還可以新增 torch_tensorrt.dynamo.conversion.converter_utils.enforce_tensor_types，它允許您指定一個字典，對映輸入位置與其可以接受的型別。在可能的情況下，裝飾器將轉換輸入以匹配這些型別，優先考慮提供的順序。字典中的 int key 將引用 args 中的位置引數。 str key 將引用 kwargs 中的關鍵字引數。
示例: Convolution¶
預設的卷積 converter 同時使用能力驗證器和型別強制來防止在不支援的情況下執行。能力驗證器在分割槽期間執行，以確定特定的卷積節點是否可以轉換為 TensorRT 或需要在 PyTorch 中執行。在此，驗證器確保卷積不超過 3D。型別強制器會在呼叫 converter 之前進行自動型別轉換，將輸入轉換為 converter 中支援的型別，從而限制作者必須處理的情況數量。
@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。
運算子分解¶
有些 converter 可以分解為 PyTorch 中的子操作，無需單獨註冊 converter。此類 converter 可以透過分解實現
示例: addmm¶
分解透過 register_torch_trt_decomposition 裝飾器註冊。我們定義 addmm_replacement 並將其替換為 torch ops，這將呼叫其相應的 converter。
@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 下層操作。
在大多數情況下，這比實現 converter 要容易得多。因此，在可能的情況下，應首先嚐試這種方法。
