【推理引擎】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當前只支援同步的執行模式,並且由其控制整個計算圖的執行。

相關文章