Torch-TensorRT 解释¶
Torch-TensorRT 是一个针对 NVIDIA GPU 的 PyTorch 模型编译器,通过 TensorRT 模型优化 SDK 实现。它旨在为 PyTorch 模型提供更好的推理性能,同时仍然保持 PyTorch 的出色人体工程学设计。
Dynamo 前端¶
Dynamo前端是Torch-TensorRT的默认前端。它利用了来自PyTorch的dynamo编译器堆栈。
torch.compile (即时编译)¶
torch.compile 是一个JIT编译器堆栈,因此,编译会延迟到首次使用时进行。这意味着随着图中条件的变化,图将自动重新编译。
这为用户提供了最大的运行时灵活性,但在序列化方面限制了选项。
在底层,torch.compile 委托它认为可以降低到 Torch-TensorRT 的子图。Torch-TensorRT 进一步将这些图降低为仅由 Core ATen Operators 或适合 TensorRT 加速的选定“高级操作”组成的操作。子图进一步划分为将在 PyTorch 中运行的组件和基于操作符支持进一步编译到 TensorRT 的组件。然后,TensorRT 引擎替换支持的块,并将混合子图返回到 torch.compile 以在调用时运行。
接受的格式¶
torch.fx 图模块 (
torch.fx.GraphModule)PyTorch 模块 (
torch.nn.Module)
返回¶
首次调用时触发编译的Boxed函数
torch_tensorrt.dynamo.compile (提前编译)¶
torch_tensorrt.dynamo.compile 是一个AOT编译器,模型在显式编译阶段进行编译。这些编译产物可以在以后序列化并重新加载。
图通过 torch.export.trace 系统被降低为由 Core ATen Operators 或适合TensoRT加速的“高级操作”组成的图。
子图根据对操作符的支持进一步划分为将在PyTorch中运行的组件和进一步编译为TensorRT的组件。然后,TensorRT引擎替换支持的块,
并将混合子图打包到可以序列化和重新加载的 ExportedProgram 中。
接受的格式¶
torch.export.ExportedProgram (
torch.export.ExportedProgram)torch.fx GraphModule (
torch.fx.GraphModule) (通过torch.export.export)PyTorch 模块 (
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 模块 (
torch.jit.ScriptModule)PyTorch 模块 (
torch.nn.Module) (通过torch.jit.script或torch.jit.trace)
返回¶
TorchScript 模块 (
torch.jit.ScriptModule)
FX 图形模块 (torch_tensorrt.fx.compile)¶
这个前端几乎已经完全被Dynamo前端所取代,Dynamo前端是FX前端可用功能的超集。出于向后兼容的原因,原始的FX前端仍然保留在代码库中。
接受的格式¶
torch.fx 图模块 (
torch.fx.GraphModule)PyTorch 模块 (
torch.nn.Module) (通过torch.fx.trace)
返回¶
torch.fx 图模块 (
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进行跟踪,并受其限制。