论文阅读笔记：DepGraph: Towards Any Structural Pruning

论文：https://arxiv.org/pdf/2301.12900

代码：https://github.com/VainF/Torch-Pruning

1 背景

主流的剪枝方法可以被分为结构化剪枝和非结构化剪枝。两者的核心区别在于，结构化剪枝通过物理地移除分组参数来改变神经网络的结构，而非结构化剪枝在不修改网络结构的情况下，对部分权重进行归零。相比于非结构化剪枝，结构化剪枝不太以来特定的AI加速器或者软件来减少内存消耗和计算成本，从而在实际应用中找到更广泛的应用领域。

模型剪枝根据粒度的不同，至少可以粗分为4个粒度：

1.细粒度剪枝(fine-grained)：即对连接或者神经元进行剪枝，它是粒度最小的剪枝。
2.向量剪枝(vector-level)：它相对于细粒度剪枝粒度更大，属于对卷积核内部(intra-kernel)的剪枝。
3.核剪枝(kernel-level)：即去除某个卷积核，它将丢弃对输入通道中对应计算通道的响应。
4.滤波器剪枝(Filter-level)：对整个卷积核组进行剪枝，会造成推理过程中输出特征通道数的改变。

细粒度剪枝(fine-grained)，向量剪枝(vector-level)，核剪枝(kernel-level)方法在参数量与模型性能之间取得了一定的平衡，但是网络的拓扑结构本身发生了变化，需要专门的算法设计来支持这种稀疏的运算，被称之为非结构化剪枝。
而滤波器剪枝(Filter-level)只改变了网络中的滤波器组和特征通道数目，所获得的模型不需要专门的算法设计就能够运行，被称为结构化剪枝。除此之外还有对整个网络层的剪枝，它可以被看作是滤波器剪枝(Filter-level)的变种，即所有的滤波器都丢弃。

参考：

https://www.jianshu.com/p/8e1209d3127a

https://zhuanlan.zhihu.com/p/613899881

https://blog.csdn.net/baidu_38262850/article/details/125101878

然而，结构化剪枝本身的性质使其本身成为一项具有挑战性的任务，特别是对于具有耦合和复杂内部结构的现代深度神经网络。其原理在于，深度神经网络建立在卷积、归一化或激活等大量基本模块的基础上，而这些模块，无论参数化与否，都通过错综复杂的连接内在耦合。因此，即使试图从CNN中只删除一个通道，都必须同时照顾到它对所有层的相互依赖关系，否则最终会得到一个破碎的网络。准确地说，残差链接要求两个卷积层的输出共享相同的通道，从而迫使他们一起被剪枝。如图1所示，在其他架构，如transformers，RNNs和GNNs上也是如此。

现有的结构化剪枝方法在很大程度上依赖于逐个案例分析来处理网络中的依赖关系，尽管取得了较好的结构，但这种方式是费力且不具有推广性的。

在本文中，作者为任意结构化剪枝寻求了一个通用方案，其中任意网络架构上的剪枝都是自动执行的。方法的核心是估计依赖图（Dependency Graph, DepGraph），它显式地建模了神经网络中成对层之间的相互依赖关系。引入DepGraph进行结构化剪枝的动机源于观察到，一层的结构化剪枝有效地"触发"了相邻层的剪枝，这进一步导致了类似{ $B{N}_2\leftarrow Con{v}_2\to B{N}_1\to Con{v}_1$ }的链式效应，如图1 ( a )所示。

另外，值得注意的是，在结构化剪枝中，同时对分组层进行剪枝，期望同一组中所有被移除的参数始终保持不重要，但由于其他层参数的纠缠，单层中参数的重要性不会被正确的显示出来。为了解决这个问题，作者充分利用DepGraph提供的依赖建模的综合能力，设计了一个“组级别”的重要性准则，该准则学习组内一致的稀疏性，从而可以安全地删除那些零化的组，而不会带来太大的性能损失。

有很多剪枝方法的思路都是在指定范围内用某种方式表示权重，并将其进行从小到大的排序，将小于某一滤值的元素置为0，然后将是0的剪枝掉。

如用于结构化剪枝算法的L2Filter Pruner，使用L2范数来对权重矩阵进行排序，并删除排列最末的 k%。权重越稀疏，排名越靠后。

例如本文代码：

https://github.com/VainF/Torch-Pruning/blob/master/torch_pruning/pruner/importance.py#L194

local_imp = w.abs().pow(self.p).sum(1) # self.p = 2

https://github.com/VainF/Torch-Pruning/blob/master/torch_pruning/pruner/algorithms/metapruner.py#L532

_pruning_indices = (imp <= thres).nonzero().view(-1)
imp_argsort = torch.argsort(imp)
if len(_pruning_indices)>0 and self.round_to: 
    n_pruned = len(_pruning_indices)
    current_channels = get_channel_fn(module)
    n_pruned = self._round_to(n_pruned, current_channels, self.round_to)
    pruning_indices = imp_argsort[:n_pruned]
pruning_indices.append(_pruning_indices)

2 创新点

针对任何结构化剪枝的通用剪枝方案，称为依赖图( Dependency Graph，DepGraph )，它允许自动参数分组，并有效地提高了结构化剪枝在各种网络架构上的泛化性，包括CNN、RNN、GNN和Vision Transformers。

3 方法

本文工作专注于在参数依赖的限制下对任何神经网络进行结构化剪枝。

如图2（a）所示，从3个全连接层组成的神经网络开始，分别用二维权重矩阵 $w_l$、$w_{l+1}$ 和 $w_{l+2}$ 进行参数化。这种简单的神经网络可以通过去除神经元进行结构修剪。在这种情况下，很容易发现参数之间存在某种依赖关系，记为 $w_l\iff w_{l+1}$，这迫使 $w_l$ 和 $w_{l+1}$ 同时被剪枝。具体来说，为了修剪连接 $w_l$ 和 $w_{l+1}$ 的第 $k$ 个神经元，$w_l[k,:]$ 和 $w_{l+1}[:,k]$ 都将被移除。

如图2（b-d）所示，这种依赖关系有很多种，以逐个案例的方式人工分析所有依赖关系是很难的，更不用说简单的依赖关系可以嵌套或组合形成更复杂的模式。为了解决结构化剪枝中的依赖问题，本文引入了依赖图( Dependency Graph )，为依赖建模提供了一种通用的、全自动的机制。

4 模块

4.1 分组

为了实现结构化剪枝，首先需要根据层的相关性对层进行分组。目标是找到一个分组矩阵 $G\in R^{L\times L}$，其中L表示待剪枝网络的层数，$G_{ij}=1$ 表示第 $i$ 层和第 $j$ 层之间存在依赖关系。为了方便起见，令 $Diag(G)=1^{1\times L}$以实现自依赖。利用分组矩阵，可以很直观地找到相关性到第 $i$ 层的所有耦合层，记为 $g(i)$：

然而，由于现代深度网络可能由数以千计的具有复杂连接的层组成，从而产生了一个庞大而复杂的分组矩阵G。这个矩阵中， $G_{ij}$ 不仅由第 $i$ 层和第 $j$ 层决定，还受到他们之间的中间层的影响。因此这种非局部和隐式的关系在大多数情况下不能用简单地规则来处理。为了克服这一挑战，作者不直接估计分组矩阵G，而是提出了一种等价但易于估计的依赖建模方法，即依赖图（Dependency graph）。从依赖图中高效的导出G。

4.2 依赖图

首先考虑一个群 g = { w 1 , w 2 , w 3 } g=\{w_1, w_2, w_3\} g={w1​,w2​,w3​}，它具有依赖关系 $w_1\iff w_2,w_2\iff w_3,w_1\iff w_3$。在仔细观察这种依赖关系建模后，可以发现存在一些冗余。哟如，从 $w_1\iff w_2,w_2\iff w_3$可以通过递归过程导出依赖关系 $w_1\iff w_3$。首先以 $w_1$ 为起点，考察它对其它层的依赖性，例如 $w_1\iff w_2$。然后，$w_2$ 为递归扩展依赖关系提供了新的起点，依赖关系触发 $w_2\iff w_3$。这个递归过程最终以一个传递关系 $w_1\iff w_2\iff w_3$结束。在这种情况下，只需要两个依赖关系来描述组g中的关系。同样地，分组矩阵G对于依赖建模也是冗余的，因此可以压缩为具有更少边数的更紧凑的形式，同时保留相同的信息。作者证明了一个新的度量相邻层之间局部相关性的图 D，称为依赖图，可以对分组矩阵G进行有效的约减。依赖图与G的不同之处在于，它只记录了具有直接连接的相邻层之间依赖关系。图D可以看成是图G的可迁约化，它包含了图G相同的顶点但具有尽可能少的边。形式上，构造D使得对所有 $G_{ij}=1$ 在D中存在一条顶点 $i$ 和 $j$ 之间的路径。因此 $G_{ij}$ 可以通过检查D中顶点 $i$ 和 $j$ 之间是否存在一条路径得到。

4.3 网络分解

然而，作者在实际构建依赖图时发现可能存在问题，因为一些基本层（如全连接层）可能由两种不同的剪枝方案，如之前讨论的 $w[k,:]$ 和 $w[:,k]$ ，他们分别压缩了输入和输出的维度。此外，网络还包括跳跃连接等非参数化操作，这也会影响层与层之间的依赖关系。为了解决这些问题，作者将网络 $F(x;w)$ 分解成更精细和更基本的组件，记为 F = { f 1 , f 2 , … , f L } F=\{f_1, f_2, …, f_L\} F={f1​,f2​,…,fL​} ，其中每个组件 $f$ 指的是参数化的层（如卷积）或非参数化的操作（如残差相加）。取代在层级别上建模关系，作者专注于层的输入和输出之间的依赖关系。具体的，将分量 $f_i$ 的输入和输出分别记为 $f_i^{-}$ 和 $f_i^{+}$。对于任意网络，最终的分解可以形式化为 $F=f_1^{-},f_1^{+},\dots ,f_L^{-},f_L^{+}$。这种表示方法有利于更容易的依赖建模，并且允许对同一层采用不同的剪枝方案。

4.4 依赖建模

将神经网络重新绘制为式（2），其中可以识别出两种主要类型的依赖关系，即层间依赖和层内依赖，如下图所示：
#pic_center

其中，↔ 表示相邻两层之间的连通性，对这两种依赖关系的检验可以得到简单但通用的依赖关系建模规则：

（1）层间依赖：在连接层$f_i^{-}\leftrightarrow f_j^{+}$中一致产生一个依赖 $f_i^{-}\iff f_j^{+}$

（2）层内依赖：如果 $f_i^{-}$ 和 $f_i^{+}$具有相同的剪枝方案，则存在依赖 $f_i^{-}\iff f_i^{+}$ ,记为 $sch(f_i^{-})=sch(f_i^{+})$ 。

首先，如果已知网络的拓扑结构，可以很容易地估计层间依赖。对于 $f_i^{-}\leftrightarrow f_j^{+}$ 的连接层，由于 $f_i^{-}$ 和 $f_j+$ 对应网络相同的中间特征，因此依赖关系是始终存在。

其次，层内依赖要求对单层的输入和和输出同时进行剪枝，许多网络层满足这个条件，例如BN层，其输入和输出共享相同的剪枝方案，记作 $sch(f_i^{-})=sch(f_i^{+})$ ，将同时进行剪枝，如图3所示。相反，卷积等层对其输入和输出具有不同的剪枝方案，如图3所示，$w[:,k,:,:]\ne w[k,:,:,:]$ ，导致 $sch(f_i^{-})\ne sch(f_i^{+})$，在这种情况下，卷积层的输入和输出之间不存在依赖。

给定上述规则，可以形式的建立如下依赖建模：

其中，$\wedge$ 和 $\vee$ 表示逻辑与和或，$1$ 是一个条件成立则返回True的指示器函数。第一项考察由网络连接引起的层间依赖，而第二项考察由层输入和输出之间的共享剪枝方案引入的层内依赖。值得注意的是，DepGraph 是一个对称矩阵，即 $D(f_i^{-},f_j^{+})=D(f_j^{+},f_i^{-})$。因此可以检查所有的输入输出对来估计依赖图。图3中可视化了具有残差连接的CNN块的依赖图。

算法1和算法2总结了依赖建模和分组的算法。

4.4 组级剪枝

评估分组参数的重要性对剪枝提出了重大挑战，因为它涉及了多个耦合层。这里，作者利用一个简单的基于范数的准则来建立一个实用的组级剪枝方法。给定一个参数组 g = { w 1 , w 2 , … w ∣ g ∣ } g=\{w_1,w_2,…w_{|g|}\} g={w1​,w2​,…w∣g∣​}，现有的L2范数重要性准则 $I(w)=||w|{|}_2$可以对每个 $w\in g$ 产生独立的评分。估计群体重要性的一个自然方法是计算一个综合得分 $I(g)={\sum }_{w\in g}I(w)$ ，不幸的是，不同层独立评估的重要性得分很可能是不可加的，因为分布和量级的差异而因此没有意义。为了使这种简单的聚合方式适用于重要性估计，作者提出了一种稀疏训练方法来稀疏化组级别的参数，如图4（c）所示，以便那些被零化的组可以安全的从网络中移除。

具体来所，对于每一个具有 $K$ 个可剪枝维度的参数 $w$，作者引入了一个简单的用于稀疏训练的正则化项，定义为：

其中， $I_{g,k}={\sum }_{w\in g}||w[k]|{|}_2^2$ 表示 第 $k$ 个可剪枝维度的重要性，${\gamma }_k$ 指的是应用于这些参数的缩放强度。作者使用一个可控的指数策略来确定 ${\gamma }_k$ 如下：

其中，收缩强度 ${\alpha }_k$ 采用归一化分数来控制，其变化范围为 $[2^0,2^{\alpha }]$。在实验中 $\alpha =4$。在稀疏训练之后，作者进一步使用了简单的相对得分 I ^ g , k = N ⋅ I g , k / ∑ { T o p N ( I g ) } \hat{I}_{g,k}=N·I_{g,k}/\sum\{TopN(I_g)\} I^g,k​=N⋅Ig,k​/∑{TopN(Ig​)},用于识别和去除不重要的参数。

这里正则化的目的就是为了让成对的参数重要性都趋向于0。

稀疏训练过程如下:

for epoch in range(epochs):
    model.train()
    pruner.update_regularizer() # <== initialize regularizer
    for i, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        out = model(data)
        loss = F.cross_entropy(out, target)
        loss.backward() # after loss.backward()
        pruner.regularize(model) # <== for sparse training
        optimizer.step() # before optimizer.step()

以out_channel为例，正则化过程如下（代码见https://github.com/VainF/Torch-Pruning/blob/master/torch_pruning/pruner/algorithms/group_norm_pruner.py#L126）：

alpha=2**4
imp = self.estimate_importance(group).sqrt()
gamma = alpha**((imp.max() - imp) / (imp.max() - imp.min()))
for i, (dep, idxs) in enumerate(group):
    layer = dep.target.module
    prune_fn = dep.handler
    if prune_fn in [
        function.prune_conv_out_channels,
        function.prune_linear_out_channels,
    ]:
        root_idxs = group[i].root_idxs
        _gamma = torch.index_select(gamma, 0, torch.tensor(root_idxs, device=gamma.device))

        w = layer.weight.data[idxs]
        g = w * _gamma.view( -1, *([1]*(len(w.shape)-1)) )
        layer.weight.grad.data[idxs]+=self.reg * g  # self.reg = 1e-4

个人理解是通过以权重的反方向和损失一起更新梯度，不重要的权重因为没有loss的梯度，所以会朝权重的反方向更新，直到趋向于0，重要的权重因为会有loss的梯度，所以不会趋向于0。

5 效果

在CIFAR10上剪枝了一个ResNet56，在CIFAR100上剪枝了VGG19，效果如表1。

一致稀疏性对于结构剪枝非常重要，因为它迫使所有被剪枝的参数一致不重要。如图5中，将图4（c）和（b）中一致和不一致策略学习到分组参数的范数可视化，对应 w/和 w/o。容易发现，本文提出的方法在组水平上产生了强稀疏性，这有利于识别不重要的参数。（被剪枝后值变小）

为了进一步验证分组的有效性，分别测试：

（1）不分组：在单个卷积层上独立进行稀疏学习和重要性评估

（2）仅卷积分组：组内卷积层以一致的方式进行稀疏化。

（3）全分组：组内的所有参数层，包括卷积，BN，全连接等都一致稀疏化

结果如表2。

层稀疏度也是剪枝的一个重要因素，它决定了剪枝后神经网络的最终结构。表2提供了一些关于层稀疏性的结果。这项工作主要关注两种类型的稀疏性，即均匀稀疏性和学习稀疏性。在具有均匀稀疏性的情况下，将对不同层施加相同的剪枝比例，假设冗余通过网络均匀分布。然而，图5中先前的实验表明，不同的层并不等同于可剪枝的。在大多数情况下，学习到的稀疏性优于均匀稀疏性，尽管有时它可能会过度剪枝某些层，从而导致精度下降。

同时表2也证明了本文框架的可推广性。

由于分组参数过程十分复杂，对大型神经网络进行剪枝是一个相当大的挑战。然而，通过使用DepGraph，可以毫不费力地获得所有耦合组。在图6中提供了DenseNet121 ，ResNet18和Vision Transformers 的DepGraph D和衍生分组矩阵G的可视化。分组矩阵由算法2中的DepGraph导出，其中G[ i , j] = 1表示第i层与第j层属于同一组。DenseNet121在同一稠密块内的层与层之间表现出很强的相关性，导致结构修剪时产生较大的组。在处理复杂网络时，所提出的依赖图被证明在处理复杂网络时很有帮助，因为手动分析此类网络中的所有依赖关系确实是一项困难的任务。

表3为剪枝结果在ImageNet上的效果，包括ResNet, DenseNet, MobileNet, ResNeXt, and Vision Transformers。

表4在其他网络结构上的泛化性如表4。