• Docs >
  • Torch-TensorRT Explained
Shortcuts

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.scripttorch.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.exporttorch_tensorrt.save 保存。

  • torchscriptts:通过 TorchScript 堆栈运行图。如果输入模块是 torch.nn.Module,则它必须是“可脚本化的”,因为该模块将通过 torch.jit.script 进行编译。返回一个 torch.jit.ScriptModule,可以立即运行或通过 torch.savetorch_tensorrt.save 保存。

  • fx: 通过 torch.fx 堆栈运行图。如果输入模块是 torch.nn.Module,它将使用 torch.fx.trace 进行跟踪,并受其限制。