【推理引擎】ONNXRuntime 的架構設計

虔誠的樹發表於2022-03-29
架構

ONNXRuntime，深度學習領域的神經網路模型推理框架，從名字中可以看出它和 ONNX 的關係：以 ONNX 模型作為中間表達（IR）的執行時（Runtime）。

本文許多內容翻譯於官方文件：https://onnxruntime.ai/docs/reference/high-level-design.html ，並適當地新增一些自己的理解，由於對ONNXRuntime的認識還不夠深入，因此可能會存在一些錯誤，希望多多指正，深入交流。

特色：

  1. 在不同平臺上，最大限度地、自動地使用定製的加速器（accelerators）和執行時（runtimes）；
  2. 針對定製的加速器和執行時，提供良好的抽象和執行時（onnxruntime）來支援執行，這裡的抽象也被稱之為EP（Execution Provider，eg. CUDA、TensorRT、OpenVINO、ROCm等）。每個EP都各自定義自己的功能，比如記憶體分配、可以執行的單個的或融合的節點（注意：本文中所說的節點就是運算元，兩者等同；conv屬於單個的運算元，conv_bn_relu屬於融合的運算元），這些功能需要以標準的API形式暴露給 ONNXRuntime，以供其呼叫；
  3. ONNXRuntime並不要求每個EP都完全支援ONNX中定義的所有運算元，這也就意味著 ONNXRuntime 可能需要在異構環境中才能完整的執行完一個模型，這裡的異構環境是指涉及到多個計算硬體，比如CPU和GPU；
  4. 支援多種圖優化（Graph Optimization），主要分為兩類：
  • 全域性變換（Global transformations）：這種優化方式需要對整張計算圖進行分析並優化；在原始碼中，每種變換都繼承自 GraphTransformer 類；
  • 區域性變換（Local transformations）：這種優化方式相當於定義一些簡單的重寫規則（rewriting rules），比如消除一些沒有具體操作的圖節點（eg.推理階段的dropout節點）；與全域性變換不同，重寫規則一般只針對圖中的部分節點，也就是說需要先判斷圖中的節點是否滿足重寫條件，然後再決定是否實施變換；在原始碼中，每種重寫規則都繼承自 RewriteRule 類，但是最後會使用 GraphTransformer 的一個派生類 RuleBasedGraphTransformer ，將所有的 RewriteRule 類聚合起來。

從更高視野看ONNXRuntime系統結構

從這張圖中，我們可以看出ONNXRuntime的執行流程。

  1. ONNXRuntime 首先將 ONNX 模型轉變為 In-memory 形式；
  2. 針對這個模型執行一些與EP無關的優化；
  3. 根據設定的EP（可能會有多個），將整體計算圖分割成多個子圖；
  4. 每個子圖都被分配到一個相應的EP中，分配過程中要確保這個EP能夠執行該子圖；

由於很多EP都會對一些特定的運算元做特殊優化，因此在分割子圖時，ONNXRuntime希望充分利用這些EP的能力，但是仍然會存在一些運算元不能被EP執行，或者高效執行，這時就需要設定一個預設的EP進行兜底，這個角色往往由CPU承擔。

計算圖分割的策略：首先設定可用的EP，比如

ort_sess = ort.InferenceSession('onnx_model/resnet50.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

然後依照providers中設定的順序為每個EP充分分配其可以執行的子圖，為了確保每個子圖都被執行，一般會講CPU EP放置在最後。ONNXRuntime當前只支援同步的執行模式，並且由其控制整個計算圖的執行。

相關文章