• 文档 >
  • Torch-TensorRT 解释
快捷方式

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 模块 (torch.nn.Module)

返回值

  • 在第一次调用时触发编译的盒装函数

torch_tensorrt.dynamo.compile (提前编译)

torch_tensorrt.dynamo.compile 是一个 AOT 编译器,模型在显式编译阶段进行编译。然后,这些编译工件可以在以后的日期序列化和重新加载。图通过 torch.export.trace 系统降低到由 核心 ATen 运算符 或适合 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 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系统。立即返回一个包装函数,该函数将在第一次调用时进行编译。

  • 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进行跟踪,并受其限制。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源