MNN 引擎的稀疏計(jì)算方案設(shè)計(jì)與實(shí)踐
推理引擎領(lǐng)域,經(jīng)過最近幾年的打磨優(yōu)化,阿里推出的MNN(Mobile Neural Network) [1][2][3] 也成為業(yè)內(nèi)領(lǐng)先的推理引擎,當(dāng)我們想進(jìn)一步提升性能時(shí),結(jié)合深度學(xué)習(xí)模型設(shè)計(jì)入手是一個(gè)有潛力的方向,結(jié)合科學(xué)計(jì)算、高性能計(jì)算領(lǐng)域的知識(shí)則是一個(gè)更具體的方法,基于這個(gè)思考路徑, 業(yè)內(nèi)逐漸開始從稀疏計(jì)算角度來提升推理引擎性能。稀疏是指數(shù)據(jù)矩陣中有部分是0元素,一般而言,矩陣稀疏化之后,非0元素?cái)?shù)據(jù)在內(nèi)存中不連續(xù),無法直接復(fù)用GEMM(general matrix multiply) [9] 算法,同時(shí)緩存命中問題或判斷開銷會(huì)降低稀疏計(jì)算性能,需要新的方法實(shí)現(xiàn)稀疏相對稠密計(jì)算的加速。
結(jié)合業(yè)內(nèi)已有的稀疏計(jì)算方法與MNN的推理框架、內(nèi)存布局、算子關(guān)系,我們在MNN框架中設(shè)計(jì)支撐多類后端的稀疏整體方案,打磨稀疏計(jì)算核心匯編實(shí)現(xiàn)達(dá)到高性能;重要性能指標(biāo)數(shù)據(jù)如下:以AI(Artificial Intelligence)領(lǐng)域經(jīng)典的分類模型MobileNetV1 [12] MobileNetV2 [13] 為例, 選取相同高通SD 835 CPU(MSM8998),與XNNPack [6] 相比,MobileNetV1 90%稀疏度時(shí)XNNPack 加速比為2.35x,MNN加速比為 2.57x - 3.96x。MobileNetV2 85% 稀疏度時(shí), XNNPack MobileNetV2(只有 block=2),加速比為1.62x,MNN加速比為 3.71x(block=4),block含義與其他模型信息見1.1、 3.1節(jié)解讀,其他信息在后續(xù)展開。
1. 稀疏布局與加速原理
1.1 可調(diào)分塊稀疏權(quán)重
深度神經(jīng)網(wǎng)絡(luò)模型的稀疏計(jì)算包括輸入稀疏 [11] 、輸出稀疏、權(quán)重稀疏,0元素占全部元素的比例就是稀疏度,通常而言輸入和輸出稀疏對具備特定屬性模型會(huì)有效果,例如激活函數(shù)為ReLU,計(jì)算出負(fù)數(shù)值會(huì)被歸零,所以可以不精確計(jì)算出這些結(jié)果,只計(jì)算符號(hào)實(shí)現(xiàn)加速。而在CPU對通用DNN模型具備加速效果的是權(quán)重稀疏,在MNN中我們聚焦權(quán)重稀疏這個(gè)通用場景,另一個(gè)重要特點(diǎn)是面向深度學(xué)習(xí)的權(quán)重稀疏比例一般在30%-90%,有別于科學(xué)計(jì)算中極度稀疏到99%的場景。即以下三點(diǎn):
MNN稀疏計(jì)算的基本選型:
- 左右矩陣方面:選擇通用權(quán)重稀疏;
 - 稀疏度方面: 在30%-90%左右需產(chǎn)生加速效果,非科學(xué)計(jì)算的極度稀疏場景;
 - 結(jié)構(gòu)方面:不采用裁剪權(quán)重矩陣通道維度的完全結(jié)構(gòu)化剪枝方法,而是隨機(jī)稀疏+分塊稀疏。
 
常規(guī)結(jié)構(gòu)化剪枝中 [4][5] ,直接刪減權(quán)重矩陣的IC或OC維度的整個(gè)維度切片,裁剪后模型仍然是稠密的,直接縮減矩陣乘法規(guī)模,由于對宏觀結(jié)構(gòu)改變較大,導(dǎo)致模型精度損失顯著。
MNN中的稀疏計(jì)算選擇對權(quán)重在OC維度做動(dòng)態(tài)分塊稀疏,下圖1為隨機(jī)稀疏分塊下矩陣乘法的左右矩陣數(shù)據(jù)布局,涂色方格表示數(shù)據(jù)非0, 白色方格表示稀疏化處理后為0,即可不存儲(chǔ),又可不計(jì)算。

下圖2為分塊稀疏的矩陣乘法的左右矩陣數(shù)據(jù)布局,左矩陣與隨機(jī)稀疏情形相同, 右矩陣為權(quán)重矩陣,以每OCBlock為單位,連續(xù)為0或非0分布于內(nèi)存中。為了靈活性考慮,我們設(shè)計(jì)支持OCBlock為可調(diào)數(shù)值。

1.2 權(quán)重矩陣數(shù)據(jù)壓縮格式
經(jīng)過對比選型,在MNN中權(quán)重矩陣采用了多種布局形式,原始權(quán)重如圖3所示,為了保障性能的同時(shí)壓縮內(nèi)存占用。性能隨機(jī)稀疏時(shí),選用圖4所示的布局樣式;分塊稀疏時(shí),選用圖5布局。原始矩陣壓縮為非0數(shù)據(jù)與索引部分。下圖4、5中,權(quán)重矩陣A的0元素被壓縮后不再存儲(chǔ),data為非0數(shù)據(jù),Row start Index 與Column Index為行索引與列索引。圖5的布局時(shí),可以節(jié)省更多索引空間,壓縮內(nèi)存占用。



2. MNN稀疏計(jì)算方案設(shè)計(jì)
2.1 推理框架層設(shè)計(jì)
在MNN已有架構(gòu)基礎(chǔ)上的設(shè)計(jì)稀疏計(jì)算,則設(shè)計(jì)與實(shí)現(xiàn)需要考慮較多現(xiàn)狀的約束,結(jié)合基礎(chǔ)軟件的定位,設(shè)計(jì)目標(biāo)主要包括以下5個(gè)方面。

說明:由于漢語時(shí)序的”前后“和英語”forward/backward”語義相反,故用了向舊兼容代替通常說的有歧義的"向前兼容"。
2.2 MNN稀疏計(jì)算架構(gòu)解析
經(jīng)過不斷設(shè)計(jì)與完善,MNN稀疏計(jì)算目前整體上包含稀疏訓(xùn)練、轉(zhuǎn)換參數(shù)、算子框架、后端kernel 四個(gè)階段,設(shè)計(jì)四部分松耦合;方便層內(nèi)擴(kuò)展,以及模塊化集成與被集成,如下圖package UML所示。為了便于理解,參考C4Model建模方法分層,圖6為container層架構(gòu)、圖7為具體一層的Component層架構(gòu)。

- 第一,屬于算法模型階段, 算法工程師結(jié)合數(shù)據(jù)搭建模型,根據(jù)自身偏好在各類框架下訓(xùn)練得到AI模型;
 - 第二,稀疏化訓(xùn)練階段,參考圖6,從float稠密權(quán)重模型開始,導(dǎo)入MNN Python壓縮工具包(mnncompress),設(shè)定mnncompress需要的參數(shù),運(yùn)行將原模型中權(quán)重部分?jǐn)?shù)值稀疏化為0。我們建議用戶使用MNNCompress工具里的訓(xùn)練插件,最大化發(fā)揮加速性能;
 - 第三,轉(zhuǎn)換模型階段,MNN convertor 主要包含三類模塊,MNN內(nèi)部buffer格式轉(zhuǎn)換、圖優(yōu)化、后處理,選擇算子,將權(quán)重矩陣統(tǒng)計(jì)處理,滿足稀疏閾值的,給添加稀疏后端識(shí)別和運(yùn)行需要的參數(shù),最終寫文件得到稀疏MNN模型;
 - 第四,MNN engine推理計(jì)算階段,部署新模型到MNN運(yùn)行環(huán)境中,和普通模型運(yùn)行一樣,MNN 運(yùn)行時(shí)會(huì)自行處理算子映射、后端microkernel選擇、執(zhí)行推理。參考一般模型部署運(yùn)行文檔。(https://www.yuque.com/mnn/cn/create_session)
 
在 MNN Engine中,參考圖7, 算子operators和后端backend做如下結(jié)構(gòu)設(shè)計(jì)與考量。
- 外部看來,算子注冊并未新增一種“稀疏卷積”算子,仍然是普通的卷積算子,這樣可以減小用戶使用的選擇成本,減少算子膨脹,在MNN內(nèi)部更靈活地選擇稀疏計(jì)算加速,或原始稠密卷積。
 - 算子層面,將原始的稠密卷積重組成兩層,合理分配可復(fù)用與需要擴(kuò)展的部分,實(shí)現(xiàn)稀疏計(jì)算壓縮各類操作。
 - 量化稀疏算子則基于量化卷積算子ConvInt8Tiled擴(kuò)充實(shí)現(xiàn),基本方法與2相近。
 - 算子中平臺(tái)有關(guān)的核心函數(shù),我們分別實(shí)現(xiàn)了ARM32 fp32,ARM32 int8, ARM64 fp32, ARM64 int8, x86 avx2 fp32, x86 avx512 fp32共6種后端的匯編代碼。
 - 測試類方面,稀疏的用例完全包含稠密卷積的用例,并增加稀疏分塊維度,遍歷不同分塊、不同稀疏度、不同后端的正確性。
 

3. 稀疏計(jì)算性能評(píng)估
3.1 典型模型稀疏加速評(píng)估
為全面評(píng)估fp32稀疏加速效果,我們從"稀疏度、分塊大小、cpu型號(hào)、模型類型"四個(gè)維度進(jìn)行評(píng)測,匯集在圖8展示結(jié)果, 單張曲線圖表示固定一個(gè)設(shè)備、一個(gè)模型時(shí),推理耗時(shí)隨稀疏度的變化曲線。
將大圖的第1行第2列取出單列置于上方,便于解讀數(shù)據(jù)。單張曲線圖包含四條曲線,"dense"為MNN當(dāng)前基準(zhǔn), "dense-im2col"為將推理算法固定為分塊im2col做加速的數(shù)據(jù),不使用winograd之類加速算法, "sparse-block1x1"表示對權(quán)重矩陣的稀疏為1x1分塊,也就是完全隨機(jī)稀疏。 "sparse-block1x4"是半結(jié)構(gòu)化稀疏,表示對權(quán)重矩陣的稀疏分塊為1x4分塊,沿著ic維度分塊為1, 沿著oc維度分塊為4。
從小圖可見在mobilenet V1、Mi6 機(jī)型、1x4分塊時(shí),對應(yīng)綠色左三角?曲線,稀疏加速比在0.9稀疏度是達(dá)到3.71x。
參考對比XNNPack模型推理性能:
XNNPack [6] 數(shù)據(jù)評(píng)測采用高通SD 835 CPU(MSM8998),mobilenetV1 為90%稀疏度,對mobilenetV2 為85% 稀疏度,
對應(yīng)SD 835 CPU,我們以小米6為例,下圖第1行第2列,為例對比數(shù)據(jù):
XNNPack mobilenetV1(sparsity=0.9, block=1),加速比為2.35x,MNN加速比為 2.57x , 分塊為1x4時(shí)加速比3.96x.
XNNPack mobilenetV2(sparsity=0.85, block=2),加速比為1.62x,MNN加速比為 3.71x(block=4).
對于卷積kernel不是1x1的層,我們?nèi)钥梢詫?shí)現(xiàn)稀疏加速,在XNNPack的論文中表明他們并未能實(shí)現(xiàn)稀疏加速。
通過數(shù)據(jù)大圖8,我們可得到幾點(diǎn)分析和結(jié)論:
- 參照設(shè)備小米6上, 在稀疏分塊1x4時(shí),稀疏加速臨界值優(yōu)化到0.3, 低端機(jī)型臨界值有升高,其他中高端機(jī)型, 分塊1x4時(shí), 稀疏度0.1的時(shí)候就達(dá)到加速臨界值了, 0.9稀疏度時(shí)加速比可達(dá)4.13x。
 - MNN推理耗時(shí)隨稀疏度增加,基本線性下降,跨模型、cpu 一致性比較好。
 - 內(nèi)存占用:經(jīng)過推導(dǎo), 稀疏相對稠密節(jié)省的內(nèi)存比例如下,
 



另一方面,我們評(píng)估了典型模型的分類精度,MobileNet V2在0.5稀疏度、1x4稀疏分塊配置下,精度損失0.6%,對應(yīng)上圖加速比為1.3x。

4. 業(yè)務(wù)模型實(shí)踐
4.1 某圖片超分業(yè)務(wù)
在某業(yè)務(wù)移動(dòng)端鏈路上,時(shí)間延遲與流量是其痛點(diǎn),業(yè)務(wù)方對接端智能 MNN與工作臺(tái),開發(fā)超分辨率任務(wù)模型,同時(shí)使用MNN稀疏計(jì)算加速方案,主要做了四個(gè)步驟,
- 第一步為超分模型算法訓(xùn)練;
 - 第二步,參考稀疏訓(xùn)練文檔,設(shè)定壓縮工具mnncompress需要的稀疏參數(shù),得到權(quán)重部分?jǐn)?shù)稀疏為0的模型;
 - 第三步,使用MNN convertor轉(zhuǎn)換模型;
 - 運(yùn)用MNN工作臺(tái)(https://www.mnn.zone/m/0.3/) ,部署模型到mnn運(yùn)行環(huán)境中。整套流程集成再M(fèi)NN工作臺(tái),大大提升了AI的開發(fā)流程效率,感興趣的同學(xué)可以試用和聯(lián)系工作臺(tái)負(fù)責(zé)人”明弈“。
 
下圖所示,在稀疏度0.45情況下,稀疏自身方案對推理的加速比約1.2x, 業(yè)務(wù)指標(biāo)為圖像信噪比,從 34.728dB少量下降到34.502dB。對業(yè)務(wù)精度影響在接受范圍內(nèi)。

4.2 某語音模型
另一個(gè)業(yè)務(wù)方模型為某語音模型,下圖為構(gòu)造稀疏后在avx512下實(shí)測加速比,不同稀疏方法對加速會(huì)有影響,0.75稀疏度時(shí),輸入序列20時(shí)(典型場景1),下圖所示,稀疏與稠密模型相比,encoder1加速比 2.82x, encoder2加速比3.02x。對于此方案達(dá)到預(yù)期的加速比。

5. 總結(jié)與展望
我們在MNN現(xiàn)有框架中,設(shè)計(jì)了通用稀疏卷積計(jì)算方案,使其數(shù)據(jù)布局、算子結(jié)構(gòu)配合既有MNN結(jié)構(gòu)、發(fā)揮出較高性能,而XNNPack的稀疏矩陣乘法只針對網(wǎng)絡(luò)pointwise類卷積有加速。
第一點(diǎn),我們設(shè)計(jì)實(shí)現(xiàn)了推理性能優(yōu)于XNNPack的MNN稀疏加速方案;0.9稀疏度時(shí),CV模型在ARM端獲得3.16x-4.13x加速比,跨機(jī)型、跨模型加速效果都比較顯著。
第二點(diǎn),在實(shí)際業(yè)務(wù)模型中驗(yàn)證了業(yè)務(wù)精度指標(biāo),損失有限、可接受。
第三點(diǎn),推理耗時(shí)隨稀疏度增加線性下降,跨模型、cpu 一致;在小米6上,稀疏分塊1x4加速臨界值優(yōu)化到0.3,中高端機(jī)型甚至稀疏度0.1的時(shí)候可達(dá)臨界值。
第四點(diǎn),降低內(nèi)存占用隨稀疏度成正比,具體數(shù)值見性能分析部分。
在引擎實(shí)現(xiàn)、稀疏計(jì)算內(nèi)核、匯編代碼開發(fā)中,經(jīng)常會(huì)從不及預(yù)想值開始,在一次次調(diào)試中不斷加深對MNN既有邏輯的理解,優(yōu)化稀疏算子代碼、SIMD代碼,優(yōu)化數(shù)據(jù)布局,最終將綜合指標(biāo)提升到高水準(zhǔn)。
稀疏計(jì)算研發(fā)工作得以完成,非常感謝團(tuán)隊(duì)的同學(xué)通力協(xié)作!移動(dòng)端或服務(wù)端,量化加速在不同領(lǐng)域都相對常見,指令集支持由來已久。而深度模型在CPU上的稀疏計(jì)算加速不斷發(fā)展,分別使用時(shí),二者各有相對優(yōu)勢。稀疏計(jì)算加速可以理解為一種 "可伸縮"等效位寬的技術(shù),這個(gè)角度可以探索更多獨(dú)特的應(yīng)用場景。















 
 
 









 
 
 
 