1. 程式人生 > >開源背後 | 面對端側推理引擎的挑戰,阿里工程師如何應對?

開源背後 | 面對端側推理引擎的挑戰,阿里工程師如何應對?

阿里妹導讀:MNN(Mobile Neural Network)已於今年5月7日在 Github 上正式開源。淘寶無線開發專家——陳以鎏(離青)在 GMTC 全球大前端技術大會為大家分享了 MNN 開發、開源中的思考與總結,通過淘寶在移動 AI 上的實踐經驗,你將會了解移動 AI 的發展狀況和應用場景,以及通過端側推理引擎瞭解移動/ IoT 深度優化策略。

開源與背景

人工智慧從 2006 年開始,迎來了第三次浪潮。隨著 AlphaGo 在 2016 年、 2017 年先後戰勝李世石和柯潔,人工智慧徹底進入了公眾的視野。人工智慧大熱的背後,是大資料的積累,是深度學習的發展,也是裝置算力的提升。與此同時,深度學習的框架也在不斷演進 —— 從 Torch、Caffe 到 TensorFlow、PyTorch ,再到更面向移動端的 CoreML、NNAPI、NCNN、MACE 等。淘寶的深度學習推理引擎 MNN 也於 2019 年 5 月宣佈開源。

MNN 是一個輕量級的深度學習端側推理引擎,核心解決深度神經網路模型在端側推理執行問題,涵蓋深度神經網路模型的優化、轉換和推理。目前,MNN已經在手淘、手貓、優酷、聚划算、UC、飛豬、千牛等 20 多個 App 中使用,覆蓋直播、短視訊、搜尋推薦、商品影象搜尋、互動營銷、權益發放、安全風控等場景,每天穩定執行上億次。此外,菜鳥自提櫃等 IoT 裝置中也有應用。在 2018 年雙十一購物節中,MNN 在天貓晚會笑臉紅包、掃一掃、明星猜拳大戰等場景中使用。

MNN 專案從 2017 年開始啟動,在經歷一年多的開發迭代並通過了淘寶雙十一的考驗後,於 2018 年底啟動開源計劃,在歷時小半年的開源改造後,今年 5 月份正式在 Github 開源。

開源首先還是因為經歷過雙十一之後,我們覺得自己做好了準備,開源有助於我們鞭策自己,把 MNN 做的更好;另一方面,業界的開源方案,不論是 TensorFlow Lite 、 NCNN 還是 Mace ,都給了我們很好的輸入和借鑑,我們也希望藉著開源,將我們的思考和創新回饋社群。

下文將主要圍繞著 MNN ,來介紹淘寶在移動 AI 上的一些實踐經驗。

挑戰與應對

端側推理引擎面臨的挑戰中,碎片化是最為顯著的,這種碎片化是多層次、多維度的 。

  • 訓練框架上, Caffe 、 TensorFlow 、 PyTorch 、 MXNet 在訓練模型時都很常用;
  • 計算裝置上, CPU 、 GPU 已是主流, NPU 、 TPU 漸漸成為標配, DSP 、 FPGA 在 IoT上也很常見;
  • 運算元層面上,眾多引數會形成不同的組合,從而對應出不同的優化方式,輕量化和通用化需要取捨;

一款優秀的端側推理引擎,就需要在這樣碎片化的環境下,利用裝置有限的資源,儘可能發揮出裝置的效能。為此,也需要在轉換、排程、執行上加入相應的優化策略。下文,會就其中的部分展開說明。

轉換工具

模型優化

在模型優化中,MNN 引入了前端的概念來統一訓練框架。不同的前端負責載入不同訓練框架的模型,統一轉換為 MNN 的模型格式。對於最常用的訓練框架 TensorFlow 和 Caffe ,我們提供了獨立的前端;其他訓練框架,比如 MXNet ,則需要先將模型轉換為 ONNX ,再通過 ONNX 前端載入。這裡,由於 TensorFlow 的運算元顆粒度相比 Caffe 和 ONNX 要更小,我們引入了圖優化的模組來對齊運算元之間的顆粒度。模型轉換之後,會經過優化器優化,包含運算元融合、運算元替換、佈局調整等等。之後,可以選擇對浮點模型執行量化壓縮。目前模型壓縮的模組還沒有開源,我們會在完善之後,將相關程式碼開源。這些步驟都完成之後,會使用 flatbuffer 來儲存部署模型。

圖優化

這裡以 RNN-GRU cell 為例,說明一下圖優化。

左圖是 RNN-GRU cell 在 TensorBoard 中的視覺化描述。它足足包含了 3584 個節點,而每一個節點都代表了一定的資料讀寫或運算,累積起來的總量非常大。然而,所有這些節點可以打包使用一個大顆粒的運算元來替代。這不僅大幅降低了部署模型的大小,還可以在大顆粒運算元的實現中聚合大量的計算,避免不必要的資料讀寫。

右圖展示的是一個實際業務模型在圖優化前後的效能對比。在華為 P10 、紅米 3x 、小米 6 上都有 1 倍左右的效能提升。而如果是雙向 GRU ,效果還會更明顯。

運算元融合

再以 Convolution、Batchnorm、Scale、ReLU 為例說明優化器中的運算元融合。

首先融合 Convolution 和 Batchnorm,Convolution 的 weight 等於 weight 乘 alpha ,而 bias 等於 bias 乘 alpha 再加 beta ;而後融合 Convolution 和 Scale ,融合過程和 Batchnorm 類似;最後融合 Convolution 和 ReLU ,在輸出結果前,計算啟用函式 ReLU 即可。

這樣,四個運算元就可以合併成一個運算元。融合的過程避免了三次 tensor 讀寫、兩次 tensor 乘加。優化效果:MobileNet V1 在小米 5 和華為 P10 上有 20 ~ 40% 的效能提升,效果還是比較明顯的。

智慧排程

整體設計

在排程上, MNN 將每一類計算裝置抽象為一個後端,將運算元在特定後端上的實現抽象為執行器。後端負責特定裝置上的資源分配和計算排程,執行器負責具體的實現。後端和運算元的新增都通過登錄檔來實現,這是一個雙層登錄檔結構,拓展起來就相對靈活。

排程時,可以為子圖選擇相應的後端,再由後端創建出相應的執行器,組成管線;也可以為子圖選擇後端組,實現混合排程。比如,在 GPU 上不宜實現排序運算元時,可以回退到 CPU 來執行。

目前, MNN 在 CPU 上實現了 76 個運算元, Metal 上有 55 個, OpenGL 覆蓋了基礎的 CNN 網路, OpenCL 和 Vulkan 分別有 29 和 31 個。

快取管理

在建立完執行器之後,子圖和管線已經就緒。下來,需要計算出所有 tensor 的形狀,在相應的後端上完成記憶體的分配。而後,在準備執行器時,再為所有的執行器預先在後端上申請好必要的 buffer。執行結束後,返回 tensor 即可。

由於推理所需的所有記憶體在準備期就已經申請完畢,在後續推理時,如果輸入的形狀不變,就可以複用 tensor 和 buffer,從而避免頻繁地申請、釋放記憶體;只有輸入形狀改變的時候,才需要從形狀計算開始,調整一次記憶體分配。同時,由於使用後端統一管理快取,後端內的執行器之間,快取就可以充分複用的,這就大大減少了記憶體的需求量。此外,MNN 分配記憶體時,預設按照32位對齊,記憶體對齊有利於資料讀寫。

執行優化

資料佈局與滑窗卷積

資料佈局對效能影響巨大。

先來看一看在 NCHW 的佈局下,怎麼利用 SIMD 加速 3x3 的 depth-wise 卷積。

首先,讀取資料時,需要一次性讀取四個 float 作為第一行的資料,後兩行的讀取也是相似的;此時,讀取出的三行資料已經足夠計算兩列輸出,即,可以複用部分資料;而後,為了提高資料複用,會再讀取出第四行資料,一次計算兩行兩列,即,可以引入迴圈展開;然而,殘留的 5~25 和 21~25 亮度眼邊界無法利用 SIMD 計算,只能逐一迴圈讀寫完成計算;按照這樣的方式,就可以相應完成後幾個通道的計算。

但是, NCHW 佈局下,無法充分利用 SIMD 進行加速,同時,實現優化分支越多,佔用包大小也就越多。

再來看一看 NC/4HW4 佈局下,利用 SIMD 加速的情況又是怎樣的。

這裡的 "C/4" 指的是按照 4 個通道對齊的方式重排資料。重排所有輸入和權重資料後,每次 SIMD 讀寫都天然是 4 個通道的輸入資料和 4 個通道的權重資料。這樣,不論 kernel、stride、dilation 怎麼變化,我們都可以簡單地使用 for 迴圈和 SIMD 的一套通用優化完成卷積計算。既不會有邊緣資料無法加速的問題,也不會對包大小造成影響。

Winograd

對於對於 KxK 卷積,可以使用 Winograd 演算法進一步加速。MNN 中支援 2x2 到 7x7 的 Winograd 實現。Winograd 計算時,需要把輸出拆分成 NxN 的小塊,把輸入拆分成 (N+K-1)x(N+K-1) 的小塊。這樣,問題就可以簡化為兩個小矩陣的卷積。

再套用 Winograd 的公式,將矩陣間的卷積運算轉換為矩陣點乘運算。在這個過程中,除了矩陣點乘外,還引入三個矩陣轉換,分別是輸入矩陣 d 、權重矩陣 g 和結果矩陣 Y’ 的轉換。其中,權重轉換時, G 矩陣可以利用中國剩餘數定理計算, GgGT 就可以在準備執行器時提前計算;輸入轉換和輸出轉換時用到的 A 和 B 矩陣需要根據 N 和 K 計算,我們在程式碼中內建了幾種優化後的組合,所以實際計算時,這兩個轉換並不需要經過複雜的矩陣乘法。

這樣,原來矩陣卷積所需要的 9x4 次乘法計算就可以用矩陣點乘的 4x4 次乘法計算代替。只考慮乘法耗時的話,加速了 2.25 倍。示例中, K=3,N=2 ,但實際使用時,可以選擇更大的 N 值,獲取高的加速倍數,但也要相應消耗更多的記憶體。

Strassen

MNN 可能是端側推理引擎中,第一個應用 Strassen 演算法優化矩陣乘法的。

Strassen 在計算矩陣乘法時,首先需要將矩陣平均拆分成四個小矩陣。這裡使用 a11 ~ a22、b11 ~ b22、c11 ~ c22 代表四個小矩陣,計算過程一共需要8次小矩陣乘法運算。

這裡可以引入中間小矩陣, s1 ~ s4、t1 ~ t4、m1 ~ m7、u1 ~ u7 。其中,只有 m1 ~ m7 包含小矩陣乘法,一共 7 次小矩陣乘法運算。而其他的,只包含小矩陣的加減法。也就是說,通過 4 + 4 + 7 次小矩陣加減法,替代了一次小矩陣乘法。

與原來的矩陣乘法相比, Strassen 的時間複雜度從 n 的 3 次方,降低到 n 的 2.81 次方。在矩陣較大時,矩陣乘法遠遠慢於矩陣加減法,收益就更明顯。

在 MNN 中,我們會遞迴使用 Strassen 。也就是說,遞迴拆分矩陣。在矩陣足夠大時,繼續拆分;在矩陣不夠大時,使用普通的矩陣演算法。這裡使用減免的矩陣乘法開銷作為收益,使用小矩陣 s 、小矩陣 t 、小矩陣 u 矩陣的加減法開銷之和作為代價,收益大於代價時,就可以考慮使用 Strassen 演算法。

鏈路優化

鏈路優化可以舉一個 19 年春節淘寶掃年貨的例子。在獲得手機相機輸入後,每一幀的影象首先需要經過一次預處理,將圖片縮放到年貨檢測模型的輸入大小上,然而再經過推理,判定影象有沒有年貨,如果有,就發放相關權益。這個過程中,圖片預處理的耗時也不容忽視。降低這個耗時,就可以幫助我們提升幀率,從而改進使用者體驗。為此,我們引入了一個輕量級的 2D 圖片處理庫,可以高效地完成色值變化、色彩空間的轉換或者仿射變換等。這樣, MNN 的使用者就不再需要為圖片處理引入 libyuv 或者 opencv 了。

效能比較

經過種種優化後,這是我們在效能上交出的答卷。

MobileNet V2 ,在 OPPO r17 和 iPhone 7Plus 上做了一系列的效能對比。

如圖, MNN 的效能在 CPU 和 GPU 上都有一定的優勢。

小結

總的來說, MNN 吸納了前人的經驗,也結合自己對端側推理引擎的認知,做了許多創新。綜合考慮效能、模型和後端的拓展性、快取、 CPU 和 GPU 的運算元實現,以及 CV 庫等方面的表現,在端側推理引擎中, MNN 是值得移動 AI 使用者嘗試的選擇。

後續規劃

在後續計劃上,轉換部分,我們計劃追加更多運算元和更多圖優化匹配模板,也計劃將模型量化工具開源;排程部分,我們計劃分步實現端側訓練和邊緣學習,計算裝置自動選擇也在籌劃中;執行部分,還是會持續優化現有各端運算元的實現,也計劃優化量化卷積、矩陣乘演算法,計劃在 CV 庫上直接支援 GPU ,我們也考慮將現有 NC/4HW4 實現的演算法,整理為獨立的高效能運算庫,演算法自動選擇同樣在籌劃中;其他部分,我們會持續建設專案的可用性,持續加入更多的文件和示例。

淘寶基礎平臺部-端智慧團隊歡迎推理引擎研發工程師、AR技術研發工程師、高效能運算研發工程師的加入。對新技術感興趣,善於創新突破,渴望用新技術給使用者帶來創新體驗的同學請聯絡我們,簡歷投遞至[email protected]

作者:陳以鎏(離青)

原文連結

本文為雲棲社群原創內容,未經