TensorRT - TensorRT 的工作原理

原创已于 2025-11-10 08:34:57 修改 · 658 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#开发语言 #人工智能

于 2025-10-13 02:57:57 首次发布

TensorRT 专栏收录该内容

31 篇文章 ¥9.90 ¥99.00

订阅专栏

超级会员免费看

TensorRT 的工作原理

本节将详细介绍 TensorRT 的工作原理。

对象生命周期

TensorRT 的 API 是基于类的，其中某些类作为其他类的工厂。对于由用户拥有的对象，工厂对象的生命周期必须覆盖它所创建的对象的生命周期。例如，NetworkDefinition 和 BuilderConfig 类是从 Builder 类创建的，这些类的对象应在 Builder 工厂对象销毁之前销毁。

一个重要的例外是从 builder 创建 engine。创建 engine 后，你可以销毁 builder、network、parser 和 build config，并继续使用 engine。

错误处理与日志

在创建 TensorRT 顶层接口（builder、runtime 或 refitter）时，必须提供 Logger（C++、Python）接口的实现。Logger 用于诊断和信息消息，其详细级别可配置。由于 Logger 可能在 TensorRT 生命周期内的任何时间点返回信息，因此它的生命周期必须覆盖该接口的使用期。此外，该实现必须是线程安全的，因为 TensorRT 可能在内部使用工作线程。

对对象的 API 调用将使用与对应顶层接口关联的 logger。例如，在调用 ExecutionContext::enqueueV3() 时，execution context 是从 runtime 创建的 engine 创建的，因此 TensorRT 会使用该 runtime 关联的 logger。

主要的错误处理方法是 ErrorRecorder（C++、Python）接口。可以实现该接口并将其附加到 API 对象，以接收与该对象相关的错误。对象的 recorder 也会传递给它创建的其他对象——例如，如果将 error recorder 附加到 engine 并从该 engine 创建 execution context，则会使用同一个 recorder。如果随后将新的 error recorder 附加到 execution context，则只会接收该 context 的错误。如果生成了错误但没有找到 error recorder，则会通过关联的 logger 发出。

注意 CUDA 错误通常是异步的，所以在单一 CUDA context 中异步执行多次推理或其他 CUDA 工作流时，可能会在与生成错误不同的 execution context 中观察到 GPU 异步错误。

内存

TensorRT 使用大量的设备内存（即 GPU 可直接访问的内存，与 CPU 附加的主机内存不同）。由于设备内存通常是受限资源，了解 TensorRT 的内存使用方式非常重要。

构建阶段

在构建过程中，TensorRT 会为定时层实现分配设备内存。有些实现会消耗大量临时内存，尤其是处理大张量时。可以通过 builder config 的内存池限制来控制最大临时内存。工作空间大小默认等于设备全局内存的大小，但在必要时可以受限。如果 builder 找到某些无法运行的内核，是由于工作空间不足，会发出日志提示。

即使工作空间较小，定时也需要为输入、输出和权重创建缓冲区。TensorRT 能够应对操作系统（OS）因分配不足而返回内存不足的情况。在某些平台，操作系统可能成功分配内存，但随后由于系统内存不足而被内存杀手进程终止。如果发生这种情况，请在重试前尽量释放系统内存。

构建阶段，主机内存中通常会有至少两份权重副本：一份来自原始网络，另一份在构建 engine 时包含在 engine 中。此外，当 TensorRT 合并权重（比如卷积与批归一化合并时），还会创建额外的临时权重张量。

运行阶段

运行时，TensorRT 使用较少的主机内存，但设备内存消耗很大。

engine 在反序列化时会分配设备内存存储模型权重。由于序列化的 engine 几乎包含所有权重，其大小近似于权重所需的设备内存量。

TensorRT 提供了 ICudaEngine::getEngineStat() API，可用于检索 engine 的详细统计信息，包括精确的权重大小。利用 EngineStat 枚举，可以查询如下内容：

kTOTAL_WEIGHTS_SIZE：返回 engine 使用的所有权重总字节数
kSTRIPPED_WEIGHTS_SIZE：对于使用 BuilderFlag::kSTRIP_PLAN 标志构建的 engine，返回被剥离权重的字节数

注意：启用权重流（BuilderFlag::kWEIGHT_STREAMING）功能时，权重可能不会全部复制到设备。getEngineStat(kTOTAL_WEIGHTS_SIZE) 返回的是 engine 使用的所有权重之和，可能与实际分配的 GPU 权重内存不同。

如果在普通 engine 上查询 kSTRIPPED_WEIGHTS_SIZE，则该函数会返回 -1，表示无效查询。

该 API 可用于程序准确监控和管理 engine 的权重内存使用情况，而不仅仅依赖序列化 engine 本身。

ExecutionContext 使用两种设备内存：

某些层实现需要持久内存——例如，一些卷积实现使用边缘掩码，该状态无法像权重一样在 context 之间共享，因为其大小依赖于层输入形状，可能因 context 而异。该内存在创建 execution context 时分配，并持续到其生命周期结束。
Enqueue 内存用于在处理网络时保存中间结果。这部分内存用于中间张量（激活内存），同时也用于层实现所需的临时存储（scratch 内存），其上限由 IBuilderConfig::setMemoryPoolLimit() 控制。TensorRT 通过如下方式优化内存使用：
- 通过在具有不相交生命周期的激活张量之间共享设备内存块
- 在可行时，允许临时（scratch）张量占用未使用的激活内存。因此，TensorRT 所需 enqueue 内存范围为 {总激活内存, 总激活内存 + 最大 scratch 内存}。

可以选择不带 enqueue 内存创建 execution context，使用 ICudaEngine::createExecutionContextWithoutDeviceMemory()，并在网络执行期间自行提供该内存。这样可在多个非并发运行的 context 之间共享，或在推理未运行时用于其他用途。所需 enqueue 内存可通过 ICudaEngine::getDeviceMemorySizeV2() 获取。

execution context 使用的持久内存和 scratch 内存的信息会在 builder 构建网络时以 kINFO 严重级别输出日志。例如：

[08/12/2021-17:39:11] [I] [TRT] Total Host Persistent Memory: 106528
[08/12/2021-17:39:11] [I] [TRT] Total Device Persistent Memory: 29785600
[08/12/2021-17:39:11] [I] [TRT] Max Scratch Memory: 9970688

默认情况下，TensorRT 直接从 CUDA 分配设备内存。不过，可以将 TensorRT 的 IGpuAllocator（C++、Python）接口的实现附加到 builder 或 runtime，从而自行管理设备内存。如果应用希望控制所有 GPU 内存并为 TensorRT 子分配，而不是让 TensorRT 直接从 CUDA 分配，这非常有用。

NVIDIA cuDNN 和 NVIDIA cuBLAS 也可能占用大量设备内存。TensorRT 允许通过 builder config 的 TacticSources（C++、Python）属性控制这些库是否用于推理。某些插件实现需要这些库，因此排除它们时网络可能无法成功编译。如果设置了正确的 tactic sources，则 cudnnContext 和 cublasContext 句柄会通过 IPluginV2Ext::attachToContext() 传递给插件。

CUDA 基础设施和 TensorRT 设备代码也会消耗设备内存。具体内存消耗依平台、设备和 TensorRT 版本而异。可以使用 cudaGetMemInfo 获取设备内存总量。

TensorRT 在 builder 和 runtime 的关键操作前后测量内存使用情况，这些统计信息会打印到 TensorRT 的信息日志。例如：

[MemUsageChange] Init CUDA: CPU +535, GPU +0, now: CPU 547, GPU 1293 (MiB)

表示 CUDA 初始化后内存变化。CPU +535, GPU +0 是初始化后增加的内存量，now: 后是 CUDA 初始化后的 CPU/GPU 内存快照。

注意：

在多租户环境中，cudaGetMemInfo 和 TensorRT 报告的内存使用容易发生竞争（race condition），即其他进程或线程分配/释放了新内存。由于 CUDA 不控制统一内存设备上的内存，cudaGetMemInfo 在这些平台上的结果可能不准确。

CUDA 延迟加载

CUDA 延迟加载是一项 CUDA 功能，可显著降低 TensorRT 的峰值 GPU 和主机内存使用量，并加速 TensorRT 初始化，性能影响很小（<1%）。具体内存和初始化时间节省量取决于模型、软件栈、GPU 平台等。通过设置环境变量 CUDA_MODULE_LOADING=LAZY 启用。更多信息见 NVIDIA CUDA 文档。

L2 持久缓存管理

NVIDIA Ampere 及更高架构支持 L2 缓存持久特性，可优先保留 L2 缓存行，减少 DRAM 访问和功耗。

缓存分配按 execution context 进行，可通过 context 的 setPersistentCacheLimit 方法启用。所有 context 的总持久缓存（及其他使用该特性的组件）不应超过 cudaDeviceProp::persistingL2CacheMaxSize。更多信息见 NVIDIA CUDA 最佳实践指南。

线程

TensorRT 对象通常不是线程安全的，客户端必须在不同线程间序列化访问对象。

预期的运行时并发模型是，不同线程操作不同的 execution context。context 在执行期间包含网络状态（激活值等），因此在不同线程并发使用同一 context 会导致未定义行为。

为支持此模型，以下操作是线程安全的：

对 runtime 或 engine 的非修改性操作
从 TensorRT runtime 反序列化 engine
从 engine 创建 execution context
注册和注销插件

在不同线程中使用多个 builder 不会有线程安全问题，但 builder 通过计时确定最优内核参数，在同一 GPU 上并行使用多个 builder 会影响计时以及 TensorRT 构建最优 engine 的能力。使用不同 GPU 并行构建则没有此问题。

确定性

TensorRT builder 通过计时选择实现层的最快内核。计时受噪声影响，如 GPU 上其他任务、GPU 时钟波动等。因此，连续构建时可能不会选择同一实现。

一般来说，不同内核实现会使用不同顺序的浮点操作，导致输出有微小差异。这些差异对最终结果影响通常很小。但当 TensorRT 配置为多精度调优时，FP16 和 FP32 内核之间的差异可能较大，尤其是网络未充分正则化或对数值漂移敏感时。

其他可能导致不同内核选择的配置包括不同的输入尺寸（如 batch size）或优化点不同（参见动态形状使用）。

可编辑定时缓存机制允许强制 builder 为某一层选择特定实现。可用于保证 builder 每次选择相同内核。更多信息见算法选择与可复现构建。

engine 构建完成后，除 IFillLayer 和 IScatterLayer 外是确定性的：在相同运行环境下输入相同数据会产生相同输出。

虽然 TensorRT 可以保证相同输入跨运行的一致性，但当将相同输入分别放在 batch 的不同 slot 或与不同输入共同 batch 时，输出可能不同。

`IFillLayer` 的确定性

当使用 RANDOM_UNIFORM 或 RANDOM_NORMAL 操作将 IFillLayer 添加到网络时，上述确定性保证不再有效。每次调用会根据 RNG 状态生成张量，并更新 RNG 状态。该状态按 execution context 存储。TensorRT 支持在构建时启用分布独立特性（BuilderFlag::kDISTRIBUTIVE_INDEPENDENCE），可消除这种差异，具体约束见分布独立确定性。

`IScatterLayer` 的确定性

如果网络中添加了 IScatterLayer 且输入张量 indices 有重复项，则对于 ScatterMode::kELEMENT 和 ScatterMode::kND，上述确定性保证无效。此外，输入更新张量中的某个值会被任意选中。

分布独立确定性

启用 BuilderFlag::kDISTRIBUTIVE_INDEPENDENCE 且某层声明输出的 axis i 为分布轴时，TensorRT 保证该层行为等同于在分布轴上每次评估都使用相同操作，即如果某些输入在分布轴上相同，则对应输出也一致。

TensorRT 对分布轴的定义如下：

对于 IMatrixMultiplyLayer：所有不是向量或矩阵维度的轴都是分布轴
对于 reduction 层：所有非 reduction 轴都是分布轴
对于 einsum 层：设 n 为最左侧 reduction 轴，n 左侧的所有轴都是分布轴

注意分布独立特性有如下约束：

启用 BuilderFlag::kDISTRIBUTIVE_INDEPENDENCE 时，禁用多 profile，因为不同 profile 的内核可能不同，结果也有微小差异
启用该特性时，尤其是动态形状情况下，可能有性能下降

运行时选项

TensorRT 提供多种运行时库以满足不同用例。C++ 应用运行 TensorRT engine 时应链接如下之一：

默认运行时为主库（libnvinfer.so/.dll）
精简运行时库（libnvinfer_lean.so/.dll）比默认库小，仅包含运行版本兼容 engine 所需代码。主要限制是无法重训练（refit）或序列化 engine。
分发运行时（libnvinfer_dispatch.so/.dll）是一个小型 shim 库，可加载精简运行时并重定向调用。分发运行时可加载旧版精简运行时，并结合 builder 配置用于新版本 TensorRT 与旧 plan 文件的兼容。使用分发运行时类似手动加载精简运行时，但它会检查 lean runtime 是否实现了 API，并在可能时做参数映射以支持 API 变更。

精简运行时包含的运算符实现比默认运行时少。由于 TensorRT 在构建时选择运算符实现，构建 engine 时需指定为精简运行时启用版本兼容。它可能比为默认运行时构建的 engine 慢。

精简运行时包含分发运行时的全部功能，默认运行时包含精简运行时的全部功能。

TensorRT 提供了对应各库的 Python 包：

tensorrt

Python 包。是默认运行时的 Python 接口。

tensorrt_lean

Python 包。是精简运行时的 Python 接口。

tensorrt_dispatch

Python 包。是分发运行时的 Python 接口。

Python 应用运行 TensorRT engine 时应导入上述包之一，以加载适合自己用例的库。