【剪枝】torch-pruning的基本使用

论文：DepGraph: Towards Any Structural Pruning
工程：https://github.com/VainF/Torch-Pruning
算法和库的使用介绍：CVPR 2023 | DepGraph 通用结构化剪枝

1 TP的简介

该算法介绍了DepGraph 如何建模结构化剪枝中的层依赖，实现任意结构的剪枝。对应实现的库为 torch-pruning。
本篇博客对作者的介绍做一个自己的梳理和记录。

• Torch-Pruning的简单介绍
  Torch-Pruning（TP）是一个结构化剪枝库，与现有框架（例如torch.nn.utils.prune）最大的区别在于，TP会物理地移除参数，同时自动裁剪其他依赖层。TP是一个纯 PyTorch 项目，实现了内置的计算图的追踪(Tracing)、依赖图(DepenednecyGraph, 见论文)、剪枝器等功能，同时支持 PyTorch 1.x 和 2.0 版本。
• 用 Torch-Pruning 剪枝的好处
  假设正在对一个卷积结构化剪枝，需要减去哪些内容，具体第几个卷积核、对应的偏置、BN中对应的维度、与其直接或间接相连的层的核的channel。我们要实现剪枝，需要对不同模型定制不同的代码实现。Torch-Pruning可让实现者跳脱出对层剪枝时最具体的操作，而关注于整体剪枝的设置。

2 TP的初尝试

---

2.1 初步尝试

以 ResNet-18 结构化剪枝为例，对【conv1】进行剪枝，同时处理对应的bn、紧临的卷积。

from torchvision.models import resnet18
import torch_pruning as tp
import torch

model = resnet18(pretrained=True).eval()
tp.prune_conv_out_channels(model.conv1, idxs=[0,1]) # 剪枝前两个通道
tp.prune_batchnorm_out_channels(model.bn1, idxs=[0,1]) # 尝试修复bn
tp.prune_conv_in_channels(model.layer1[0].conv1, idxs=[0,1]) # 尝试修复紧邻的conv
output = model(torch.randn(1,3,224,224)) # 尝试运行剪枝后的网络

会报错如下：问题出在残差结构上。残差的相加操作要求传入的两个tensor具有相同的空间尺寸，也就意味着剪枝后的Tensor通道数62和另一个tensor的通道数64不再匹配。

---

2.2 使用TP对 conv1进行剪枝

手动设置DependencyGraph是Torch-Pruning框架的底层算法，设计目标就是"自动寻找耦合层"，并自动化处理。
使用TP对ResNet-18的conv1进行剪枝，代码如下：

import torch
from torchvision.models import resnet18
import torch_pruning as tp

model = resnet18(pretrained=True).eval()

# 1. 构建依赖图
DG = tp.DependencyGraph()
DG.build_dependency(model, example_inputs=torch.randn(1,3,224,224))

# 2. 获取与model.conv1存在依赖的所有层，并指定需要剪枝的通道索引（此处我们剪枝第[2,6,9]个通道）
group = DG.get_pruning_group( model.conv1, tp.prune_conv_out_channels, idxs=[2, 6, 9] )
print(model, group)

# 3. 执行剪枝操作
if DG.check_pruning_group(group): # 避免将通道剪枝到0
    group.prune()
print(model, group)
output = model(torch.randn(1,3,224,224)) # 尝试运行剪枝后的网络

上述过程一共三步:

• 1 对网络进行依赖图构建;
• 2 选取需要剪枝的层，指定剪枝通道，获得分组group；这里的group，是所有与conv1相依赖的层。
• 3 执行剪枝操作，按组移除通道。那剪枝过程具体操作了哪些层呢？
   

左图为剪枝前的conv1 的group，右图为剪枝后的conv1 的group。
怎么去看这个group呢，在下图右侧进行了简单的标注，可以发现conv1的group都会进行剪枝，从而适应conv1的卷积核的维度发生的变化

左图为剪枝前的resnet结构部分，右图为剪枝后的resnet结构部分。

---

2.3 使用TP对网络中每个层进行剪枝

在实际实现时，我们希望是对整个网络结构进行剪枝，而非特定的某几层，这就涉及到如何不重复地遍历网络中所有分组的问题。
DepGraph提供了接口DG.get_all_groups来实现这以目标。该接口仅实现层的分组，并不会分辨通道的重要性。该接口包含两个参数

• ignored_layers：指定忽略 某些希望被剪枝的层。通常包括最后的分类层、以及报错的层（也可以使用其它正确的层进行替换）
• root_module_types：指定了每个组的起始层的类型。比如想剪枝所有的卷基层，而不想剪枝全连接层，只需要只传入对应的卷积类即可。
  值得注意的是，不同层可能出现在同一分组中，Depgraph会自动去除重复分组。
   

下面先提前设定好需要剪枝的通道，来展示DG.get_all_groups的使用：

import torch
import torch.nn as nn
import torch_pruning as tp
from torchvision.models import resnet18

model = resnet18(pretrained=True).eval()

# 1. 构建依赖图
DG = tp.DependencyGraph()
DG.build_dependency(model, example_inputs=torch.randn(1,3,224,224))

# 2. 获取与model.conv1存在依赖的所有层，并指定需要剪枝的通道索引（此处我们剪枝第[2,6,9]个通道）
Groups = DG.get_all_groups(ignored_layers=[model.conv1], root_module_types=[nn.Conv2d, nn.Linear])

# 3. 执行剪枝操作
for group in Groups:
    idxs = [2,4,6] # your pruning indices
    group.prune(idxs=idxs)
    print(group)

output = model(torch.randn(1,3,224,224)) # 尝试运行剪枝后的网络

但该段代码剪枝，在TP实际剪枝也是较少使用，这里是展示一个剪枝底层的基本操作。

3 TP对完整网络的剪枝

---

3.1 常用的结构化剪枝原理

结构化和非结构化剪枝方向，已发表的有较多的论文。但在工业上较常用的为结构化剪枝。实际中最常用的结构化剪枝方法：

• 利用权值进行filter剪枝：Pruning Filters for Efficient ConvNets
  在这张图中，我们可以找到两个卷积参数矩阵（Kernel Matrix）：第一个卷积层以 $x_i$ 作为输入，输出特征图 $x_{i+1}$；第二个卷积层以 $x_{i+1}$作为输入，生成特征图 $x_{i+2}$。
  在结构化剪枝中，这两个卷基层之间存在非常直观的依赖关系，即当我们调整第一层的输出通道时，第二个卷积层的输入通道也需要相应的进行调整，这使得蓝色高亮的参数需要同时被剪枝。
  
  此外，作者指出网络中可能存在更复杂的依赖，例如残差结构依赖：
  
• 利用bn进行剪枝：Learning Efficient Convolutional Networks through Network Slimming。
  
  BN会按通道对输入特征进行归一化，使得不同的特征处于比较接近的范围内。我们将缩放因子（从批量归一化层重用）与卷积层中的每个通道相关联。稀疏正则化在训练期间被施加在这些缩放因子上，以自动识别不重要的通道。缩放因子值较小（橙色）的通道将被修剪（左侧）。修剪后，我们获得紧凑模型（右侧），然后对其进行微调，以实现与正常训练的全网络相当（甚至更高）的精度。
  在任何一个网络中，BN的scale参数都具备一定的绝对值大小（也就是不会过小），这意味着各个通道都具有不可忽略的重要性。解决这类问题的一种有效方法是使用稀疏训练，通过对scale参数施加正则化项来稀疏化一部分通道。在slimming论文中，作者对scale参数施加了一个额外的L1正则化项，从而实现了这一过程。整个流程如下所示：稀疏训练–>剪枝–>微调
  

---

3.2 TP剪枝示例

网络中存在大量复杂依赖的情况下，如何进行剪枝呢？
【1】计算网络每个group中每层的重要性

• Torch-Puning 库内置了处理依赖的功能，并提供了可扩展的接口用于自定义剪枝器。
  类 tp.importance.Importance 要求我们实现一个非常简单的接口 __call__
  • 入参为一个 group，包含了多个相互耦合的层。
  • 输出为一个一维的重要性的得分向量，其含义是每个通道的重要性，因此他的维度和通道数量是相同的。
    由于输入的group通常会包含多个可剪枝层，因此我们首先对这些层进行独立的重要性计算，然后通过求平均值得到最终结果。
• Torch-Puning也提供了常用重要性评估策略：
  tp.importance.MagnitudeImportance(p=2)：p=2表示使用L2正则，对每个group中的每个层的权值，独立的计算重要性 。
  tp.importance.BNScaleImportance()：利用BN计算每个group中的每个层的权值的重要性
  tp.importance.GroupNormImportance()：与继承于MagnitudeImportance，且没做任何的添加和修改。
   

【2】对网络进行剪枝

• Torch-Pruning库定义了一个元剪枝器 tp.pruner.MetaPruner，能够完成除了重要性评估之外的所有工作。一般常在自定义的重要性评估后，执行剪枝时使用
• Torch-Puning也提供了常用的剪枝策略
  tp.pruner.MagnitudePruner()
tp.pruner.BNScalePruner()
tp.pruner.GroupNormPruner() Depgraph 提出的基于全局重要性的剪枝
   

【3】例子
为了增加难度，这里我们对一个DenseNet模型进行剪枝。
这里只展示了稀疏训练和微调使用的位置，仅剪枝部分能够有效跑通。

import torch
import torch.nn as nn
import torch_pruning as tp
from torchvision.models import densenet121

model = densenet121(pretrained=True)
example_inputs = torch.randn(1, 3, 224, 224)

# 1. 使用我们上述定义的重要性评估
# imp = tp.importance.MagnitudeImportance(p=2)
# imp = tp.importance.BNScaleImportance()
imp = tp.importance.GroupNormImportance()

# 2. 忽略无需剪枝的层，例如最后的分类层
ignored_layers = []
for m in model.modules():
    if isinstance(m, torch.nn.Linear) and m.out_features == 1000:
        ignored_layers.append(m) # DO NOT prune the final classifier!

# 3. 初始化剪枝器
iterative_steps = 5 # 迭代式剪枝，重复5次Pruning-Finetuning的循环完成剪枝。  
# pruner = tp.pruner.MagnitudePruner(
# pruner = tp.pruner.BNScalePruner(
pruner = tp.pruner.GroupNormPruner(
    model,
    example_inputs, # 用于分析依赖的伪输入
    importance=imp, # 重要性评估指标
    iterative_steps=iterative_steps, # 迭代剪枝，设为1则一次性完成剪枝
    ch_sparsity=0.5, # 目标稀疏性，这里我们移除50%的通道 ResNet18 = {64, 128, 256, 512} => ResNet18_Half = {32, 64, 128, 256}
    ignored_layers=ignored_layers, # 忽略掉最后的分类层
)

# 4. 稀疏训练（为了节省时间我们假装在训练，实际应用时只需要在optimizer.step前插入regularize即可）
for _ in range(100):
    pass
    # optimizer.zero_grad() 
    # ...
    # loss.backward()
    # pruner.regularize(model, reg=1e-5) # <== 插入该行进行稀疏化
    # optimizer.step()
    
# 4. Pruning-Finetuning的循环
base_macs, base_nparams = tp.utils.count_ops_and_params(model, example_inputs)
for i in range(iterative_steps):
    pruner.step() # 执行裁剪，本例子中我们每次会裁剪10%，共执行5次，最终稀疏度为50%
    macs, nparams = tp.utils.count_ops_and_params(model, example_inputs)
    print("  Iter %d/%d, Params: %.2f M => %.2f M" % (i+1, iterative_steps, base_nparams / 1e6, nparams / 1e6))
    print("  Iter %d/%d, MACs: %.2f G => %.2f G"% (i+1, iterative_steps, base_macs / 1e9, macs / 1e9))
    # finetune your model here
    # finetune(model)
    # ...
print(model)