OnnxRuntime EP 上下文缓存功能设计

目录

背景

OnnxRuntime 执行提供程序使用户能够在由后端 SDK（如 QNN、OpenVINO、Vitis AI 等）支持的不同类型的硬件加速器上推理 Onnx 模型。执行提供程序将 Onnx 模型转换为后端 SDK 所需的图形格式，并将其编译为硬件所需的格式。特别是对于 NPU 领域，转换和编译过程需要很长时间才能完成，尤其是对于 LLM 模型。在某些情况下，会话创建时间长达数十分钟，这严重影响了用户体验。为了避免转换和编译成本，大多数后端 SDK 都提供了将预编译模型转储到二进制文件的功能。预编译模型可以直接由后端 SDK 加载并在目标设备上执行。通过这种方式，可以大大缩短会话创建时间。为了实现这一点，OnnxRuntime 在 MS 域中定义了一个名为 EPContext 的贡献 Op。

EPContext Op 模式

Op 域: com.microsoft 节点输入和输出: variadic 域: com.microsoft 属性

属性 数据类型 描述
main_context int64 1（默认值）：此节点指向包含此节点引用的图形的 EP 上下文内容。
0：此节点不指向任何 EP 上下文内容。期望从此字段为 1 的节点获取图形。
一些 EP 支持 1 个单一上下文包含多个图形。main_context=1 的 EPContext 节点指的是真实上下文。并且上下文包含由 main_context=0 的其他节点引用的图形。
ep_cache_context 字符串 如果 embed_mode=1，则为 EP 上下文的有效负载；如果 embed_mode=0，则为上下文文件的路径。
该路径是 Onnx 模型文件的相对路径。它可以是文件名，也可以是子文件夹/文件名
embed_mode int64 1（默认值）：ep_cache_context 包含上下文内容的有效负载。
0：ep_cache_context 是上下文二进制文件路径。
ep_sdk_version 字符串 可选。用于生成节点的 SDK 版本。
onnx_model_filename 字符串 可选。原始 Onnx 模型文件名。
hardware_architecture 字符串 可选。硬件架构。
partition_name 字符串 可选。OnnxRuntime 分区图名称。
source 字符串 可选。用于生成节点的源。应该是 EP 标识的键，以便 OnnxRuntime 可以支持使用不同 EP 运行的多个 EPContext 节点。例如，QNN EP 仅接受源=QNN 或 QnnExecutionProvider 的节点，OpenVINO EP 仅接受源=OpenVINOExecutionProvider 的节点。
notes 字符串 可选。特定 EP 所需的附加信息。
max_size int64 可选。上下文中的最大大小。用法取决于 EP。默认为 0。

EP Context node example

会话选项 描述
ep.context_enable 仅用于上下文模型生成。
1：启用 OnnxRuntime 以转储上下文缓存模型。
0（默认值）：禁用。
ep.context_file_path 指定转储模型的的文件路径。
对于上下文模型生成，默认为 original_file_name_ctx.onnx。
对于模型推理，如果用户从内存缓冲区加载模型，并且 EP 上下文二进制文件在 Onnx 模型之外，则用户需要设置此选项。OnnxRuntime EP 使用此路径来获取文件夹路径，并与 ep_cache_context（指向上下文二进制路径）一起获取上下文二进制文件的绝对路径。
ep.context_embed_mode 仅用于上下文模型生成。
1：将 EP 上下文内容转储到 Onnx 模型中，位于 ep_cache_context 节点属性内。
0（默认值）：将 EP 上下文内容转储到单独的文件中，将文件名保留在 Onnx 模型中。文件路径在 ep_cache_context 节点属性中跟踪。
ep.context_node_name_prefix 仅用于上下文模型生成。
指定 EPContext 节点名称（也是 partition_name 属性，内部图形名称）前缀，使其在节点之间是唯一的，以防用户在一个模型中粘合多个 EPContext 节点以避免冲突。
ep.context_model_external_initializers_file_name 这适用于某些节点在 CPU EP 上分区，并且这些节点具有外部初始化器的情况。生成 EP 上下文模型时，新生成的模型不应依赖于源 Onnx 模型使用的旧外部数据文件。
当使用外部初始化器文件转储 EP 上下文模型时，请使用此配置。如果指定，所有初始化器都将位于外部数据文件中，否则都将位于生成的 Onnx 文件中。
默认情况下未设置，因此所有初始化器都将位于 Onnx 文件中。

EP 上下文缓存模型生成工作流程

OnnxRuntime EP 应遵循这些规则来创建 EP 上下文缓存模型，以维护统一的用户界面。

  • ep.context_enable
    • 如果 ep.context_enable = 1，则 OnnxRuntime 创建 EP 上下文缓存模型。否则，ep.context_enable = 0（默认值），只需执行正常工作流程。
  • ep.context_file_path
    • 如果没有提供 ep.context_file_path，OnnxRuntime 只需将原始输入文件名中的“.onnx”替换为“_ctx.onnx”作为输出文件名。否则，只需使用用户提供的文件路径。
    • 如果用户从内存缓冲区加载模型，则需要 ep.context_file_path，因为在这种情况下 OnnxRuntime 无法获取输入文件路径。
  • ep.context_embed_mode
    • 1（默认值）：将 EP 上下文内容转储到 Onnx 模型中。
    • 0：将 EP 上下文内容转储为单独的文件。EP 决定文件名，并在 EPContext 节点属性 ep_cache_context 中跟踪文件名。单独的文件应始终与转储的 Onnx 模型文件位于同一位置。并且在 EPContext 节点中跟踪的文件路径是 Onnx 模型文件的相对路径。注意：允许子文件夹。
  • ep.context_node_name_prefix
    • 如果用户想要在 EPContext 节点名称（也是 partition_name 属性和图形名称）内添加特殊标签，则 EP 应在 EP 创建 EPContext 节点时提供此功能。
    • 如果用户想要将来自多个模型的多个 EPContext 节点粘合到一个模型中，并且存在跨模型发生节点名称（图形名称）冲突的风险，这将非常有用。取决于 EP 实现。QNN EP 支持多个 EPContext 节点，因此用户可以合并和重新连接来自不同模型的 EPContext 节点。

从 EP 上下文缓存模型推理工作流程

支持从具有 EPContext 节点的 Onnx 模型加载的 OnnxRuntime EP 应遵循模型推理的工作流程/规则。

  • EP 应该能够识别具有 EPContext 节点的模型。
    • 如果模型内部没有 EPContext 节点，则 EP 遵循其正常工作流程。
    • 如果它是具有 EPContext 节点的 Onnx 模型。
      • EP 应检查所有 EPContext 节点的源节点属性，以确保此 EP 存在任何 EPContext 节点（源节点属性与 EP 要求的键匹配）。
      • EP 仅在源节点属性与 EP 要求的键匹配的 EPContext 节点中进行分区。
      • EP 从 EPContext 节点内部的缓存上下文中加载
  • 如果上下文缓存 Onnx 模型以 embed_mode = 1 转储，则在同一文件夹中的 Onnx 模型旁边有单独的上下文二进制文件。
    • OnnxRuntime EP 从 EPContext ep_cache_context 节点属性中获取上下文二进制文件相对路径。
    • 如果用户从 Onnx 模型文件路径加载模型，则 EP 应获取输入模型文件夹路径，并将其与从步骤 a) 获取的相对路径组合，作为上下文二进制文件的完整路径。
    • 如果用户从内存缓冲区加载模型，则用户需要提供会话选项 ep.context_file_path。EP 从 ep.context_file_path 获取文件夹路径，并将其与从步骤 a) 获取的相对路径组合，作为上下文二进制文件的完整路径。

EP Context nodes with different EPs

ExecutionProvider 接口 GetEpContextNodes()，以帮助生成 EP 上下文缓存模型

执行提供程序很难在执行提供程序代码中生成分区图，因为执行提供程序没有整个分区图的良好视图。添加了新的 ExecutionProvider 接口 GetEpContextNodes() 以支持此功能。

  virtual const InlinedVector<const Node*> GetEpContextNodes() const {
    return InlinedVector<const Node*>();
  }

此 API 返回 EPContext 节点的指针数组。如果执行提供程序需要生成上下文缓存模型，则需要实现此接口。否则，请保留它。执行提供程序有责任创建具有其依赖项的 EPContext 节点（例如，如果不是 embed_mode，则为上下文二进制文件）。OnnxRuntime GraphPartitioner 使用此接口来获取 EPContext 节点并生成分区 Onnx 模型。 代码详情请点击此处