Torch-TensorRT 解释¶
Torch-TensorRT 是一个 PyTorch 模型编译器,通过 TensorRT 模型优化 SDK 定位 NVIDIA GPU。它旨在为 PyTorch 模型提供更好的推理性能,同时保持 PyTorch 出色的易用性。
Dynamo 前端¶
Dynamo 前端是 Torch-TensorRT 的默认前端。它利用 PyTorch 的 Dynamo 编译器堆栈。
torch.compile
(即时编译)¶
torch.compile
是一个 JIT 编译器堆栈,因此,编译会推迟到首次使用时。这意味着当图中的条件发生变化时,图将自动重新编译。这为用户提供了最大的运行时灵活性,但限制了序列化选项。
在底层,torch.compile 将它认为可以降低到 Torch-TensorRT 的子图委派给它。Torch-TensorRT 进一步将这些图降低为仅包含 核心 ATen 运算符 或选定的“高级运算符”的运算,这些运算符适用于 TensorRT 加速。子图根据对运算符的支持进一步划分为在 PyTorch 中运行的组件和进一步编译为 TensorRT 的组件。然后 TensorRT 引擎替换受支持的块,并将混合子图返回给 torch.compile
以在调用时运行。
支持的格式¶
torch.fx GraphModule (
torch.fx.GraphModule
)PyTorch Module (
torch.nn.Module
)
返回¶
在首次调用时触发编译的封装函数
torch_tensorrt.dynamo.compile
(提前编译)¶
torch_tensorrt.dynamo.compile
是一个 AOT 编译器,模型在显式编译阶段进行编译。这些编译工件随后可以序列化并在以后重新加载。图通过 torch.export.trace
系统降低为包含 核心 ATen 运算符 或选定的适用于 TensorRT 加速的“高级运算符”的图。子图根据对运算符的支持进一步划分为在 PyTorch 中运行的组件和进一步编译为 TensorRT 的组件。然后 TensorRT 引擎替换受支持的块,并将混合子图打包到 ExportedProgram 中,该程序可以序列化和重新加载。
支持的格式¶
torch.export.ExportedProgram (
torch.export.ExportedProgram
)torch.fx GraphModule (
torch.fx.GraphModule
) (通过torch.export.export
)PyTorch Module (
torch.nn.Module
) (通过torch.export.export
)
返回¶
torch.fx.GraphModule (可使用
torch.export.ExportedProgram
序列化)
传统前端¶
由于 PyTorch 生态系统多年来存在多种编译器技术,Torch-TensorRT 具有一些针对它们的传统功能。
TorchScript (torch_tensorrt.ts.compile)¶
TorchScript 前端是 Torch-TensorRT 最初的默认前端,目标是 TorchScript 格式的模型。提供的图将被划分为受支持和不受支持的块。受支持的块将被降低到 TensorRT,不受支持的块将保留以与 LibTorch 运行。结果图将作为 ScriptModule
返回给用户,可以使用 Torch-TensorRT PyTorch 运行时扩展加载和保存。
支持的格式¶
TorchScript Module (
torch.jit.ScriptModule
)PyTorch Module (
torch.nn.Module
) (通过torch.jit.script
或torch.jit.trace
)
返回¶
TorchScript Module (
torch.jit.ScriptModule
)
FX 图模块 (torch_tensorrt.fx.compile)¶
此前端几乎已被 Dynamo 前端完全取代,Dynamo 前端是 FX 前端可用功能的超集。原始 FX 前端出于向后兼容性原因仍保留在代码库中。
支持的格式¶
torch.fx GraphModule (
torch.fx.GraphModule
)PyTorch Module (
torch.nn.Module
) (通过torch.fx.trace
)
返回¶
torch.fx GraphModule (
torch.fx.GraphModule
)
torch_tensorrt.compile
¶
由于存在许多不同的前端和支持的格式,我们提供了一个便利层 torch_tensorrt.compile
,它允许用户访问所有不同的编译器选项。您可以通过设置 ir
选项来指定 torch_tensorrt.compile
使用哪个编译器路径,告诉 Torch-TensorRT 尝试通过特定的中间表示来降低提供的模型。
ir
选项¶
torch_compile
:使用torch.compile
系统。立即返回一个封装函数,该函数将在首次调用时进行编译。dynamo
:通过torch.export
/torchdynamo 堆栈运行图。如果输入模块是torch.nn.Module
,则它必须是“可导出跟踪的”,因为该模块将通过torch.export.export
进行跟踪。返回一个torch.fx.GraphModule
,该模块可以立即运行或通过torch.export.export
或torch_tensorrt.save
保存。torchscript
或ts
:通过 TorchScript 堆栈运行图。如果输入模块是torch.nn.Module
,则它必须是“可脚本化的”,因为该模块将通过torch.jit.script
进行编译。返回一个torch.jit.ScriptModule
,该模块可以立即运行或通过torch.save
或torch_tensorrt.save
保存。fx
:通过torch.fx
堆栈运行图。如果输入模块是torch.nn.Module
,它将通过torch.fx.trace
进行跟踪并受其限制。