深度學習模型編譯技術

前言

深度學習模型的開發周期，包括訓練階段和部署階段。訓練階段，用戶需要收集訓練數據，定義自己的模型結構，在CPU或者GPU硬件上進行訓練，這個過程反復優化，直到訓練出滿意精度的模型。有了模型之后，我們需要將模型服務部署運行，我們期望服務延遲越低越好，吞吐越高越好。這里會從編譯優化的角度來分析一下如何進行深度學習模型推理優化。

1. 推理優化現狀

1.1 原生框架部署 - TF/PT

模型使用原生框架進行推理非常便捷，一般不需要進行太多環境適配和開發工作就能將模型部署起來，對于開發者非常友好。但是訓練框架（目前主流的TensorFlow/PyTorch/MindSpore等）為了訓練開發的便利，存在大量瑣碎的算子，使得在batch比較小的時候，推理的性能不是那么高效。

為了提高模型推理的性能，業界近幾年開始不斷研究模型編譯優化技術，旨在不同硬件不同場景中自動的優化推理性能。目前大體可以分成兩個大的方向：推理引擎和推理編譯器。兩者有很多相通的地方，也有各自獨特的優勢。

1.2 推理引擎 - Library

在人工智能發展的浪潮中，許多公司推出了各家代表的推理優化解決方案，例如NVIDIA基于自家顯卡的TensorRT，Intel基于自家處理器的OpenVINO，Huawei基于昇騰芯片的MindSporeLite等等。這些推理引擎往往以library的方式供用戶使用，將模型的IR轉化為引擎內部的IR表示，然后映射綁定到對應硬件的算子實現，最后使用引擎的runtime來啟動推理計算。基于自家的處理器和高度定制的算子實現，推理引擎能夠達到非常極致的優化性能。但是推理引擎的架構決定了不能很好的覆蓋所有原生框架的算子（模型轉化+算子實現覆蓋度），在一部分模型推理優化中需要一定的算子開發量，往往也需要有經驗的優化工程師的參與才能達到比較好的性能。

為了解決推理引擎算子覆蓋度和不同硬件適配的問題，大家逐漸開始了對算子自動代碼生成的研究，即推理編譯器范疇。

1.3 推理編譯器 - Compiler

推理編譯器主要是為了解決幾個問題：1. 算子代碼自動生成；2. IR與不同硬件的適配。

一些常用的ElementWise/Reduction等算子，往往可以通過模板或者一些IR的定義來生成不同硬件的代碼，省去了手寫算子的人力投入。為了適配不同的硬件，往往編譯器也會將IR表示進行分層，大體分為High Level IR和Low Level IR，其中High Level IR表示與硬件解耦，盡可能的表達圖和算子的特性，編譯器可以基于這一層的IR進行與硬件無關的圖優化；Low Level IR則設計和硬件相關，用于針對不同硬件進行圖優化和代碼生成。

1.4 寫在最后

推理引擎和推理編譯器沒有優劣之分，兩者有各自的優勢和適合的場景。兩者很多技術是共同的，例如圖融合優化；兩者也在不斷的相互借鑒和融合，例如TensorRT內部的Meylin編譯器也能夠進行算子代碼自動生成。后文中編譯優化囊括了推理引擎和推理編譯器兩種優化。

2. 編譯優化

編譯優化技術一般可以分為前端優化和后端優化，分別對應High Level IR和Low Level IR，因此前端優化大都與硬件解耦，而后端優化與具體的硬件緊密相關。接下來，會分別介紹一下編譯優化前端和后端有哪些優化方式。

2.1 編譯前端

2.1.1 Algebraic Simplification

模型其實是算子的DAG圖，算子與算子之間往往可以通過算數等價的方式來進行計算量簡化，達到優化推理性能的目的。常見的算術簡化包括如下幾種：

algebraic identification

這里使用layer norm中均值和方差計算進行舉例。通過最后一行簡化的表示，可以給予硬件同時計算x和x^2均值的可能，提高硬件的利用率。而原始的表示決定了要先計算x的均值后，再計算方差。由于兩者是算術等價的，因此選擇最佳的算子圖表示。

strength reduction

在元素計算中，往往可以通過更高效的算子代替比較慢的算子計算，提高算子計算效率。如移位運算代替乘法或者除法操作。

computation order

算子的運算順序可能能夠省去不必要的轉化操作，如圖中的矩陣乘法等價變換。

constant folding

當算子輸入都為常量tensor的時候，可以離線計算出算子的輸出，消減掉該算子，提高運行時推理效率。

2.1.2 Dataflow-level optimizations

除了算子和算術等價的優化外，還可以基于圖的結構進行推理優化。例如通過前后算子的內存復用來節省內存占用；通過轉化Tensor的組織結構（NCHW）來獲得更高效的算子實現；取消推理時無效的算子；更改算子前后位置創造融合機會等。

2.1.3 Operation Fusion

算子融合是非常常用的圖優化技術。算子融合可以減少不必要的線性計算（例如BN可以提前把scale/var信息更新到conv的weight/bias中），同時可以減少片外帶寬存儲空間的反復讀寫，盡可能的將計算在片上計算完成，也減少了kernel的啟動開銷。

2.2 編譯后端

2.2.1 Hardware Specific Opt

針對特定的硬件，往往可以做計算和訪存的定制優化，例如：

原語或指令的映射：將算子拆分到硬件的原語級別，將拆分后的計算映射到對應的原語，硬件的定制的原語具有非常高的計算效率。

訪存優化：通過多級存儲空間，將DRAM的數據讀取到片上寄存器或者共享內存進行計算操作，并通過并發將Load/Store指令與計算指令進行隱藏，提高吞吐。

循環展開：通過將循環展開，可以給予編譯器更多指令流水線優化的可能，提高計算效率。

向量化優化：很多硬件都有SIMD或者SIMT結構，可以將并行展開，使用向量化指令提高并發效率。

2.2.2 Lowering

在后端最重要的是lowering設計，即將高層次的表示映射到低層次硬件相關的表示。如果硬件廠商提供了一些底層高性能算子庫實現，那么可以直接將算子表示映射到library上面，畢竟手動實現一個高效的gemm/conv還是很有難度的。另外一些ElementWise/Reduction/Transpose等操作，可以通過模板自動生成對應的代碼，不用反復實現。最終編譯器會有Just-In-Time (JIT) 和Ahead-Of-Time (AOT) 兩種方式來輸出編譯優化后的結構。JIT是在運行時進行優化，AOT是提前做好優化（生成so或者binary）。

3. 優化無止境

高性能并行計算領域有一個經典的定律叫阿姆達爾定律，阿姆達爾定律這一樸素的定律可以幫助我們評估出優化的上限；并指導我們要針對瓶頸進行優化。針對瓶頸進行優化才能事半功倍。優化的流程包括了應用瓶頸分析，針對瓶頸進行并行化加速，優化實現，部署。這個循環可以往復，當一個瓶頸得到優化后，又會有新的瓶頸，直到優化滿足我們的需求才結束。深度學習推理優化也如此，一個應用可能瓶頸在邏輯的處理，也可能在模型的計算，在優化前需要對整體進行分析后再開始針對的優化。

推理優化技術也在朝著自動化、低精度、多硬件方向發展中，推動了人工智能深度學習技術在我們的日常生活中的普及，希望人工智能越來越智能和低成本，早日實現科幻電影中的未來世界。

深度學習 高性能計算

版權聲明：本文內容由網絡用戶投稿，版權歸原作者所有，本站不擁有其著作權，亦不承擔相應法律責任。如果您發現本站中有涉嫌抄襲或描述失實的內容，請聯系我們jiasou666@gmail.com 處理，核實后本網站將在24小時內刪除侵權內容。