基于一階泰勒展開式的結(jié)構(gòu)化剪枝

【GiantPandaCV導(dǎo)語】文章是Nvidia在2017年ICLR發(fā)表的，核心思路是將結(jié)構(gòu)化通道剪枝判別依據(jù)視為一個優(yōu)化問題，即最小化剪枝后前后的代價函數(shù)之間的差。把優(yōu)化問題基于一階泰勒展開式近似，得到通道重要性的判據(jù)只需要考慮激活函數(shù)與對應(yīng)的梯度的乘積，再求其均值。為了跨層的公平性加上重要性歸一化，以及加入Flops正則化。

0.引言

文章是NVIDIA在2017 ICLR上面發(fā)表的，《Pruning convolution Neural Network for resource efficient inference》。其核心也是將結(jié)構(gòu)化通道剪枝判別依據(jù)視為一個優(yōu)化問題，最小化剪枝后前后的代價函數(shù)之間的差。

四個要點:

1、優(yōu)化方程是最小化已減去權(quán)重的代價函數(shù)與未剪枝之前權(quán)重的代價函數(shù)之差，用一階泰勒展開式去近似優(yōu)化方程，計算通道重要性只需要激活函數(shù)與對應(yīng)的梯度的乘積，再求其均值；

2、對計算出來的重要性進行歸一化；

3、加入FLOPs等正則化，使其硬件更友好；

4、可以聯(lián)合其他通道重要性判據(jù)

圖片描述：三步剪枝流程圖

一.剪枝的優(yōu)化方程

這一部分是為了推導(dǎo)出公式(6)、公式(7)和公式(8)。

訓(xùn)練樣本集合

表示輸入， 表示輸出

網(wǎng)絡(luò)模型的參數(shù)

模型參數(shù)的代價函數(shù)為 ，其中 是負對數(shù)似然函數(shù)(negative log-likelihood function)

剪枝的目標是最小化已減去的權(quán)重的代價函數(shù)與未剪枝之前的代價函數(shù)的差，記為公式(1)：

其中 ， 是范數(shù)，也即非0值的個數(shù)， 是被剪掉的權(quán)重參數(shù)。

按照最樸素遍歷的方法，剪去一組參數(shù)需要次計算，顯然不科學(xué)的。比如VGG-16有4224個卷積特征圖（可以對照VGG模型數(shù)出來）。

利用泰勒展開式的方法來近似優(yōu)化方程：

一組卷積核：

，是一組特征圖， 表示layer

單張?zhí)卣鲌D

,表示剪枝與否。

所以

特征圖集合記為：

是由一組參數(shù)得到的

公式(2)

表示被剪去參數(shù)的代價函數(shù)是原來參數(shù)的代價函數(shù)

公式(3) 為泰勒公式：

用一階泰勒展開式來逼近，即，

公式(4)：

余項采用拉格朗日（Lagrange）余項：

所以 :

忽略余項，將公式(4)代入公式(2)，

公式(5)：

公式(6)：

其中 ?是特征圖向量的長度。公式5和公式6本質(zhì)上沒有什么區(qū)別，只是等價變換。

到這里，優(yōu)化方程已經(jīng)得到了，就是公式6，得出的就是通道重要性的評分，這里計算通道重要性只需要計算激活函數(shù)與對應(yīng)的梯度的乘積，再求其均值。

另外，為了保證跨層的通道之間的公平性，對特征圖作了歸一化，

公式(7)

其中是第層中第個通道（參數(shù)組）的重要性。

最后一點，做了FLOPs正則化，這個是為了減少FLOPs，這里也可以對其他參數(shù)做正則化，比如存儲大小，內(nèi)存，卷積核。

公式(8)

二、核心代碼實現(xiàn)

整個代碼不算大，但是全部貼上來也占據(jù)不少位置，于是我就貼上文章核心觀點的重要代碼，全部代碼在參考鏈接。

圖片描述：整體代碼流程圖

1、通道重要性計算

def?compute_rank(self,?grad,?lamda=1):
????activation_index?=?len(self.activations)?-?self.grad_index?-?1
????activation?=?self.activations[activation_index]
????
????taylor?=?activation?*?grad????????????##計算基于一階泰勒展開的值???
????taylor?=?taylor.mean(dim=(0,?2,?3)).data????##?以通道維度計算得分

????flops?=?self.flops_dict[activation_index]?*?torch.ones_like(taylor)?###FLOPs?正則化,見下面get_flops()
????taylor?=?taylor?-?lamda?*?flops?#?公式(8)

????if?activation_index?not?in?self.filter_ranks:
????????self.filter_ranks[activation_index]?=?torch.FloatTensor(activation.size(1)).zero_().cuda()

????????self.filter_ranks[activation_index]?+=?taylor
????????self.grad_index?+=?1

2、計算FLOPs

def?get_flops(self):
????self.flops_dict?=?{}
????def?hook(module,?inputs,?outputs):??#hook函數(shù)，用來計算每層的flops
????????params?=?module.weight.size().numel()??#torch.numel()?返回一個tensor變量內(nèi)所有元素個數(shù)，可以理解為矩陣內(nèi)元素的個數(shù)
????????W?=?outputs.size(2)
????????H?=?outputs.size(3)
????????module.flops?=?params?*?W?*?H??#計算flops
????????
????for?m?in?model.modules():
????????if?isinstance(m,?nn.Conv2d):
????????????m.register_forward_hook(hook)?#為每一層注冊hook
????????x?=?torch.randn(1,?3,?224,?224)?#隨便定義一個數(shù)據(jù)，跑一次forward來獲取flops
????????_?=?self.model(x)
????????activation_index?=?0
????????for?m?in?self.model.modules():
????????????if?isinstance(m,?nn.Conv2d):
????????????????self.flops_dict[activation_index]?=?m.flops??#獲取每層的flops，記錄到字典
????????????????activation_index?+=?1

3、對基于一階泰勒展開式計算得到重要性，進行通道歸一化

def?normalize_ranks_per_layer(self):
????for?i?in?self.filter_ranks:
????????v?=?torch.abs(self.filter_ranks[i])
????????v?=?v?/?torch.sqrt(torch.sum(v?*?v))
????????self.filter_ranks[i]?=?v

三、文章的一些討論

1、泰勒展開的討論

1990年 LeCun提出Optimal Brain Damage (OBD) 也是采用泰勒展開式的方法來近似。因為經(jīng)過多次訓(xùn)練迭代后，梯度會趨向于0，那么，則，所以LeCun認為提供的信息很有限， ，所以采用泰勒二階展開式。

雖然的期望為0，但是他的方差不為0，其，所以，所以，可以用一階泰勒展開式來近似。

2、通道重要性

雖然文章提出了一種基于一階泰勒展開的方法計算通道重要性，從flops正則化中我們可以，將其他通道重要性判別也加入最后的重要性排序，也就是在公式8中加入其他項，這樣可以更加全面的考慮通道重要性；當然，隨之增加的就是整個剪枝算法的計算量和耗時也增加。

四、參考鏈接

• 論文：https://arxiv.org/abs/1611.06440

• 代碼：https://github.com/azhe198827/channel_prune

歡迎關(guān)注GiantPandaCV, 在這里你將看到獨家的深度學(xué)習(xí)分享，堅持原創(chuàng)，每天分享我們學(xué)習(xí)到的新鮮知識。( ? ?ω?? )?

有對文章相關(guān)的問題，或者想要加入交流群，歡迎添加BBuf微信：

為了方便讀者獲取資料以及我們公眾號的作者發(fā)布一些Github工程的更新，我們成立了一個QQ群，二維碼如下，感興趣可以加入。