YOLOv5系列(四十六) 模型剪枝与Pytorch剪枝策略

在yolov5项目中的torch_utils.py文件下，有prune这个函数，用来实现模型的剪枝处理。对模型裁剪，模型剪枝这方面之前没有接触到，这里用这篇笔记来学习记录一下这方面内容。

---

1. Yolov5代码实现

在yolov5项目中提供了两个函数：sparsity与prune，前者可以返回模型的稀疏性，后者实现对模型的裁剪处理。

• 原理介绍：

对于模型稀疏性的判断，其实现的思路是遍历每一层模块的参数量，当当前层的参数值非0，表示当前的神经元是被激活的在网络的前向传播中是使用到的；而当当前层的参数值为0时，表示这个为0的参数所控制的神经元是没有被激活的，也就是不参与到网络的训练上，或者说是被dropout掉了。通过这样的一种原理，将失活的神经元与整个模型的全部参数的比值，就可以判断当前模型的稀疏性。因为当失活的神经元占比越大，表示这个模型的参数是小的，训练量也是比较小的，从而实现轻量化模型的想法，加速了前向传播与后向传播的速度，加速了整个训练过程。

但是，尽管神经元被失活没有使用上，其在网络中所初始化的位置还是被保留下来的。也就说，所以这个参数在训练网络的过程中全程都没有使用上，但是模型初始化后的参数量大小是没有被改变的（通过后序实验thop计算得知）。只能说，这里yolov5实现的模型剪枝将更多的参数置0，不参与网络更新，不被激活。

• 剪枝代码与稀疏性测试代码如下：

def sparsity(model):
    # Return global model sparsity
    # a用来统计使用的神经元的个数, 也就是参数量个数
    # b用来统计没有使用到的神经元个数, 也就是参数为0的个数
    a, b = 0., 0.
    for p in model.parameters():
        a += p.numel()        # numel()返回数组A中元素的数量
        b += (p == 0).sum()   # 参数为0 表示没有使用到这个神经元参数
    # b / a 即可以反应模型的稀疏程度
    return b / a


def prune(model, amount=0.3):
    # Prune model to requested global sparsity
    import torch.nn.utils.prune as prune
    print('Pruning model... ', end='')
    # 对模型中的nn.Conv2d参数进行修剪
    for name, m in model.named_modules():
        if isinstance(m, nn.Conv2d):
            # 这里会对模块原来的weight构建两个缓存去, 一个是weight_orig(原参数), 另外一个是weight_mask(原参数的掩码)
            # weight_mask掩码参数有0/1构成, 1表示当前神经元不修剪, 0表示修剪当前神经元
            prune.l1_unstructured(m, name='weight', amount=amount)  # prune
            # 将name+'_orig'与name+'_mask'从参数列表中删除, 也就是将掩码mask作用于原参数上
            # 使name保持永久修剪, 同时去除参数的前向传播钩子(就是不需要前向传播)
            prune.remove(m, 'weight')  # make permanent

    # 测试模型的稀疏性
    print(' %.3g global sparsity' % sparsity(model))

• 测试代码：

# 功能: 测试模型参数
def model_parms(model):
    from thop import profile

    input = torch.randn(1, 3, 640, 640)
    flops, params = profile(model, inputs=(input,))
    print('flops:{}G'.format(flops / 1e9))
    print('params:{}M'.format(params / 1e6))


# 功能: 测试模型剪枝
def model_prune():
    from utils.torch_utils import prune, sparsity, model_info
    from thop import profile

    model = load_model()
    model_parms(model)
    # model_info(model, verbose=True)

    result = sparsity(model)
    print("prune before:{}".format(result))

    prune(model)
    result = sparsity(model)
    print("prune after:{}".format(result))
    model_parms(model)
    # model_info(model, verbose=True)

• 输出结果：

# 剪枝前的结果:(浮点计算量, 模型参数量, 模型稀疏性)
flops:7.9331296G
params:7.02772M
prune before:2.418992153252475e-06

# 剪枝后的结果:(浮点计算量, 模型参数量, 模型稀疏性)
Pruning model...  0.299 global sparsity
flops:7.9331296G
params:7.02772M
prune after:0.29918551445007324

分析：可以看见，剪枝后的模型参数为0的比例增加，失活的神经元比例增加，模型的稀疏性增加。但是模型的参数量与浮点计算量的大小没有改变。

---

2. 模型剪枝介绍

详细内容见参考资料1,2

2.1 剪枝方法简介

剪枝就是通过去除网络中冗余的channels,filters, neurons, or layers以得到一个更轻量级的网络，同时不影响性能。网络剪枝的步骤神经网络中的一些权重和神经元是可以被剪枝的，这是因为这些权重可能为零或者神经元的输出大多数时候为零，表明这些权重或神经元是冗余的。

模型剪枝并不是一个新的概念，其实我们从学习深度学习的第一天起就接触过，Dropout和DropConnect代表着非常经典的模型剪枝技术，模型剪枝不仅仅只有对神经元的剪枝和对权重连接的剪枝，根据粒度的不同，至少可以粗分为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)的变种，即所有的滤波器都丢弃。

深度学习网络模型从卷积层到全连接层存在着大量冗余的参数，大量神经元激活值趋近于0，将这些神经元去除后可以表现出同样的模型表达能力，这种情况被称为过参数化，而对应的技术则被称为模型剪枝。

• 剪枝步骤

网络剪枝的过程主要分以下几步：

①训练网络；
②评估权重和神经元的重要性：可以用L1、L2来评估权重的重要性，用不是0的次数来衡量神经元的重要性；
③对权重或者神经元的重要性进行排序然后移除不重要的权重或神经元；
④移除部分权重或者神经元后网络的准确率会受到一些损伤，因此我们要进行微调，也就是使用原来的训练数据更新一下参数，往往就可以复原回来；
⑤为了不会使剪枝造成模型效果的过大损伤，我们每次都不会一次性剪掉太多的权重或神经元，因此这个过程需要迭代，也就是说剪枝且微调一次后如果剪枝后的模型大小还不令人满意就回到步骤后迭代上述过程直到满意为止

2.2 剪枝合理性解释

在我之前的一篇笔记中，记录过曾经上李宏毅老师关于为什么可以进行网络剪枝的解释，笔记见：学习笔记——神经网络压缩

现在有一个问题，既然大的网络需要剪枝处理，那么为什么一开始就不训练一个小的网络呢？一个可能的感受是，小的网络比较难以去训练，然后大的网络比较容易去优化。一般来说，训练过程中存在鞍点或者局部最优解的问题。而如果网络够大，那么这种情况就不会太严重。现在有足够多的文献可以证明，只要网络够大够深，就可以用gradient descent直接找到全局最优解。所以训练一个大的网络，再剪枝处理是比较好的。解释这一想象的一个假设是大乐透假设（Lottery Ticket Hypothesis）

• 大乐透假设（Lottery Ticket Hypothesis）

在一个大的网络结构中，其实可以看成是有很多个小的网络组成的，每一个小的网络就有可能是一种初始化的参数，而这些小的网络有些可以train起来，而有些会train不起来。所以大的网络结构中容易训练的原因可能是，其中这么多个小网络，只有有一个可以train起来了，那么大的网络就可以train起来了。paper链接：https://arxiv.org/abs/1803.03635

这是因为剪枝做成了网络结构的不规则，因此难以用GPU进行加速。在进行实验需要使用weight pruning时可以使用将被剪枝的权重设置成0的方法，也就是掩码设计的方法。

---

3. Pytorch剪枝策略

官方文档：https://pytorch.org/docs/stable/nn.html#utilities

剪枝可以在单层（a single layer），多层(multiple layer)或整个模型(an entire model)中进行。主要的剪枝策略如下所示：（详细见参考资料3）

• 类方法实现：
  prune.Identity 实用剪枝方法，不剪枝任何单元，但生成带有掩码的剪枝参数化。
  prune.RandomUnstructured 随机修剪（当前未修剪的）张量中的单元。
  prune.L1Unstructured 通过将具有最低 L1 范数的单元归零来修剪（当前未修剪）张量中的单元。
  prune.RandomStructured 随机修剪张量中的整个（当前未修剪的）通道。
  prune.LnStructured 根据 Ln范数在张量中修剪整个（当前未修剪的）通道。
• 函数方法实现：
  prune.identity 将修剪重新参数化应用于与调用的参数对应的张量name，module而不实际修剪任何单位。
  prune.random_unstructured 通过删除随机选择的指定的（当前未修剪的）单元来修剪与调用name的参数相对应的张量。
  prune.l1_unstructured 通过删除具有最低 L1 范数的指定数量的（当前未修剪的）单元来修剪与调用name的参数相对应的张量。
  prune.random_structured 通过沿随机选择的指定删除指定的（当前未修剪的）通道来修剪与调用name的参数相对应的张量。
  prune.ln_structured 通过沿着具有最低 L范数的指定通道移除指定的（当前未修剪的）通道，修剪与调用name的参数相对应的张量。
  prune.global_unstructured parameters通过应用指定的来全局修剪与所有参数对应的张量pruning_method。
  prune.custom_from_mask name通过在 中module应用预先计算的掩码来修剪与调用的参数相对应的张量mask。
  prune.remove 从模块中删除修剪重新参数化，从前向钩子中删除修剪方法。

以下内容详细见参考资料4.

• 构建卷积网络：这里最简单的卷积网络，LeNet5为例：

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        # 1 input image channel, 6 output channels, 3x3 square conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 6 * 6, 120)  # 5x5 image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

3.1 局部剪枝

局部剪枝 主要是对个别模块（如某一层，模块等）进行剪枝操作

import torch.nn.utils.prune as prune

# 只对conv1层的权重进行随机剪枝操作
prune.random_unstructured(model.conv1, name='weight', amount=0.25)
# 也可对偏置进行剪枝操作
prune.random_unstructured(model.conv1, name='bias', amount=0.25)

3.2 迭代剪枝

剪枝可以迭代式运用，因此在实际应用中可以针对不同维度用不同的方法进行剪枝操作。
同时剪枝不仅可应用于模块中，也可以对参数进行剪枝的。

示例：对网络结构中的卷积层进行权重剪枝，对线性连接层进行不同的权重比值剪枝操作

model = LeNet5().to(device)

for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        # Prune all 2D convolutional layers by 30%
        prune.random_unstructured(module,name='weight', amount=0.3)
    # Prune all linear layers by 50%.
    elif isinstance(module, torch.nn.Linear):
        prune.random_unstructured(module,  name='weight', amount=0.5)

3.3 全局剪枝

全局剪枝就是对整个模型进行剪枝操作。

示例：对模型的参数进行25%剪枝操作

model = LeNet5().to(device)
parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.fc1, 'weight'),
    (model.fc2, 'weight'),
    (model.fc3, 'weight'),
)

# prune 25% of all the parameters in the entire model
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.25
)

3.4 自定义剪枝

如果找不到适合您需求的修剪方法，您可以创建自己的修剪方法。 为此，请从 torch.nn.utils.prune 中提供的 BasePruningMethod 类创建一个子类。
您将需要编写自己的 _ init _() 构造函数和 compute_mask() 方法来描述您的修剪方法如何计算掩码。 此外，您需要指定修剪的类型（结构化、非结构化或全局）。

If you can’t find a pruning method that suits your needs, you can create your own pruning method. To do so, create a subclass from the BasePruningMethod class provided in torch.nn.utils.prune.
you will need to write your own init() constructor and compute_mask() method to describe how your pruning method computes the mask. In addition, you’ll need to specify the type of pruning (structured, unstructured, or global).

示例：以下自定义了一个剪枝策略，就是间隔地将掩码赋值为0，mask.view(-1)[::2] = 0

class MyPruningMethod(prune.BasePruningMethod):
    PRUNING_TYPE = 'unstructured'
    def compute_mask(self, t, default_mask):
        mask = default_mask.clone()
        mask.view(-1)[::2] = 0
        return mask

def my_unstructured(module, name):
    MyPruningMethod.apply(module, name)
    return module

# 对模型进行自定义剪枝操作
model = LeNet5().to(device)
my_unstructured(model.fc1, name='bias')

查看剪枝后的编制属性与缓存区结果：

# 查看缓存区结果
print(list(model.fc1.named_buffers()))
# 输出：
[('bias_mask', tensor([0., 1., 0., 1., 0., 1., 0.,