系统概述

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 运行时中运行，并使用用户提供的数据。