系统概述

Torch-TensorRT 主要是一个 C++ 库，计划提供 Python API。我们使用 Bazel 作为构建系统，目前的目标平台是 Linux x86_64 和 Linux aarch64（仅限原生）。我们使用的编译器是 GCC 7.5.0，该库尚未在早于该版本的编译器上进行测试，因此如果您尝试使用较旧的编译器，可能会出现编译错误。

仓库结构分为

• core：主编译器源代码

• cpp：C++ API

• tests：C++ API、核心和转换器的测试

• py：Python API

• notebooks：使用 Torch-TensorRT 构建的示例应用程序

• docs：文档

• docsrc：文档源文件

• third_party：依赖库的 BUILD 文件

• toolchains：不同平台的工具链

C++ API 尚不稳定，在库成熟之前可能会发生更改，但大多数工作都在核心部分幕后完成。

核心部分有几个主要组成部分：顶层编译器接口，它协调模块的摄取、降低、转换和生成新模块，并将其返回给用户。编译器主要有三个阶段：降低阶段、转换阶段和执行阶段。

编译器阶段

降低

降低阶段

降低由一组 pass（一些来自 PyTorch，一些特定于 Torch-TensorRT）组成，这些 pass 在图 IR 上运行，将大型 PyTorch 算子集映射到简化的算子集，后者更易于转换为 TensorRT。

分区

分区阶段

此阶段是可选的，由用户启用。它指示编译器将节点分离为应在 PyTorch 中运行的节点和应在 TensorRT 中运行的节点。分离标准包括：缺少转换器、算子被用户显式设置为在 PyTorch 中运行，或者节点具有一个标志，该标志告诉分区通过模块回退 pass 在 PyTorch 中运行。

转换

转换阶段

在转换阶段，我们遍历降低后的图并构建等效的 TensorRT 图。转换阶段由三个主要组成部分组成：用于管理编译时数据的上下文、将执行可在编译时解析的操作的求值器库，以及将 JIT 中的算子映射到 TensorRT 的转换器库。

编译和运行时

部署 Torch-TensorRT 程序

最终编译阶段构建一个 TorchScript 程序来运行转换后的 TensorRT 引擎。它接受一个序列化的引擎，并在引擎管理器中实例化它，然后编译器将构建一个 JIT 图，该图引用此引擎并将其包装在一个模块中以返回给用户。当用户执行该模块时，JIT 程序将在 JIT 运行时中运行，该运行时由 Torch-TensorRT 使用用户提供的数据进行扩展。