ONNX Runtime 源碼閱讀:引擎運行過程總覽

推薦閱讀
ONNX Runtime 源碼閱讀:模型推理過程概覽
ONNX Runtime 源碼閱讀:模型結點串行執行順序的確定
ONNX Runtime 源碼閱讀:模型分區的實現

之前的文章中,我們一直從源碼細節的角度去看某個功能的實現。總會有一種一葉障目不見泰山的感覺,對ONNX Runtime缺乏一個比較總體的認識。通過前面的幾篇文章,在看細節源碼的過程中,腦海漸漸浮現出了ONNX Runtime總體的輪廓,但是並不是很清晰。今天打算暫時跳出細節,站在高處,看看ONNX Runtime它的筋骨、它的脈絡,也更能方便以後進行局部分析。
下圖展現的是ONNX Runtime主要的一些類以及它們之間的關係,這裏只展現了比較重要的一些類而不是全部。通過UML圖可以看到,InferenceSession, SessionState這兩個類佔據着舉足輕重的地位。如果要打個比方,那麼InferenceSession就是一個運籌帷幄之中的統帥,負責統籌全局、制定計劃;而SessionState就相當於保障物資的後勤保障系統。所謂兵馬未動糧草先行,在小兵們(OpKernel)衝鋒陷陣(推理)之前,統帥一定要準備好各種糧草輜重、武器裝備並制定戰術,不僅讓士兵喫飽、裝備精良,還要儘量發揮各個兵種的優勢,該遠程消耗就讓弓箭手上、該攻城就讓步兵來,不然讓騎兵去打水戰註定是悲劇。
下面就讓我們站在上帝的視角,看一看一場戰鬥是怎麼打贏的。
overview.jpg

之前多次說過,模型推理的過程可以簡單分爲三步:實例化、初始化、推理。今天就這上面這張圖,分析一下整個推理過程。

實例化

實例化的時候,InferenceSession會構建起一個框架。當實例化結束,InferenceSession已經持有SessionState, KernelRegistryManager, ExecutionProviders等的引用,只不過他們也都只是個空架子,裏面並沒有什麼實質性內容。緊接着,就進入了初始化階段。
初始化階段是模型推理的重頭戲,它最終讓一個靜靜躺在磁盤中的模型文件中的描述內容再內存中一一建立了起來。

初始化

模型加載

首先,InferenceSession將模型文件從磁盤讀入並解析。如果模型內容已經存在內存之中,就直接使用這塊內存;如果模型還不在內存,就通過一個Load()方法去加載。值得注意的是,雖然解析到的模型可以通過一個Protobuffer編譯器編譯出來的model->MainGraph可以直接訪問到模型的內容,但是ONNXRuntime還是在這之上添加了一層抽象層,這個抽象層就是按照ONNX標準來的。主要目的就是使得整個引擎和模型文件的格式隔離開,引擎最終與模型文件個格式解耦,也就是和模型的解析過程解耦。雖然現在ONNX用的是Protobuffer,但是有了則層抽象層,後面如果換成了FlatBuffer或者其他的格式,也不會對整個框架有影響。並且,最終Load()方法會委託用戶自定義的一個loader去執行具體將模型加載到內存的操作,這就讓從哪裏加載模型獲得了極大的自由。突然想到Java的類加載器,只要你最終加載到的模型符合標準,並不管你是從本地磁盤還是網絡獲取到的。
模型加載並解析之後,得到的是一個原始的模型在內存中的表示。原始的意思就是,整個數據結構是直接通過Protobuffer編譯器編譯得到的,後續還需要對她進行一層封裝處理。都說巧婦難爲無米之炊,現在米已經有了,就看怎麼做好這頓飯。

註冊Provider以及Kernel

在實例化階段,已經生成了KernelRegistryManager的實例。此時,InferenceSession將會實例化所有可用Provider,並將他們保存在ExecutionProviders當中。之後,KernelRegistryManager會通過ExecutionProviders拿到InferenceSession傳遞過來的Provider,並委託他們提供自己所支持的所有OpKernel的信息。這些信息並不是OpKernel的實例,而是KernelRegisty的實例。KernelRegisty總持有生成每一OpKernel的實例所需要的信息和方法,這些信息和方法通過一個結構體KernelCreateInfo保存着。當這一步完成,有關鍋碗瓢盆也準備停當。

模型分區

InferenceSession委託GraphPartitioner去做模型分區工作。GraphPartitioner初始化的時候InferenceSession會給到它ExecutionProvidersKernelRegistryManager。具體分區過程在ONNX Runtime 源碼閱讀:模型分區的實現中有介紹。

模型節點執行順序的確定

模型加載並解析之後,所的到的知識原始的模型結構,ONNX Runtime需要對它按照ONNX的標準進行轉換和封裝,包括將所有原始節點編號、確定執行順序等,然後存入Graph當中,最後Graph被存放於SessionState當中。具體執行順序的確定的過程在ONNX Runtime 源碼閱讀:模型結點串行執行順序的確定中有介紹。

實例化Kernel

Kernel的實例化過程中,InferenceSessionKernelRegistryManager傳遞給了SessionState.CreateKernels()。反過來,SessionState.CreateKernels()拿到KernelRegistryManager後委託它通過來實例化。最終通過層層委託,最終是KernelCreateInfoKernelCreateFn根據KernelDef最終生成了OpKernel的實例。這部分就留着以後細看吧。

推理

推理,也就是依次調用OpKernel.Compute(),利用面向對象的多態機制,此時基類OpKernel的引用綁定的是每一個具體的子類,因此會調用到子類重寫的Compute()方法。

另外還有一大塊內容:運行過程中的內存分配,我們從來沒有觸碰過。可能對於運行在服務器端的引擎還沒那麼敏感,但是內存對於移動設備那可真是寸土寸金。這些都值得一看,留坑後續慢慢填了。

公衆號二維碼

首發於個人微信公衆號TensorBoy。微信掃描上方二維碼或者微信搜索TensorBoy並關注,及時獲取更多最新文章!
C++ | Python | 推理引擎 | AI框架源碼,有一起玩耍的麼?

Resources

https://github.com/Microsoft/onnxruntime
https://github.com/zmychou/onnx-runtime-code-reading

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章