使用 Torch-TensorRT 編譯匯出程式

Pytorch 2.1 引入了 torch.export API，該 API 可以將 PyTorch 程式中的圖匯出為 ExportedProgram 物件。Torch-TensorRT dynamo 前端會編譯這些 ExportedProgram 物件，並使用 TensorRT 對其進行最佳化。以下是 dynamo 前端的一個簡單用法

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224), dtype=torch.float32).cuda()]
exp_program = torch.export.export(model, tuple(inputs))
trt_gm = torch_tensorrt.dynamo.compile(exp_program, inputs) # Output is a torch.fx.GraphModule
trt_gm(*inputs)

注意

torch_tensorrt.dynamo.compile 是使用者與 Torch-TensorRT dynamo 前端互動的主要 API。模型的輸入型別應為 ExportedProgram（理想情況下是 torch.export.export 或 torch_tensorrt.dynamo.trace（將在下一節討論）的輸出），輸出型別是 torch.fx.GraphModule 物件。

可配置設定

使用者可以透過許多選項自定義使用 TensorRT 進行最佳化的設定。以下是一些常用的選項：

• inputs - 對於靜態形狀，它可以是 torch 張量的列表或 torch_tensorrt.Input 物件。對於動態形狀，這應該是一個 torch_tensorrt.Input 物件的列表。

• enabled_precisions - TensorRT 構建器在最佳化期間可以使用的精度集合。

• truncate_long_and_double - 將 long 和 double 值分別截斷為 int 和 float。

• torch_executed_ops - 強制由 Torch 執行的運算子。

• min_block_size - 作為 TensorRT 段執行所需的連續運算子的最小數量。

完整的選項列表可以在這裡找到

注意

我們目前在 Dynamo 中不支援 INT 精度。目前對 INT 精度的支援存在於我們的 Torchscript IR 中。我們計劃在下一個版本中為 dynamo 實現類似的支援。

原理

從底層來看，torch_tensorrt.dynamo.compile 在圖上執行以下操作。

• 降級 - 應用降級流程以新增/移除運算子，從而實現最佳轉換。

• 分割槽 - 根據 min_block_size 和 torch_executed_ops 欄位將圖劃分為 PyTorch 和 TensorRT 段。

• 轉換 - 在此階段，PyTorch 運算子被轉換為 TensorRT 運算子。

• 最佳化 - 轉換後，我們構建 TensorRT 引擎並將其嵌入到 PyTorch 圖中。

追蹤

torch_tensorrt.dynamo.trace 可用於追蹤 PyTorch 圖並生成 ExportedProgram。這在內部執行一些運算子分解，以便進行下游最佳化。然後可以將 ExportedProgram 與 torch_tensorrt.dynamo.compile API 一起使用。如果您的模型具有動態輸入形狀，您可以使用此 torch_tensorrt.dynamo.trace 匯出具有動態形狀的模型。此外，您也可以直接使用帶約束的 torch.export。

import torch
import torch_tensorrt

inputs = [torch_tensorrt.Input(min_shape=(1, 3, 224, 224),
                              opt_shape=(4, 3, 224, 224),
                              max_shape=(8, 3, 224, 224),
                              dtype=torch.float32)]
model = MyModel().eval()
exp_program = torch_tensorrt.dynamo.trace(model, inputs)