模型壓縮:剪枝演算法

語言: CN / TW / HK

過引數化主要是指在訓練階段,在數學上需要進行大量的微分求解,去捕抓資料中的微小變化資訊,一旦完成迭代式的訓練之後,網路模型推理的時候就不需要這麼多引數。而剪枝演算法正是基於過引數化的理論基礎而提出的。

剪枝演算法核心思想就是減少網路模型中引數量和計算量,同時儘量保證模型的效能不受影響。

那在AI框架中,實際上剪枝主要作用在右下角的端側模型推理應用場景中,為的就是讓端側模型更小,無論是平板、手機、手錶、耳機等小型IOT裝置都可以輕鬆使用AI模型。而實際在訓練過程更多體現在剪枝演算法和框架提供的剪枝API上面。

剪枝演算法分類

實際上大部分剛接觸剪枝演算法的時候,都會從從巨集觀層面去劃分剪枝技術,主要是分為Drop Out和Drop Connect兩種經典的剪枝演算法,如下圖所示。

1)Drop Out:隨機的將一些神經元的輸出置零,稱之為神經元剪枝。

2)Drop Connect:隨機將部分神經元間的連線Connect置零,使得權重連線矩陣變得稀疏。

結構化剪枝 VS 非結構化剪枝

下面會把剪枝的更多種方式呈現出來,可能會稍微複雜哈。從剪枝的粒度來劃分,可以分為結構化剪枝和非結構化剪枝,2個剪枝結構方法。下面來看看具體的剪枝方法有4種:

1) 細粒度剪枝(fine-grained) :即對連線或者神經元進行剪枝,是粒度最小的剪枝,上面Drop Out和Drop Connect都是屬於細粒度剪枝。

2) 向量剪枝(vector-level) :它相對於細粒度剪枝粒度稍大,屬於對卷積核內部(intra-kernel) 的剪枝。

3) 核剪枝(kernel-level) :即去除某個卷積核,丟棄對輸入通道中對應卷積核的計算。

4) 濾波器剪枝(Filter-level) :對整個卷積核組進行剪枝,推理過程中輸出特徵通道數會改變。

細粒度剪枝、向量剪枝、核剪枝在引數量與模型效能之間取得了一定的平衡,但是網路模型單層的神經元之間的組合結構發生了變化,需要專門的演算法或者硬體結構來支援稀疏的運算,這種叫做結構化剪枝(Unstructured Pruning)

其中,非結構化剪枝能夠實現更高的壓縮率,同時保持較高的模型效能,然而會帶來網路模型稀疏化,其稀疏結構對於硬體加速計算並不友好,除非底層硬體和計算加速庫對稀疏計算有比較好的支援,否則剪枝後很難獲得實質的效能提升。

濾波器剪枝(Filter-level)主要改變網路中的濾波器組和特徵通道數目,所獲得的模型不需要專門的演算法和硬體就能夠執行,被稱為結構化剪枝(Structured Pruning) 。結構化剪枝又可進一步細分:可以是channel-wise,也可以是filter-wise,還可以是在shape-wise。

結構化剪枝與非結構化剪枝恰恰相反,可以方便改變網路模型的結構特徵,從而達到壓縮模型的效果,例如知識蒸餾中的student網路模型、NAS搜尋或者如VGG19和VGG16這種裁剪模型,也可以看做變相的結構化剪枝行為。

剪枝演算法流程

雖然剪枝演算法的分類看上去很多,但是核心思想還是對神經網路模型進行剪枝,目前剪枝演算法的總體流程大同小異,可以歸結為三種:標準剪枝、基於子模型取樣的剪枝、以及基於搜尋的剪枝,如下圖所示。

標準剪枝演算法流程

標準剪枝是目前最流行的剪枝流程,在Tensorflow、Pytroch都有標準的介面。主要包含三個部分:訓練、剪枝、以及微調。

1)訓練:首先是對網路模型進行訓練。在剪枝流程中,訓練部分主要指預訓練,訓練的目的是為剪枝演算法獲得在特定基礎SOTA任務上訓練好的原始模型。

2) 剪枝:在這裡面可以進行如細粒度剪枝、向量剪枝、核剪枝、濾波器剪枝等各種不同的剪枝演算法。其中很重要的就一點,就是在剪枝之後,對網路模型結構進行評估。確定一個需要剪枝的層,設定一個裁剪閾值或者比例。實現上,通過修改程式碼加入一個與引數矩陣尺寸一致的Mask矩陣。Mask矩陣中只有0和1,實際上是用於微調網路。

3)微調:微調是恢復被剪枝操作影響的模型表達能力的必要步驟。結構化模型剪枝會對原始模型結構進行調整,因此剪枝後的模型引數雖然保留了原始的模型引數,但是由於模型結構的改變,剪枝後模型的表達能力會受到一定程度的影響。實現上,微調網路模型,引數在計算的時候先乘以該Mask,Mask為1的引數值將繼續訓練通過BP調整梯度,而Mask為0的部分因為輸出始終為0則不對後續部分產生影響。

4)再剪枝:再剪枝過程將微調之後的網路模型再送到剪枝模組中,再次進行模型結構評估和執行剪枝演算法。目的是使得每次剪枝都在效能更優的模型上面進行,不斷迭代式地進行優化剪枝模型,直到模型能夠滿足剪枝目標需求。

最後輸出模型引數儲存的時候,因為有大量的稀疏,所以可以重新定義儲存的資料結構, 僅儲存非零值以及其矩陣位置。重新讀取模型引數的時候,就可以還原矩陣。

基於子模型取樣流程

除標準剪枝之外,基於子模型取樣的剪枝《EagleEye: Fast sub-net evaluation for efficient neural network pruning》最近也表現出比較好的剪枝效果。得到訓練好的模型之後,進行子模型取樣過程。一次子模型取樣過程為:

1)對訓練好的原模型中可修剪的網路結構,按照剪枝目標進行取樣,取樣過程可以是隨機的,也可以按照網路結構的重要性或者通過KL散度計算進行概率取樣。

2)對取樣後的網路結構進行剪枝,得到取樣子模型。子模型取樣過程通常進行 次,得到 個子模型( ≥1), 之後對每一個子模型進行效能評估。子模型評估結束之後,選取最優的子模型進行微調以得倒最後的剪枝模型。

基於搜尋的剪枝流程

基於搜尋的剪枝主要依靠強化學習等一系列無監督學習或者半監督學習演算法,也可以是神經網路結構搜尋相關理論。

給定剪枝目標之後,基於搜尋的剪枝在網路結構中搜索較優的子結構,這個搜尋過程往往伴隨著網路引數的學習過程,因此一些基於搜尋的剪枝演算法在剪枝結束後不需要再進行微調。

剪枝的發展

這幾年神經網路剪枝pruning作為模型壓縮技術的四小龍之一,正在受到越來越多的關注。當然,各種更好的pruning引數選取方法一定還會層出不窮。另外,從趨勢來看,以下幾個方向值得關注:

打破固定假設:挑戰已有的固有的假設,例如ICLR2019會議的best paper彩票假說《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks 》的出現。還有一開始提到的對於over-parameterization,與重用已有引數是否有有益的反思非常有意思。這樣的工作會給剪枝演算法非常大的啟發,從而根本改變解決問題的思路。

自動化剪枝:隨著AutoML的大潮,越來越多的演算法開始走向自動化。模型壓縮能拉下嗎?當然不能。經過前面的介紹我們知道,像ADC,RNP,N2N Learning這些工作都是試圖將剪枝中部分工作自動化。如量化中的《HAQ: Hardware-Aware Automated Quantization》考慮網路中不同層資訊的冗餘程度不一樣,所以自動化使用混合量化位元進行壓縮。

與NAS融合:如前面模型剪枝流程中提到,剪枝演算法與神經網路搜尋NAS的界限已經模糊了。NAS有針對結構化剪枝進行搜尋方法,如One-Shot Architecture Search是先有一個大網路,然後做減法。NAS與模型壓縮兩個一開始看似關係不是那麼大的分支,在近幾年的發展過程中因為下游任務和部署場景的需求,最後似乎會走到一塊去。這兩個分支今天有了更多的交集,也必將擦出更多的火花。

與GAN融合:這幾年機器學習最火熱的分支之一GAN,正在不斷滲透到已有領域,在pruning中也開始有它的身影。如2019年《Towards Optimal Structured CNN Pruning via Generative Adversarial Learning》讓generator生成裁剪後網路,discrimintor來判別是否屬於原網路還是裁剪後網路,從而進行更有效的網路結構化裁剪。

硬體稀疏性支援:剪枝會給神經網路模型帶來稀疏性特徵,引數稀疏性在計算中會有大量的索引,所以並不能加速。現在雖然有像cuSPARSE這樣的計算庫,但底層硬體AI晶片本身設計並不是專門為稀疏資料處理打造的。如果能將稀疏計算和處理能力做進晶片那必將極大提高計算效率。僅2021年中國就推出了10+款基於ASIC的AI加速晶片,相信針對稀疏性場景的支援在未來會有所突破。

總結

模型壓縮演算法中針對已有的模型,有:張量分解,模型剪枝,模型量化。針對新構建的網路,有:知識蒸餾,緊湊網路設計等方法。

剪枝只是模型壓縮方法中的一種,它與其它模型壓縮方法並不衝突,因此會與量化、蒸餾、NAS、強化學習等方法慢慢融合,這些都是很值得研究的方向。另外在上面的發展來看,打破固有的假設定義,與NAS、GAN、AutoML、RL等技術進行相互的融合,可能到最後會模糊purning方式,出現新的正規化或者壓縮模式也是很吸引的。