系统概览

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 不稳定，可能会发生变化，直到库成熟为止，尽管大部分工作都在核心部分进行。

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

编译器阶段

降低

降低阶段

降低阶段由一组传递（一部分来自 PyTorch，一部分是 Torch-TensorRT 特有的）组成，这些传递在图 IR 上运行，将大的 PyTorch 操作集映射到一个简化的操作集，以便更容易转换为 TensorRT。

分区

分区阶段

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

转换

转换阶段

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

编译和运行时

部署 Torch-TensorRT 程序

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