【阅读笔记】DMCP: Differentiable Markov Channel Pruning for Neural Networks

前言

这篇论文来自商汤研究院, 为CVPR2020 Oral论文, 提出了基于马尔可夫模型的通道剪枝策略, 该方法具有较好的通用性和可操作性. 论文PDF

DMCP算法步骤:

  1. 将通道剪枝作为马尔可夫过程进行建模
  2. 使用网络原始损失函数和预算正则化(限制模型的FLOPs)对剪枝参数进行优化
  3. 使用DS或ES策略对马尔可夫模型进行剪枝,得到剪枝模型
  4. 对剪枝后的模型进行重头开始训练,评估子网络的性能

文中提出的可微体现在, 预算正则化无法被直接优化,提出了可微分的方法在训练时对其进行优化

对照:
预算正则化: Budget Regularization

马尔可夫建模

剪枝过程的定义

文中使用 M ( L ( 1 ) , L ( 2 ) , … , L ( N ) ) M\left(L^{(1)}, L^{(2)}, \ldots, L^{(N)}\right) M(L(1),L(2),,L(N))表示N层未剪枝网络模型,其中 L ( i ) L^{(i)} L(i)表示网络模型中的第 i i i层。网络中的第 i i i层第 k k k个通道的输出可以由公式1给出,
O k ( i ) = w k ( i ) ⊙ x , k = 1 , 2 , … , C o u t ( i )   ( 1 ) O_{k}^{(i)}=w_{k}^{(i)} \odot x, k=1,2, \ldots, C_{o u t}^{(i)} (1) Ok(i)=wk(i)x,k=1,2,,Cout(i) (1)
O k ( i ) O_{k}^{(i)} Ok(i)表示第 i i i层中的第 k k k个通道的输出, w k ( i ) w_{k}^{(i)} wk(i)表示 L ( i ) L^{(i)} L(i)层中的第 k k k个Filter, ⊙ \odot 表示卷积操作. C o u t ( i ) C_{o u t}^{(i)} Cout(i)表示 L ( i ) L^{(i)} L(i)层中的第 k k k个Filter的个数. 文中为了方便, 后面的符号中都省略了上标.

说明:
" O k ( i ) O_{k}^{(i)} Ok(i)表示第 i i i层中的第 k k k个通道的输出, w k ( i ) w_{k}^{(i)} wk(i)表示 L ( i ) L^{(i)} L(i)层中的第 k k k个Filter"的通道Filter说明, 如下图所示, 网络的输入为6x6x3大小特征图, Filter(卷积核)大小为3x3x3x2, 其中x2表示卷积核(Filter)的个数为2, 经过卷积操作后, 输出为4x4x2大小的特征图, 这里的2即为输出特征图的通道. 综上, Filter即为卷积核, 通道为输出特征图的’深度’, Filter的个数等于输出特征图的通道数.
在这里插入图片描述

通道和Filter说明图片
[图片来源](https://indoml.com/2018/03/07/student-notes-convolutional-neural-networks-cnn-introduction)

网络第 i i i层的马尔可夫模型描述如下, 下图中 1 − p k 1-p_k 1pk表示达到终止状态的概率, 状态 S k ( 1 ≤ k ≤ C out ) S_{k}\left(1 \leq k \leq C_{\text {out}}\right) Sk(1kCout)表示剪枝中保留的第 k k k个通道, 从 S k S_k Sk S k + 1 S_{k+1} Sk+1的转移概率 p k p_k pk表示在第 k k k个通道保留的前提下,保留地 k + 1 k+1 k+1个通道的概率. 当从他状态转移到终止状态T, 即表明剪枝过程结束.

该过程有如下性质, 如果 L L L层总数为 C c o u t C_{cout} Ccout个通道中有 k k k个通道被保留了, 那么这k个通道必须是前 k k k个通道. 即, 如果第 k k k个通道被保留了, 那么前 k − 1 k-1 k1个通道也会被保留. 可得出, 保留 k + 1 k+1 k+1个通道条件独立于当第 k k k个通道保留的前提下前 k − 1 k-1 k1个通道
在这里插入图片描述

通过马尔可夫过程进行通道剪枝

文中使用一堆可学习的参数对马尔可夫转移状态进行参数化,这些可学习的参数被称为结构参数. p ( w 1 , w 2 , … , w k − 1 ) p\left(w_{1}, w_{2}, \ldots, w_{k}-1\right) p(w1,w2,,wk1)用来表示前 k − 1 k-1 k1个通道的保留概率, 前 k k k个通道的保留概率可以如下表示,
p ( w 1 , … , w k ) = p ( w k ∣ w 1 , … , w k − 1 ) p ( w 1 , … , w k − 1 )   ( 2 ) p\left(w_{1}, \ldots, w_{k}\right)=p\left(w_{k} \mid w_{1}, \ldots, w_{k-1}\right) p\left(w_{1}, \ldots, w_{k-1}\right) (2) p(w1,,wk)=p(wkw1,,wk1)p(w1,,wk1) (2)
公式2中 p ( w k ∣ w 1 , … , w k − 1 ) p\left(w_{k} \mid w_{1}, \ldots, w_{k-1}\right) p(wkw1,,wk1)表示在保留前 k − 1 k-1 k1个通道的前提下, 保留的 k k k个通道的概率. 由于 w k w_k wk w k − 1 w_{k-1} wk1保留的前提下, 条件独立于 { w 1 , w 2 , … w k − 2 } \left\{w_{1}, w_{2}, \ldots w_{k-2}\right\} {w1,w2,wk2}, 所以可以将公式2改写成,
p k = p ( w k ∣ w 1 , w 2 , … , w k − 1 ) = p ( w k ∣ w k − 1 )   ( 3 ) p_{k}=p\left(w_{k} \mid w_{1}, w_{2}, \ldots, w_{k-1}\right)=p\left(w_{k} \mid w_{k-1}\right) (3) pk=p(wkw1,w2,,wk1)=p(wkwk1) (3)
p ( w k ∣ ¬ w k − 1 ) = 0   ( 4 ) p\left(w_{k} \mid \neg w_{k-1}\right)=0 (4) p(wk¬wk1)=0 (4)
¬ w k − 1 \neg w_{k-1} ¬wk1表示舍弃第 k − 1 k-1 k1个通道. 转移概率 P = { p 1 , p 2 , . . p C o u t } P=\left\{p_{1}, p_{2}, . . p_{C_{o u t}}\right\} P={p1,p2,..pCout}可以通过公式3定义, 并且使用结构参数 A = { α 1 , α 2 , … , α c out } A= \left\{\alpha_{1}, \alpha_{2}, \ldots, \alpha_{c_{\text {out}}}\right\} A={α1,α2,,αcout}来参数化转移概率 P P P, 结构参数A在网络参数初始化时会同时被初始化. p k p_k pk由公式5计算得出,
p k = { 1 k = 1 sigmoid ( α k ) = 1 1 + e − α k k = 2 , … , C out , α k ∈ A   ( 5 ) p_{k}=\left\{\begin{array}{ll} 1 & k=1 \\ \text {sigmoid}\left(\alpha_{k}\right)=\frac{1}{1+e^{-\alpha_{k}}} & k=2, \ldots, C_{\text {out}}, \alpha_{k} \in A \end{array}\right. (5) pk={1sigmoid(αk)=1+eαk1k=1k=2,,Cout,αkA (5)
文中设置每个卷积操作至少保留一个通道, 所以有, p 1 = p ( w 1 ) = 1 p_{1}=p\left(w_{1}\right)=1 p1=p(w1)=1 .通道 w k w_k wk的边缘概率使用 p ( w k ) p(w_k) p(wk)来表示, 并使用公式6计算,
p ( w k ) = p ( w k ∣ w k − 1 ) p ( w k − 1 ) + p ( w k ∣ ¬ w k − 1 ) p ( ¬ w k − 1 ) = p ( w k ∣ w k − 1 ) p ( w k − 1 ) + 0 = p ( w 1 ) ∏ i = 2 k p ( w i ∣ w i − 1 ) = ∏ i = 1 k p i ( 6 ) \begin{aligned} p\left(w_{k}\right) &=p\left(w_{k} \mid w_{k-1}\right) p\left(w_{k-1}\right)+p\left(w_{k} \mid \neg w_{k-1}\right) p\left(\neg w_{k-1}\right) \\ &=p\left(w_{k} \mid w_{k-1}\right) p\left(w_{k-1}\right)+0 \\ &=p\left(w_{1}\right) \prod_{i=2}^{k} p\left(w_{i} \mid w_{i-1}\right)=\prod_{i=1}^{k} p_{i} \end{aligned}(6) p(wk)=p(wkwk1)p(wk1)+p(wk¬wk1)p(¬wk1)=p(wkwk1)p(wk1)+0=p(w1)i=2kp(wiwi1)=i=1kpi(6)
为了对结构参数进行优化, 文中将结构参数与网络中某一个通道的输出进行结合, 使得结构参数可以通过梯度回传进行优化, 结构参数和通道输出的结合通过公式7实现,
O ^ k = O k × p ( w k )   ( 7 ) \hat{O}_{k}=O_{k} \times p\left(w_{k}\right)  (7) O^k=Ok×p(wk) (7)
O ^ k \hat{O}_{k} O^k表示第 k k k个通道输出, 如果要对 w k w_k wk(第 k k k个Filter)进行剪枝, 将 p ( w k ) p(w_k) p(wk)置0即可.

注意, 由于BN会对第 i i i层的输出进行缩放, 所以公式7不能在卷积层之后实现, 而是将其放在BN层之后.

短连接问题

MobileNetV2和ResNet都有含短连接的残差块, 由于element-wise操作, 最后一个卷积层的通道数必须和前一个块的通道数相等.文中使用权重共享解决此问题.

预算正则化

预算正则化目的是约束网络模型FLOPs. L L L层中所期望的通道数 E ( c h a n n e l ) E(channel) E(channel)可以用公式8表示,
E ( channel ) = ∑ i = 1 C out p ( w i ) ( 8 ) E(\text {channel})=\sum_{i=1}^{C_{\text {out}}} p\left(w_{i}\right)(8) E(channel)=i=1Coutp(wi)(8)
L L L层中期望输入通道数 E ( i n ) E(in) E(in)和输出通道数 E ( o u t ) E(out) E(out)可以通过公式8计算, 期望FLOPs E ( L F L O P s ) E(L_{FLOPs}) E(LFLOPs)可以通过以下公式计算,
E ( L F L O P s ) = E ( out ) × E ( kernel_op ) ( 9 ) E\left(L_{F L O P s}\right)=\quad E(\text {out}) \times E(\text {kernel\_op}) (9) E(LFLOPs)=E(out)×E(kernel_op)(9)
E ( kernel_op ) = E ( in ) groups × #  channel_op  ( 10 ) E(\text {kernel\_op})=\frac{E(\text {in})}{\text {groups}} \times \# \text { channel\_op } (10) E(kernel_op)=groupsE(in)×# channel_op (10)
#  channel_op  = ( S I + S P − S K stride + 1 ) × S K × S K ( 11 ) \# \text { channel\_op }=\left(\frac{S_{I}+S_{P}-S_{K}}{\text {stride}}+1\right) \times S_{K} \times S_{K}(11) # channel_op =(strideSI+SPSK+1)×SK×SK(11)
这里的groups为分组卷积的组数. S I S_I SI S K S_K SK分别表示输入特征图和卷积核的宽(或者高). S P S_P SP表示填充大小,stride表示卷积步长. 整个模型的期望FLOPs用 E ( N F L O P s ) E(N_{FLOPs}) E(NFLOPs)表示, 可由公式12计算,
E ( N F L O P s ) = ∑ l = 1 N E ( l ) ( L F L O P s ) ( 12 ) E\left(N_{F L O P s}\right)=\sum_{l=1}^{N} E^{(l)}\left(L_{F L O P s}\right) (12) E(NFLOPs)=l=1NE(l)(LFLOPs)(12)
公式12中的 N N N表示网络中卷积层的数量. 通过该公式, 可以通过梯度下降来优化FLOPs.

损失函数

设置要达到的FLOPs为 F L O P s t a r g e t FLOPs_{target} FLOPstarget, 使用公式12来公表示可微分的预算正则化的损失函数 l o s s r e g loss_{reg} lossreg,
loss ⁡ r e g = log ⁡ ( ∣ E ( N F L O P s ) − F L O P s target ∣ ) ( 13 ) \operatorname{loss}_{r e g}=\log \left(\left|E\left(N_{F L O P s}\right)-F L O P s_{\text {target}}\right|\right) (13) lossreg=log(E(NFLOPs)FLOPstarget)(13)
为了让 E ( N F L O P s ) E(N_{FLOPs}) E(NFLOPs)严格小于 F L O P s t a r g e t FLOPs_{target} FLOPstarget但又对设置的目标FLOPs不太敏感, 所以给损失函数添加了单边距, 当满足 γ × F L O P s target ≤ \gamma \times F L O P s_{\text {target}} \leq γ×FLOPstarget E ( N F L O P s ) ≤ F L O P s target E\left(N_{F L O P s}\right) \leq F L O P s_{\text {target}} E(NFLOPs)FLOPstarget 时,损失函数会变成0. γ < 1 \gamma < 1 γ<1是公差比可由用户给定.

在更新权重时, F L O P s FLOPs FLOPs损失函数对权重没有影响, 权重损失为,
L o s s w e i g h t = l o s s c l s ( 14 ) Loss_{weight}= loss_{cls} (14) Lossweight=losscls(14)
l o s s c l s loss_{cls} losscls为分类任务的交叉熵损失函数.当更新结构参数时, 损失函数如下, λ reg \lambda_{\text {reg}} λreg是用于平衡两个损失的超参数.
L o s s a r c h = l o s s c l s + λ reg loss reg Loss_{arch}=loss_{cls}+\lambda_{\text {reg}} \text {loss}_{\text {reg}} Lossarch=losscls+λreglossreg

注意, 文中没有在结构参数中添加权重衰减, 由于当保留通道的概率接近0或者1时, 可学习参数 α α α的范数将变得非常大,这会使得他们趋近于0并进行优化.

训练

训练分为未剪枝网络权重更新和结构参数更新两个步骤, 这两个步骤会在训练过程中被迭代调用.

步骤1 未剪枝网络权重更新

由公式6的定义可知, 保留第 k k k个通道概率的也可以被认为是保留第前 k − 1 k-1 k1个通道的概率. 文中的方法相当于在网络前向传播更新结构参数时对所有子网络结构的软采样. 由于网络的通道的权重都是相等的, 这样就算进行马尔可夫建模但也不好提取子网络结构. 因此, 文中改进了’'三明治法则", 命名为"变体三明治法则", 并将该方法加入训练过程中, 使得在网络模型中通道组比他之后的通道组更重要, 让通道的权重变得不相等. 根据通道的权重排序, 可以选择要保留的通道从而选择出理论性能最好的子网络结构.

文中所提的与原始的"三明治法则"有两点不同: 1. 每一层随机采样的阈值不同; 2. 随机采样的阈值服从马尔可夫过程中的结构参数分布. 注意, 文中选择出现频率高的子网络结构作为最终的剪枝模型.
在这里插入图片描述

步骤2 结构参数更新

这个步骤中只更新结构参数, 由公式6可知, 结构参数与每个通道的输出结合, 使得梯度反向传播能够对结构参数进行优化. 通过下面的公式可以让梯度回传到 α \alpha α,
∂ L o s s ∂ α j ( i ) = ∑ k = 1 C o u t ∂ L o s s ∂ O k ( i ) ^ × ∂ O k ( i ) ^ α j ( i ) ( 16 ) \frac{\partial L o s s}{\partial \alpha_{j}^{(i)}}=\sum_{k=1}^{C_{o u t}} \frac{\partial L o s s}{\partial \hat{O_{k}^{(i)}}} \times \frac{\partial \hat{O_{k}^{(i)}}}{\alpha_{j}^{(i)}} (16) αj(i)Loss=k=1CoutOk(i)^Loss×αj(i)Ok(i)^(16)
∂ O k ( i ) ^ α j ( i ) = { 0 , k < j ∂ p k ∂ α j O k ( i ) ∏ r ∈ { r ∣ r ≠ j  and  r ≤ k } p r , k ≥ j ( 17 ) \frac{\partial \hat{O_{k}^{(i)}}}{\alpha_{j}^{(i)}}=\left\{\begin{array}{ll} 0 & , k<j \\ \frac{\partial p_{k}}{\partial \alpha_{j}} O_{k}^{(i)} \prod_{r \in\{r \mid r \neq j \text { and } r \leq k\}} p_{r} & , k \geq j \end{array}\right. (17) αj(i)Ok(i)^={0αjpkOk(i)r{rr=j and rk}pr,k<j,kj(17)
∂ p k ∂ α j = ( 1 − p k ) p k ( 18 ) \frac{\partial p_{k}}{\partial \alpha_{j}}=\left(1-p_{k}\right) p_{k} (18) αjpk=(1pk)pk(18)
为了进一步减小搜索空间, 文中将通道均匀划分成组(>10), 每个 α \alpha α参数代表一个组, 并设置每层的组数是相同的.

网络预热

在迭代调用步骤1和步骤2之前, 文中先运行几个epoch步骤1来进行预热, 这个过程中会通过随机初始化结构参数的马尔可夫过程进行采样得到子网络结构, 防止网络陷入局部最优.

剪枝模型采样

直接采样DS

直接采样时使用马尔可夫过程中的转移概率在网络的每一层进行独立的采样, 在这些采样中只保存满足FLOPs约束的网络结构.

期望采样ES

期望采样是使用公式8计算出来的结果作为每一层需要采样的通道数, 试验中 l o s s r e g loss_{reg} lossreg会被优化到0, 所以FLOPs是满足约束的.

DS得到的子网络结构的性能要略微优于ES, 但是由于该方法要测试多个子网络, 开销比较大, 所以文中评估采用的ES方式.下图为这两种方式的对比, 评价指标是Top1准确率, 1.5x-210M表示是最初始网络的1.5倍, FLOPs为210MB, 由图可以得出结论, ES和DS两种提取的子网络结构的性能相当.
在这里插入图片描述

实验结果

在这里插入图片描述

代码分析

先挖个坑

评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值