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 进一步将这些图降低为仅由 Core ATen Operators 或适用于 TensorRT 加速的精选“高级操作”组成的运算。子图进一步划分为将在 PyTorch 中运行的组件和将进一步编译到 TensorRT 的组件，这基于对运算符的支持。然后，TensorRT 引擎替换受支持的块，并将混合子图返回给 torch.compile 以在调用时运行。
接受的格式¶
torch.fx GraphModule (
torch.fx.GraphModule)PyTorch 模块 (
torch.nn.Module)
返回¶
在首次调用时触发编译的 Boxed-function
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 Graph Modules (torch_tensorrt.fx.compile)¶
此前端几乎完全被 Dynamo 前端取代，Dynamo 前端是 FX 前端可用功能的超集。原始 FX 前端出于向后兼容性的原因保留在代码库中。
接受的格式¶
torch.fx GraphModule (
torch.fx.GraphModule)PyTorch 模块 (
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系统。立即返回一个 boxed-function，该函数将在首次调用时编译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进行跟踪，并受其限制。
