深度學習模型輕量化（下）
2.4 蒸餾
2.4.1 蒸餾流程
蒸餾本質是student對teacher的擬合，從teacher中汲取養分，學到知識，不僅僅可以用到模型壓縮和加速中。蒸餾常見流程如下圖所示
-
老師和學生可以是不同的網絡結構，比如BERT蒸餾到BiLSTM網絡。但一般相似網絡結構，蒸餾效果會更好。
-
總體loss爲 soft_label_loss + hard_label_loss。soft_label_loss可以用KL散度或MSE擬合
-
soft label爲teacher模型的要擬合的對象。可以是模型預測輸出，也可以是embeddings, 或者hidden layer和attention分佈。
針對軟標籤的定義，蒸餾的方案也是百花齊放，下面分享兩篇個人認爲非常經典的文章。
2.4.2 distillBERT
DistillBERT: A distilled version of BERT: smaller,
faster, cheaper and lighter
DistillBERT由大名鼎鼎的HuggingFace出品。主要創新點爲：
-
Teacher 12層，student 6層，每兩層去掉一層。比如student第二層對應teacher第三層
-
Loss= 5.0 * Lce+2.0 * Lmlm+1.0 * Lcos
·
Lce: soft_label 的KL散度
·
Lmlm: mask LM hard_label
的交叉熵
·
Lcos：hidden state 的餘弦相似度
DistilBERT 比 BERT 快 60%，體積比 BERT 小 60%。在glue任務上，保留了 95% 以上的性能。在performance損失很小的情況下，帶來了較大的模型壓縮和加速效果。
2.4.3 TinyBERT
TinyBERT: Distilling BERT for Natural Language
Understanding
總體結構
重點來看下 TinyBERT，它是由華爲出品，非常值得深入研究。TinyBERT 對 embedding 層，transformer層（包括hidden layer和attention）和 prediction 層均進行了擬合。如下圖所示。
TinyBERT蒸餾過程
其中Embeddings採用MSE, Prediction採用KL散度, Transformer層的hidden layer和attention，均採用MSE。loss如下
其中m爲層數。
效果分析
表2: glue任務上的performance。在glue任務上，可達到bert-base的96%，幾乎無損失。表3: tinyBERT模型大小和推理速度。縮小7.5倍，加速9.4倍。壓縮和加速效果十分明顯。
消融分析
表6：分析embedding、prediction、attention、hidden layer軟標籤作用，其中attention和hidden layer作用最大。這個也很好理解，transformer層本來就是整個BERT中最關鍵的部分。
表7：分析老師學生不同層對應方法的效果，uniform爲隔層對應，top爲全部對應老師頂部幾層，bottom爲全部對應老師底部幾層。Uniform效果明顯好很多。這個也很好理解，淺層可以捕捉低階特徵，深層可以捕捉高階特徵。全是低階或者高階顯然不合適，我們要儘量葷素搭配。
3 框架層加速
3.1 手機端AI能力
目前移動端AI框架也比較多，包括谷歌的tf-lite，騰訊的NCNN，阿里的MNN，百度的PaddleLite,
小米的MACE等。他們都不同程度的進行了模型壓縮和加速的支持。特別是端上推理的加速。這個可以參考“手機端AI性能排名“。
3.2 端側AI框架加速優化方法
個人總結的主要方法如下，可能有遺漏哈，各位看官請輕拍：
- 基於基本的C++編譯器優化。
a. 打開編譯器的優化選項，選擇O2等加速選項。
b. 小函數內聯，概率大分支優先，避免除法，查表空間換時間，函數參數不超過4個等。
-
利用C，而不是C++，C++有不少冗餘的東西。
-
緩存優化
a. 小塊內存反覆使用，提升cache命中率，儘量減少內存申請。比如上一層計算完後，接着用作下一層計算。
b. 連續訪問，內存連續訪問有利於一次同時取數，相近位置cache命中概率更高。比如縱向訪問數組時，可以考慮轉置後變爲橫向訪問。
c. 對齊訪問，比如224224的尺寸，補齊爲256224，從而提高緩存命中率。
d. 緩存預取，CPU計算的時候，preload後面的數據到cache中。
- 多線程。
a. 爲循環分配線程。
b. 動態調度，某個子循環過慢的時候，調度一部分循環到其他線程中。
- 稀疏化
a. 稀疏索引和存儲方案，採用eigen的sparseMatrix方案。
- 內存複用和提前申請
a. 掃描整個網絡，計算每層網絡內存複用的情況下，最低的內存消耗。推理剛開始的時候就提前申請好。避免推理過程中反覆申請和釋放內存，避免推理過程中因爲內存不足而失敗，複用提升內存訪問效率和cache命中率。
-
ARM NEON指令的使用，和ARM的深度融合。NEON可以單指令多取值（SIMD），感興趣可針對學習，這一塊水也很深。
-
手工彙編，畢竟機器編譯出來的代碼還是有不少冗餘的。可以針對運行頻次特別高的代碼進行手工彙編優化。當然如果你彙編功底驚天地泣鬼神的強，也可以全方位手工彙編。
-
算子支持：比如支持GPU加速，支持定點化等。有時候需要重新開發端側的算子。
4 硬件層加速
硬件層加速比較硬核，小編就連半瓢水都達不到了，爲了保證整個方案的全面性，還是硬着頭皮東施效顰下。目前AI芯片廠家也是百花齊放，誰都想插一腳，不少互聯網公司也來趕集，如下圖所示。
AI 芯片目前三種方案。GPU目前被英偉達和AMD牢牢把控。ASIC目前最火，TPU、NPU等屬於ASIC範疇。
5 總結
本文對深度學習模型壓縮和加速的幾類常用的方法進行了介紹。
參考文獻
-
ALBERT: A Lite BERT for Self-supervised Learning of
Language Representations -
MobileNets: Efficient Convolutional Neural Networks for
Mobile Vision Applications -
Are Sixteen Heads Really Better than One?
-
DistillBERT: A distilled version of BERT: smaller, faster,
cheaper and lighter -
TinyBERT: Distilling BERT for Natural Language
Understanding -
手機端AI性能排名