TensorRT 的功能#

本节概述了 TensorRT 的功能。旨在为所有 TensorRT 用户提供帮助。

C++ 和 Python API#

TensorRT 的 API 具有 C++ 和 Python 的语言绑定,功能几乎相同。该 API 促进了与 Python 数据处理工具包和库(如 NumPy 和 SciPy)的互操作性。C++ API 可以更高效,并且可能更好地满足某些合规性要求,例如在汽车应用中。

注意

Python API 并非在所有平台上都可用。有关更多信息,请参阅 支持矩阵

编程模型#

TensorRT 分两个阶段运行。在第一阶段(通常离线执行),您向 TensorRT 提供模型定义,TensorRT 会针对目标 GPU 对其进行优化。在第二阶段,您使用优化的模型运行推理。

构建阶段#

TensorRT 构建阶段的最高级接口是 Builder (C++, Python)。构建器负责优化模型并生成 Engine (引擎)。

要构建引擎,您必须

  1. 创建网络定义。

  2. 为构建器指定配置。

  3. 调用构建器以创建引擎。

NetworkDefinition 接口 (C++, Python) 定义了模型。将模型传输到 TensorRT 最常见的方法是从框架中以 ONNX 格式导出模型,并使用 TensorRT 的 ONNX 解析器填充网络定义。但是,您也可以使用 TensorRT 的 Layer (层) (C++, Python) 和 Tensor (张量) (C++, Python) 接口逐步构建定义。

无论您选择哪种方式,您还必须定义哪些张量是网络的输入和输出。未标记为输出的张量被视为构建器可以优化掉的瞬态值。输入和输出张量必须命名,以便 TensorRT 知道如何在运行时将输入和输出缓冲区绑定到模型。

BuilderConfig 接口 (C++, Python) 用于指定 TensorRT 应如何优化模型。在可用的配置选项中,您可以控制 TensorRT 降低计算精度的能力、控制内存和运行时执行速度之间的权衡,以及约束 CUDA 内核的选择。由于构建器可能需要几分钟或更长时间才能运行,因此您还可以控制构建器如何搜索内核以及缓存搜索结果以供后续运行使用。

在您拥有网络定义和构建器配置后,您可以调用构建器来创建引擎。构建器消除无效计算、折叠常量、重新排序和组合操作,以便在 GPU 上更高效地运行。它可以通过简单地以 16 位浮点运行浮点计算或通过量化浮点值来降低浮点计算的精度,以便可以使用 8 位整数执行计算。它还为每个层提供了多种具有不同数据格式的实现。然后,它计算执行模型的最佳计划,从而最大限度地减少内核执行和格式转换的综合成本。

构建器以序列化形式(称为 plan (计划))创建引擎,该计划可以立即反序列化或保存到磁盘以供以后使用。

注意

  • 默认情况下,TensorRT 引擎特定于 TensorRT 版本和创建它们的 GPU。请参阅 版本兼容性硬件兼容性 部分,以配置引擎以实现向前兼容性。

  • TensorRT 的网络定义没有深拷贝参数数组(例如卷积的权重)。因此,在构建阶段完成之前,您不得释放这些数组的内存。使用 ONNX 解析器导入网络时,解析器拥有权重,因此在构建阶段完成之前不得销毁它。

  • 构建器对算法进行计时以确定最快的算法。与其他 GPU 工作并行运行构建器可能会扰乱计时,从而导致优化效果不佳。

运行时阶段#

TensorRT 执行阶段的最高级接口是 Runtime (运行时) (C++, Python)。

使用运行时,您通常会执行以下步骤

  1. 反序列化计划以创建引擎。

  2. 从引擎创建执行上下文。

然后,重复

  1. 填充推理的输入缓冲区。

  2. 在执行上下文中调用 enqueueV3() 以运行推理。

Engine (引擎) 接口 (C++, Python) 表示优化的模型。您可以查询引擎以获取有关网络输入和输出张量的信息 - 预期维度、数据类型、数据格式等等。

从引擎创建的 ExecutionContext (执行上下文) 接口 (C++, Python) 是调用推理的主要接口。执行上下文包含与特定调用关联的所有状态 - 因此,您可以将多个上下文与单个引擎关联并并行运行它们。

调用推理时,您必须在适当的位置设置输入和输出缓冲区。根据数据的性质,这可能位于 CPU 或 GPU 内存中。如果不明显,根据您的模型,您可以查询引擎以确定要提供缓冲区的内存空间。

设置缓冲区后,可以将推理排队 (enqueueV3)。所需的内核在 CUDA 流上排队,并且控制权会尽快返回给应用程序。某些网络需要在 CPU 和 GPU 之间进行多次控制传输,因此控制权可能需要更长时间才能返回。要等待异步执行完成,请使用 cudaStreamSynchronize 在流上同步。

插件#

TensorRT 具有 Plugin (插件) 接口,允许应用程序实现 TensorRT 本身不支持的操作。在转换网络时,ONNX 解析器可以找到在 TensorRT 的 PluginRegistry 中创建和注册的插件。

TensorRT 附带一个插件库;其中许多插件和一些其他插件的源代码可以在 GitHub: TensorRT plugin 上找到。

您还可以编写自己的插件库并将其与引擎一起序列化。

如果需要 cuDNN 或 cuBLAS,请安装该库,因为 TensorRT 不再附带它们或依赖于它们。要获取 cudnnContext*cublasContext*,必须使用 nvinfer1::IBuilderConfig::setTacticSource() 设置相应的 TacticSource 标志。

有关更多详细信息,请参阅 使用自定义层扩展 TensorRT 部分。

类型和精度#

支持的类型#

TensorRT 支持 FP32、FP16、BF16、FP8、FP4、INT4、INT8、INT32、INT64、UINT8 和 BOOL 数据类型。有关层 I/O 数据类型规范,请参阅 TensorRT 算子文档

  • FP32、FP16、BF16:非量化浮点类型

  • INT8:低精度整数类型

    • 隐式量化

      • 解释为量化整数。具有 INT8 类型的张量必须具有关联的比例因子(通过校准或 setDynamicRange API)。

    • 显式量化

      • 解释为有符号整数。与 INT8 类型之间的转换需要显式的 Q/DQ 层。

  • INT4:用于权重压缩的低精度整数类型

    • INT4 用于仅权重量化。在执行计算之前需要反量化。

    • 与 INT4 类型之间的转换需要显式的 Q/DQ 层。

    • INT4 权重应通过每字节打包两个元素进行序列化。有关更多信息,请参阅 量化权重 部分。

  • FP8:低精度浮点类型

    • 8 位浮点类型,其中 1 位用于符号,4 位用于指数,3 位用于尾数

    • 与 FP8 类型之间的转换需要显式的 Q/DQ 层。

  • FP4:窄精度浮点类型

    • 4 位浮点类型,其中 1 位用于符号,2 位用于指数,1 位用于尾数

    • 与 FP4 类型之间的转换需要显式的 Q/DQ 层。

    • FP4 权重应通过每字节打包两个元素进行序列化。有关更多信息,请参阅 量化权重 部分。

  • UINT8:无符号整数 I/O 类型

    • 数据类型仅可用作网络 I/O 类型。

    • UINT8 中的网络级输入必须在使用 CastLayer 从 UINT8 转换为 FP32 或 FP16 后才能在其他操作中使用。

    • UINT8 中的网络级输出必须由显式插入到网络中的 CastLayer 生成(仅支持从 FP32/FP16 到 UINT8 的转换)。

    • 不支持 UINT8 量化。

    • ConstantLayer 不支持 UINT8 作为输出类型。

  • BOOL

    • 布尔类型用于受支持的层。

强类型与弱类型#

向 TensorRT 提供网络时,您可以指定它是强类型还是弱类型,默认情况下为弱类型。

对于强类型网络,TensorRT 的优化器将根据网络输入类型和算子规范静态推断中间张量类型,这与框架中的类型推断语义相匹配。然后,优化器将严格遵守这些类型。有关更多信息,请参阅 强类型网络 部分。

对于弱类型网络,如果可以提高性能,TensorRT 的优化器可能会为张量替换不同的精度。在这种模式下,TensorRT 默认对所有浮点运算使用 FP32,但有两种方法可以配置不同的精度级别

  • 要在模型级别控制精度,BuilderFlag 选项 (C++, Python) 可以向 TensorRT 指示,在搜索最快实现时,它可以选择较低精度的实现(并且由于较低精度通常更快,因此通常会这样做)。

    例如,通过设置一个标志,您可以轻松指示 TensorRT 对整个模型使用 FP16 计算。对于输入动态范围约为 1 的正则化模型,这通常会产生显着的加速,而精度变化可忽略不计。

  • 为了更精细的控制,如果某个层必须以更高的精度运行,因为网络的一部分对数值敏感或需要高动态范围,则可以为该层指定算术精度。

有关更多信息,请参阅 弱类型网络中的降低精度 部分。

量化#

TensorRT 支持量化浮点,其中浮点值被线性压缩并舍入为低精度量化类型(INT8、FP8、INT4、FP4)。这显着提高了算术吞吐量,同时降低了存储需求和内存带宽。量化浮点张量时,TensorRT 必须知道其动态范围 - 要表示的重要值范围 - 超出此范围的值在量化时会被钳制。

构建器可以基于代表性输入数据(当前仅支持 INT8)计算动态范围信息(calibration (校准))。或者,您可以在框架中执行量化感知训练,并将模型与必要的动态范围信息导入到 TensorRT。

有关更多信息,请参阅 使用量化类型 部分。

张量和数据格式#

定义网络时,TensorRT 假定多维 C 风格数组表示张量。每个层对其输入都有特定的解释:例如,2D 卷积将假定其输入的最后三个维度采用 CHW 格式 - 没有选项使用例如 WHC 格式。有关每个层如何解释输入,请参阅 TensorRT 算子文档

请注意,张量最多限制为 2^31-1 个元素。

在优化网络时,TensorRT 在内部执行转换(包括转换为 HWC,但也包括更复杂的格式)以使用最快的 CUDA 内核。格式通常被选择用于优化性能,应用程序无法控制这些选择。但是,底层数据格式在 I/O 边界(网络输入和输出以及插件之间的数据传递)处公开,以允许应用程序最大限度地减少不必要的格式转换。

有关更多信息,请参阅 I/O 格式 部分。

动态形状#

默认情况下,TensorRT 根据定义模型时的输入形状(批大小、图像大小等)优化模型。但是,可以将构建器配置为在运行时调整输入维度。要启用此功能,请在构建器配置中指定一个或多个 OptimizationProfile (C++, Python) 实例,其中包含每个输入的最小和最大形状以及该范围内的优化点。

TensorRT 为每个配置文件创建一个优化的引擎,选择适用于 [最小,最大] 范围内的所有形状且对于优化点最快的 CUDA 内核 - 通常每个配置文件都有不同的内核。然后,您可以在运行时在配置文件之间进行选择。

有关更多信息,请参阅 使用动态形状 部分。

DLA#

TensorRT 支持 NVIDIA 的深度学习加速器 (DLA),它是许多 NVIDIA SoC 上的专用推理处理器,支持 TensorRT 层的一个子集。TensorRT 允许您在 DLA 上执行部分网络,其余部分在 GPU 上执行;对于可以在任一设备上执行的层,您可以在构建器配置中按层选择目标设备。

有关更多信息,请参阅 使用 DLA 部分。

更新权重#

构建引擎时,您可以指定其权重稍后可能会更新。如果您经常更新模型的权重而不更改结构,例如在强化学习中或在重新训练模型时保持相同的结构,这将非常有用。权重更新是使用 Refitter (C++, Python) 接口执行的。

有关更多信息,请参阅 重新拟合引擎 部分。

流式传输权重#

可以将 TensorRT 配置为在网络执行期间将网络的权重从主机内存流式传输到设备内存,而不是在引擎加载时将它们放置在设备内存中。这使得权重大于可用 GPU 内存的模型能够运行,但可能会显着增加延迟。权重流式传输是一个选择加入功能,在构建时 (BuilderFlag::kWEIGHT_STREAMING) 和运行时 (ICudaEngine::setWeightStreamingBudgetV2) 都是如此。

注意

权重流式传输仅在强类型网络中受支持。有关更多信息,请参阅 权重流式传输 部分。

trtexec#

samples 目录中包含一个名为 trtexec 的命令行包装器工具。trtexec 是一个允许您在不开发应用程序的情况下使用 TensorRT 的工具。trtexec 工具主要有三个用途

  1. 在随机或用户提供的输入数据上对网络进行基准测试

  2. 从模型生成序列化引擎。

  3. 从构建器生成序列化计时缓存

有关更多信息,请参阅 trtexec 部分。

Polygraphy#

Polygraphy 是一个旨在帮助在 TensorRT 和其他框架中运行和调试深度学习模型的工具包。它包括一个 Python API 和一个使用此 API 构建的 命令行界面 (CLI)

除了其他功能外,使用 Polygraphy,您可以

  • 在多个后端(如 TensorRT 和 ONNX-Runtime)之间运行推理并比较结果 (API, CLI)。

  • 将模型转换为 TensorRT 引擎等格式,并进行训练后量化 (API, CLI)。

  • 查看有关各种模型类型的信息 (CLI)。

  • 在命令行修改 ONNX 模型

    • 提取子图 (CLI)

    • 简化和清理 (CLI)

    • 在 TensorRT 中隔离错误策略 (CLI)

有关更多信息,请参阅 Polygraphy 存储库