使用 Torch-TensorRT 編譯匯出的程式¶
Pytorch 2.1 引入了 torch.export 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 tensors 或 torch_tensorrt.Input 物件的列表。對於動態形狀，這應該是torch_tensorrt.Input物件的列表。enabled_precisions- TensorRT 建構器在最佳化期間可以使用的精度集合。truncate_long_and_double- 將 long 和 double 值分別截斷為 int 和 floats。torch_executed_ops- 強制由 Torch 執行的運算子。min_block_size- 作為 TensorRT 段執行所需的連續運算子的最小數量。
完整的選項列表可以在這裡找到。
注意
我們目前在 Dynamo 中不支援 INT 精度。目前在我們的 Torchscript IR 中支援此功能。我們計劃在下一個版本中為 dynamo 實施類似的支援。
底層原理¶
在底層，torch_tensorrt.dynamo.compile 在圖形上執行以下操作。
降低(Lowering) - 應用降低傳遞以新增/移除運算子以進行最佳轉換。
分割(Partitioning) - 根據
min_block_size和torch_executed_ops欄位將圖形分割成 Pytorch 和 TensorRT 段。轉換(Conversion) - 在此階段，Pytorch 運算會轉換為 TensorRT 運算。
最佳化(Optimization) - 轉換後，我們建構 TensorRT 引擎並將其嵌入到 pytorch 圖形中。
追蹤(Tracing)¶
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)
