快捷方式

系统概述

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

仓库结构如下

  • core: 主要编译器源代码

  • cpp: C++ API

  • tests: C++ API、核心和 Converter 的测试

  • py: Python API

  • notebooks: 使用 Torch-TensorRT 构建的示例应用

  • docs: 文档

  • docsrc: 文档源文件

  • third_party: 依赖库的 BUILD 文件

  • toolchains: 不同平台的 Toolchain

C++ API 尚不稳定,在库成熟之前可能会发生变化,但大部分工作都在核心内部进行。

核心包含几个主要部分:顶层编译器接口,负责协调模块的摄取、Lowering、转换以及生成新模块并将其返回给用户。编译器有三个主要阶段:Lowering 阶段、转换阶段和执行阶段。

编译器阶段

Lowering

Lowering 阶段

Lowering 阶段由一系列 Pass 组成(部分来自 PyTorch,部分特定于 Torch-TensorRT),这些 Pass 在图 IR 上运行,将庞大的 PyTorch opset 映射到更容易转换为 TensorRT 的精简 opset。

Partitioning

Partitioning 阶段

此阶段是可选的,由用户启用。它指示编译器将节点分离为应在 PyTorch 中运行的节点和应在 TensorRT 中运行的节点。分离的标准包括:缺少 Converter、用户明确将 operator 设置为在 PyTorch 中运行,或者节点具有通过模块回退 Pass 告知 Partitioning 在 PyTorch 中运行的标志。

转换

转换阶段

在转换阶段,我们遍历 Lowering 后的图并构建等效的 TensorRT 图。转换阶段由三个主要组件组成:用于管理编译时数据的上下文、用于执行可在编译时解析的操作的 Evaluator 库,以及用于将 JIT 的 op 映射到 TensorRT 的 Converter 库。

编译与运行时

部署 Torch-TensorRT 程序

最终的编译阶段构建一个 TorchScript 程序来运行转换后的 TensorRT Engine。它获取一个序列化的 Engine,并在 Engine Manager 中实例化它,然后编译器将构建一个引用此 Engine 并将其封装在 Module 中的 JIT 图,最后将其返回给用户。当用户执行该 Module 时,JIT 程序将在由 Torch-TensorRT 扩展的 JIT 运行时中运行,并使用用户提供的数据。

文档

获取 PyTorch 全面的开发者文档

查看文档

教程

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

查看教程

资源

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

查看资源