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 GraphModule (
torch.fx.GraphModule)PyTorch Module (
torch.nn.Module)
返回值¶
触发首次调用时进行编译的 boxed 函数
torch_tensorrt.dynamo.compile (提前编译)¶
torch_tensorrt.dynamo.compile 是一个提前（AOT）编译器，模型在显式的编译阶段进行编译。这些编译产物随后可以被序列化并在以后重新加载。图会通过 torch.export.trace 系统进行降低，生成一个由 Core ATen Operators 或适合 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 Graph Modules (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 尝试通过特定的中间表示（IR）来降低提供的模型。
ir 选项¶
torch_compile: 使用torch.compile系统。立即返回一个 boxed 函数，该函数将在首次调用时进行编译。dynamo: 将图通过torch.export/ torchdynamo 栈运行。如果输入模块是torch.nn.Module，则它必须是“可导出跟踪”（export-traceable）的，因为模块将使用torch.export.export进行跟踪。返回一个torch.fx.GraphModule，该模块可以立即运行，或通过torch.export.export或torch_tensorrt.save保存。torchscript或ts: 将图通过 TorchScript 栈运行。如果输入模块是torch.nn.Module，则它必须是“可脚本化”（scriptable）的，因为模块将使用torch.jit.script进行编译。返回一个torch.jit.ScriptModule，该模块可以立即运行，或通过torch.save或torch_tensorrt.save保存。fx: 将图通过torch.fx栈运行。如果输入模块是torch.nn.Module，它将使用torch.fx.trace进行跟踪，并受其限制。
