模型壓縮 | 結構性剪枝Data-Driven Sparse Structure Selection 以及實際剪枝實現流程

前言

剪枝的文章很多,目前稀疏化的方法有主流的:group lasso; 控制網絡結構收縮constrain the structure scale;對網絡結構正則化處理regularizing multiple DNN structures。本文主要側重於channel prune 乃至group/block等結構上稀疏化的方法。這裏主要是基於圖森的文章: Data-Driven Sparse Structure Selection for Deep Neural Networks。此外,如題,模型壓縮是很實用的一個方向,言外之意就是工程意義很大。但是論文往往不會提及如何實際上將“繁枝”剪掉。提醒一句,就算乘以0也算FLOPs。只有要剪掉的枝從搭建的網絡中消失纔算哦。所以博主最後會將github代碼也附上。另外,水平一般能力有限,歡迎大家斧正。 (更多模型壓縮文章如各種知識蒸餾方法與實現,以後會從筆記本上整理陸續放出)。

一.文獻解讀

I. Networks Slimming

此文思路還是很清楚的,我想從相關文獻說起。Networks Slimming-Learning Efficient Convolutional Networks through Network Slimming
該論文pytorch實現代碼地址https://github.com/mengrang/Slimming-pytorch
思路歸納爲兩點:
1. 利用BN中的縮放因子gamma 作爲評價上層輸出貢獻大小(下層輸入)的因子,即gamma越小,所對應的神經元越不重要,就可以裁剪掉。

2. 利用L1正則化gamma,使其稀疏化,這樣就可以在訓練過程中自動評價神經元貢獻大小,爲0的銀子可以安全剪掉。

這個思路還是很巧妙的。在這裏我不得不佩服作者對深度學習的思考深度。但是,這個方法工程上會有一個需要注意的地方。根據BN原理,需要訓練的因子有兩種,scale 和shift。gamma是相乘因子,之後還要加上一個shift,beta.那麼,gamma很小的時候,beta是不是很大?一定要保留beta.

II. Sparse Structure Selection

下面介紹圖森的SSS。這裏沒有使用BN帶的gamma,而是新增加一個參數作爲衡量神經元輸出,每一group/block輸出重要性的因子。原理很清晰,就是在每個神經元輸出上,每個殘差結構卷積分支上以及一些group卷積的非identity group上之輸出乘上一個縮放因子。訓練過程中對這些因子施加L1正則化使得其稀疏化。這樣,爲0的因子對應的神經元/group/block自然而然的認爲對結果不重要,安全移除。
文章的思路是非常棒的,但是我認爲解決L1正則化不可導的思路也很好。首先我們知道,對於不可導的L1 norm 一般使用Proximal Gradient來反傳。(這裏插一句,強化學習中的PPO,近端策略優化也是解決不可導問題,第一個P的含義)
根據推導過程發現,我們要想使用該方法,顯然要進行兩次前傳代價太大。(原因就是要對重要性因子之差再進行一次然後根據牛頓動量優化器原理髮現)唉,其實這個APG算法只不過比牛頓動量優化算法多出一個投影函數(PG算法中的概念,不懂的可以網上搜資料補補),而根據L1正則化Proximal Gradient原理,這個投影函數就是軟閾值函數(有數學證明,這裏就不寫了,因爲我也沒證明)。所以,只需要在牛頓動量的優化器中對縮放因子施加軟閾值函數即可,軟閾值兩個參數是學習率和對應縮放因子值。

二. 剪枝的實現

至此,文章的思路就介紹完畢了。源碼是MXNET寫的。但是,這裏我們發現一個問題。正如博主開頭說的,要想達到模型參數量減少目的,這只是一個開端。因爲即使相乘參數爲0,不把這一神經元/group/block從網絡結構中拿掉,計算量不變。這是所有論文裏不給說的部分,從工程上講,這纔是實際的剪枝。下面介紹一下Tensorflow如何實現剪枝。github代碼也附上

1.前文加入縮放因子的網絡已經訓練完畢,稀疏化完畢。從ckpt中取出重要性因子的值,以及各神經元權重。

2.用numpy數組操作,記錄小於設定閾值的元素位置,找到對應神經元數組位置,直接按channel整個切除。

3.參考以上切除之後的對應的權值數組搭建網絡,以此設定張量大小。注意,block級別的稀疏,直接讓輸出跨過這一block即可。然而,channel級別的稀疏,由於殘差結構的存在,相加位置張量channel要一致,是個坑。解決方案是不要對block最後一層的bottleneck稀疏即可。

4.應用指定變量初始化方法對重新定義的網絡結構初始化,fineturn一下。

水平有限,歡迎大家斧正

HongYuSuiXinLang
Zhejiang University
Email: [email protected]

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