ONNX Runtime 中的 Cache Mode

原创于 2025-09-03 11:10:15 发布 · 515 阅读

CC 4.0 BY-SA版权

文章标签：

核心概念

Cache Mode（缓存模式） 是 ONNX Runtime 的一个特性，主要用于在 TensorRT 和 CUDA 等执行提供程序（Execution Providers, EP）中显著减少模型首次推理的延迟。

它的核心思想是：将引擎编译/模型优化的结果序列化并保存到磁盘上。这样，下次在相同硬件和相同模型上运行时，就可以直接加载缓存，跳过耗时的编译阶段。

为什么需要 Cache Mode？

在使用像 TensorRT 这样的 EP 时，ONNX Runtime 需要执行一个“准备”阶段：

模型优化：TensorRT 会读取 ONNX 模型。
计算图优化：进行算子融合、常量折叠等。
引擎编译：为模型中的每一层选择最优的 kernel（内核算法），并为目标 GPU 编译出高性能的推理引擎（inference engine）。

这个过程非常耗时，尤其是对于大模型，首次推理（即包含这个编译过程的推理）可能需要几十秒甚至几分钟，但后续的推理速度会极快。

Cache Mode 就是为了消除每次启动程序时都要经历的“首次推理”的长延迟问题。

它是如何工作的？

缓存模式的工作流程可以分为两个阶段：缓存缺失（Cache Miss） 和缓存命中（Cache Hit）。

1. 缓存缺失 (Cache Miss) - “创建缓存”

当第一次在某个硬件（GPU型号）、某个模型、以及特定的配置（如精度FP16/FP32、输入尺寸等）下运行时，缓存文件不存在。

ONNX Runtime 会执行完整的 TensorRT 编译过程。
编译完成后，它会将生成的引擎计划（Engine Plan） 序列化成一个文件，保存到你指定的缓存路径下。
这次运行的延迟会很高。

2. 缓存命中 (Cache Hit) - “使用缓存”

当你再次运行完全相同的模型和配置时：

ONNX Runtime 会先检查缓存目录。
如果找到匹配的缓存文件，它就会直接反序列化并加载已编译好的引擎。
跳过耗时的编译阶段，首次推理的延迟变得非常低，几乎与后续推理一样快。

如何启用和配置？

你需要在创建 ONNX Runtime 会话选项（SessionOptions）时，为对应的执行提供程序（如 TensorRT）设置缓存路径。

Python 示例

python

import onnxruntime as ort

# 创建会话选项
so = ort.SessionOptions()

# 配置 TensorRT Provider 的选项
trt_ep_options = {
    'trt_engine_cache_enable': True,   # 启用缓存功能
    'trt_engine_cache_path': './trt_cache',  # 指定缓存文件的存放目录
    # 其他 TRT 选项，例如：
    # 'trt_fp16_enable': True,
}

# 创建会话，并指定使用 TensorRT execution provider
providers = [
    ('TensorrtExecutionProvider', trt_ep_options),
    'CUDAExecutionProvider',  # 可备选的其他EP
    'CPUExecutionProvider',
]

session = ort.InferenceSession(
    'your_model.onnx',
    sess_options=so,
    providers=providers
)

C++ 示例

cpp

Ort::SessionOptions session_options;
OrtTensorRTProviderOptionsV2* trt_options = nullptr;

// 创建一个 TensorRT 选项结构并配置它
// ... (代码细节较多，以下是概念性伪代码)
trt_options->trt_engine_cache_enable = 1;
trt_options->trt_engine_cache_path = "./trt_cache";

// 将选项添加到会话中
session_options.AppendExecutionProvider_TensorRT(*trt_options);

// 创建会话
Ort::Session session(env, "your_model.onnx", session_options);

关键注意事项

缓存的一致性：缓存文件与硬件、模型、配置紧密绑定。以下变化都会导致缓存失效，需要重新生成：
- GPU 型号：在 V100 上生成的缓存不能在 A100 上使用。
- TensorRT 版本：升级或降级 TensorRT 后，缓存通常需要重新生成。
- ONNX Runtime 版本：更改 ORT 版本也可能导致不兼容。
- 模型结构：模型有任何更改（算子、图结构等）。
- 提供商配置：任何影响引擎编译的选项更改，如：精度（FP32/FP16/INT8）、输入动态尺寸的范围、最大工作空间大小等。
输入动态尺寸：如果模型接受动态输入尺寸，TensorRT 会为不同尺寸的输入生成多个“优化配置文件”（optimization profiles）。每个配置文件都可能被缓存。你需要确保 trt_engine_cache_path 有足够的磁盘空间来存储这些可能的多份缓存。
并非所有 EP 都支持：缓存模式主要是为 TensorRT EP 设计的强大功能。其他 EP（如 CUDA、OpenVINO、DML）可能有自己类似的缓存或序列化机制，但接口和行为可能不同。
磁盘空间：缓存文件可能很大（对于大模型可能达到几百MB甚至GB级），请确保指定的缓存目录有足够的空间。

总结

特性	描述
目的	消除 TensorRT 等 EP 的首次推理编译延迟，极大加快模型加载速度。
工作原理	将编译好的引擎计划序列化到磁盘，下次运行时直接反序列化加载。
启用方式	在 Session Options 中为对应 EP 设置 `trt_engine_cache_enable` 和 `trt_engine_cache_path`。
适用场景	生产环境部署，需要频繁创建会话的场景。模型和硬件配置固定的环境。
主要限制	缓存与硬件、软件版本、模型、配置严格绑定，任何更改都可能导致缓存失效。