ONNX Runtime 架构
本文档概述了 ONNX Runtime 的高级设计。
目录
关键目标
- 最大限度地自动利用不同平台上可用的自定义加速器和运行时。
- 为自定义加速器和运行时提供正确的抽象和运行时支持。我们将这种抽象称为执行提供程序。它定义并向 ONNX Runtime 暴露其一组能力：它可以执行的单个或融合节点集、其内存分配器等。自定义加速器和运行时都是执行提供程序的实例。
- 我们不期望执行提供程序总是能够在其设备上完全运行 ONNX 模型。这意味着 ONNX Runtime 必须能够在涉及多个执行提供程序的异构环境中执行单个模型。
- 通过图转换 API 提供对可表达为模型到模型转换的高级优化的支持。此类转换分为两类：全局转换，需要分析和转换整个图；局部转换，可捕获为简单的（代数）重写规则。
高级系统架构
流程非常简单。
- 从 ONNX 模型开始，ONNX Runtime 首先将模型图转换为其内存中的图表示。
- 它执行一系列与提供程序无关的优化。
- 它根据可用的执行提供程序将图划分为一组子图。
- 每个子图都被分配给一个执行提供程序。我们通过使用
GetCapability()API 查询执行提供程序的能力来确保子图可以由执行提供程序执行。
关于分区
ONNX Runtime 根据可用的执行提供程序将模型图划分为子图，每个不同的提供程序对应一个子图。ONNX Runtime 提供了一个默认执行提供程序，用作无法推送到更专业但更高效的执行提供程序的算子的回退执行。直观上，我们希望尽可能将计算推送到更专业的执行提供程序。
我们使用一种简单的图分区技术。将按照特定顺序考虑可用的执行提供程序，每个提供程序将被分配它可以处理的最大子图（可能不止一个）。ONNX Runtime 提供的默认执行提供程序将是最后一个考虑的，它确保了完整性。未来可以考虑更复杂的优化（甚至可以实现为复合执行提供程序）。
概念上，每个分区都被简化为一个单一的融合算子。它是通过调用执行提供程序的 Compile() 方法创建的，并将其包装为自定义算子。目前我们只支持同步执行模式。执行提供程序公开其内存分配器，用于为执行提供程序分配输入张量。重写和分区将初始模型图转换为一个新图，该图由分配给默认执行提供程序或其他注册执行提供程序的算子组成。ONNX Runtime 执行引擎负责运行此图。
关键设计决策
- 多个线程可以在同一个推理会话对象上调用
Run()方法。更多详细信息请参见API 文档。 - 为了方便实现这一点，所有内核的
Compute()函数都是 const，这意味着内核是无状态的。 - 执行提供程序实现的算子称为内核。每个执行提供程序支持 ONNX 算子/内核的一个子集。
- ONNX Runtime 保证所有算子都由默认执行提供程序支持。
- 张量表示：ONNX Runtime 对张量运行时值使用标准表示。执行提供程序如果选择，可以在内部使用不同的表示，但它们有责任在其子图边界处将值从/转换到标准表示。