YOLOv5的Tricks | 【Trick9】模型剪枝处理与Pytorch实现的剪枝策略
如有错误，恳请指出。
文章目录
- 1. Yolov5代码实现
- 2. 模型剪枝介绍
-
- 2.1 剪枝方法简介
- 2.2 剪枝合理性解释
- 3. Pytorch剪枝策略
-
- 3.1 局部剪枝
- 3.2 迭代剪枝
- 3.3 全局剪枝
- 3.4 自定义剪枝
- 4. Pytorch剪枝测试
-
- 4.1 剪枝前的参数
- 4.2 对卷积层权重剪枝
- 4.3 对卷积层偏置剪枝
- 4.4 对卷积层权重删除修剪
- 4.5 对卷积层偏置删除修剪
- 4.6 剪枝对参数量与浮点计算量的变化
在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., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.,
0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.,
0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.,
0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.,
0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.,
0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.,
0., 1., 0., 1., 0., 1., 0., 1., 0., 1., 0., 1.]))]
# 查看属性结果
print(model.fc1.bias)
# 输出:
tensor([ 0.0000, -0.0242, -0.0000, 0.0286, -0.0000, 0.0308, -0.0000, -0.0145,
0.0000, -0.0377, 0.0000, 0.0256, 0.0000, -0.0133, 0.0000, 0.0201,
0.0000, 0.0028, 0.0000, -0.0362, -0.0000, -0.0119, 0.0000, 0.0405,
-0.0000, -0.0305, -0.0000, 0.0322, 0.0000, -0.0379, 0.0000, 0.0219,
0.0000, -0.0133, 0.0000, 0.0224, -0.0000, -0.0180, 0.0000, 0.0015,
0.0000, 0.0129, 0.0000, 0.0166, 0.0000, 0.0306, 0.0000, -0.0251,
-0.0000, -0.0309, 0.0000, -0.0002, -0.0000, 0.0115, 0.0000, 0.0412,
0.0000, -0.0337, -0.0000, -0.0362, -0.0000, 0.0347, -0.0000, -0.0321,
0.0000, -0.0399, 0.0000, 0.0241, -0.0000, -0.0186, -0.0000, 0.0114,
0.0000, -0.0283, -0.0000, 0.0292, -0.0000, -0.0048, 0.0000, -0.0317,
-0.0000, 0.0176, -0.0000, 0.0135, 0.0000, 0.0222, -0.0000, -0.0249,
0.0000, 0.0299, -0.0000, 0.0278, -0.0000, 0.0166, -0.0000, -0.0232,
0.0000, 0.0028, -0.0000, -0.0310, -0.0000, -0.0015, 0.0000, 0.0247,
0.0000, 0.0283, -0.0000, -0.0093, -0.0000, 0.0262, -0.0000, -0.0153,
0.0000, -0.0087, 0.0000, 0.0101, 0.0000, 0.0340, -0.0000, 0.0280],
grad_fn=<MulBackward0>)
分析：可以看见，由于自定义的掩码操作是间隔地将掩码赋值为0，所以这里的属性值与掩码缓存都是间隔数值就为0，符合我们的自定义操作想法。
4. Pytorch剪枝测试
具体见参考资料5.
这里使用LeNet5来作为模型剪枝的例子，来查看具体的一个卷积层模型去权重weight与偏置bias的变化。
# 查看模型各层结构
model = LeNet5()
for name, param in model.named_parameters():
print(name, param.dtype)
# 输出：
conv1.weight torch.float32
conv1.bias torch.float32
conv2.weight torch.float32
conv2.bias torch.float32
fc1.weight torch.float32
fc1.bias torch.float32
fc2.weight torch.float32
fc2.bias torch.float32
fc3.weight torch.float32
fc3.bias torch.float32
4.1 剪枝前的参数
model = LeNet5()
module = model.conv1
print(list(module.named_parameters()))
输出：权重有6个矩阵，每个矩阵的size是3*3，偏差为6个value
[('weight', Parameter containing:
tensor([[[[-0.0541, 0.3125, -0.2231],
[-0.2518, 0.0718, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.1638, -0.1794],
[ 0.0089, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0759],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0013, 0.2150],
[-0.1323, 0.2928, 0.2983],
[-0.2250, 0.0447, -0.0897]]],
[[[-0.2509, -0.1535, 0.2116],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.1805, 0.1331],
[-0.0635, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], requires_grad=True)), ('bias', Parameter containing:
tensor([ 0.0172, -0.0222, 0.0460, 0.2165, 0.0469, 0.0145],
requires_grad=True))]
检测是否有缓冲区
print(list(module.named_buffers()))
输出为空矩阵：[]
ps：parameter和buffer的区别
模型中需要保存下来的参数包括两种：
1）一种是反向传播需要被optimizer更新的，称之为 parameter
2）一种是反向传播不需要被optimizer更新，称之为 buffer
4.2 对卷积层权重剪枝
目标：我们将在conv1层中名为weight的参数中随机修剪 30％的连接
- 从torch.nn.utils.prune选择修建技术
- 指定模块和该模块中需要修剪的参数名称
- 使用所选修剪技术所需的适当关键字参数，指定修剪参数。
# 参数说明: 对model.conv1中的权重"weight"剪枝30%的参数
prune.random_unstructured(module, name="weight", amount=0.3)
# Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1))
修剪函数执行时候的内部原理：
修剪是通过从参数中删除weight并将其替换为名为weight_orig的新参数(即，将"_orig"附加到初始参数name）来进行的。 weight_orig存储未修剪的张量版本。 bias未修剪，因此它将保持完整。
print(list(module.named_parameters()))
# 输出：
[('bias', Parameter containing:
tensor([ 0.0172, -0.0222, 0.0460, 0.2165, 0.0469, 0.0145],
requires_grad=True)), ('weight_orig', Parameter containing:
tensor([[[[-0.0541, 0.3125, -0.2231],
[-0.2518, 0.0718, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.1638, -0.1794],
[ 0.0089, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0759],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0013, 0.2150],
[-0.1323, 0.2928, 0.2983],
[-0.2250, 0.0447, -0.0897]]],
[[[-0.2509, -0.1535, 0.2116],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.1805, 0.1331],
[-0.0635, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], requires_grad=True))]
通过以上选择的修剪技术生成的修剪掩码将保存为名为weight_mask的模块缓冲区，存储在named_buffers中，因为掩码不需要反向传播，永久保存。
print(list(module.named_buffers()))
# 输出：
[('weight_mask', tensor([[[[1., 1., 0.],
[0., 0., 1.],
[1., 1., 1.]]],
[[[1., 0., 0.],
[0., 1., 1.],
[1., 1., 1.]]],
[[[1., 1., 0.],
[1., 1., 1.],
[1., 1., 1.]]],
[[[1., 0., 1.],
[0., 1., 1.],
[0., 1., 0.]]],
[[[0., 0., 0.],
[1., 1., 1.],
[1., 1., 1.]]],
[[[1., 0., 1.],
[0., 1., 1.],
[1., 1., 1.]]]]))]
需要注意：mask里标记为0的位置对应的weight是被pruned掉的，在retrained的时候保持为0。而如果这时候打印weight，会得到掩码和原始参数结合的版本（即pruned的权重变为0）。注意这里的weight不是一个参数，只是一个属性。
print(module.weight)
# 输出：
tensor([[[[-0.0541, 0.3125, -0.0000],
[-0.0000, 0.0000, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.0000, -0.0000],
[ 0.0000, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0000],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0000, 0.2150],
[-0.0000, 0.2928, 0.2983],
[-0.0000, 0.0447, -0.0000]]],
[[[-0.0000, -0.0000, 0.0000],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.0000, 0.1331],
[-0.0000, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], grad_fn=<MulBackward0>)
剪枝需要在每次前向传播之前被应用。通过PyTorch 的forward_pre_hooks可以应用剪枝。
当模型被剪枝时，它将为与该模型关联的每个参数获取forward_pre_hook进行修剪。（注意，在这里模型不是指整个网络模型，而是指被剪枝的子模型，比如在这里是指conv1）
在这种情况下，由于到目前为止我们只修剪了名称为weight的原始参数，因此只会出现一个钩子。
print(module._forward_pre_hooks)
# 输出：
OrderedDict([(35, <torch.nn.utils.prune.RandomUnstructured object at 0x7f935dfe1c70>)])
4.3 对卷积层偏置剪枝
现在同样对bias进行剪枝处理，以查看module的参数parameter，缓冲区buffer，挂钩hook和属性property如何变化。
在这里我们尝试另一种修剪方法，按 L1 范数修剪掉最小的3个偏差bias
prune.random_unstructured(module, name="bias", amount=3)
现在查看各参数变化：
print(list(module.named_parameters()))
# 输出：
[('weight_orig', Parameter containing:
tensor([[[[-0.0541, 0.3125, -0.2231],
[-0.2518, 0.0718, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.1638, -0.1794],
[ 0.0089, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0759],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0013, 0.2150],
[-0.1323, 0.2928, 0.2983],
[-0.2250, 0.0447, -0.0897]]],
[[[-0.2509, -0.1535, 0.2116],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.1805, 0.1331],
[-0.0635, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], requires_grad=True)),
('bias_orig', Parameter containing:
tensor([ 0.0172, -0.0222, 0.0460, 0.2165, 0.0469, 0.0145],
requires_grad=True))]
缓冲区：
print(list(module.named_buffers()))
# 输出：
[('weight_mask', tensor([[[[1., 1., 0.],
[0., 0., 1.],
[1., 1., 1.]]],
[[[1., 0., 0.],
[0., 1., 1.],
[1., 1., 1.]]],
[[[1., 1., 0.],
[1., 1., 1.],
[1., 1., 1.]]],
[[[1., 0., 1.],
[0., 1., 1.],
[0., 1., 0.]]],
[[[0., 0., 0.],
[1., 1., 1.],
[1., 1., 1.]]],
[[[1., 0., 1.],
[0., 1., 1.],
[1., 1., 1.]]]])),
('bias_mask', tensor([1., 1., 0., 0., 0., 1.]))]
属性值：
print(module.bias)
# 输出：
tensor([ 0.0172, -0.0222, 0.0000, 0.0000, 0.0000, 0.0145],
grad_fn=<MulBackward0>)
钩子：可以看到有两个钩子
print(module._forward_pre_hooks)
# 输出：
OrderedDict([(35, <torch.nn.utils.prune.RandomUnstructured object at 0x7f935dfe1c70>),
(36, <torch.nn.utils.prune.RandomUnstructured object at 0x7f935dfe1b20>)])
4.4 对卷积层权重删除修剪
现在对利用prune.remove函数从模块中删除修剪重新参数化，从前向钩子中删除修剪方法。来查看个参数的变化，以了解prune.remove函数的作用。
# 参数说明: 删除module中对"weight"的剪枝处理
prune.remove(module, "weight")
# Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))
各参数变化：weight替换掉了原来的weight_orig，表示参数的更新，只不过由于剪枝后所以部分的参数为0
print(list(module.named_parameters()))
# 输出：
[('bias_orig', Parameter containing:
tensor([ 0.0172, -0.0222, 0.0460, 0.2165, 0.0469, 0.0145],
requires_grad=True)), ('weight', Parameter containing:
tensor([[[[-0.0541, 0.3125, -0.0000],
[-0.0000, 0.0000, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.0000, -0.0000],
[ 0.0000, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0000],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0000, 0.2150],
[-0.0000, 0.2928, 0.2983],
[-0.0000, 0.0447, -0.0000]]],
[[[-0.0000, -0.0000, 0.0000],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.0000, 0.1331],
[-0.0000, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], requires_grad=True))]
缓冲区：
print(list(module.named_buffers()))
# 输出：
[('bias_mask', tensor([1., 1., 0., 0., 0., 1.]))]
属性值：由grad_fn=变为requires_grad=True，现在变成了需要更新的参数
print(module.weight)
# 输出：
Parameter containing:
tensor([[[[-0.0541, 0.3125, -0.0000],
[-0.0000, 0.0000, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.0000, -0.0000],
[ 0.0000, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0000],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0000, 0.2150],
[-0.0000, 0.2928, 0.2983],
[-0.0000, 0.0447, -0.0000]]],
[[[-0.0000, -0.0000, 0.0000],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.0000, 0.1331],
[-0.0000, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], requires_grad=True)
钩子：只剩下对bias剪枝的钩子
print(module._forward_pre_hooks)
# 输出：
OrderedDict([(36, <torch.nn.utils.prune.RandomUnstructured object at 0x7f935dfe1b20>)])
4.5 对卷积层偏置删除修剪
prune.remove(module, "bias")
同样，四部曲分别查看参数，缓存区，属性，钩子，变化如下：
# 1. 参数查看
print(list(module.named_parameters()))
# 输出：
[('weight', Parameter containing:
tensor([[[[-0.0541, 0.3125, -0.0000],
[-0.0000, 0.0000, 0.1464],
[-0.2506, 0.0921, 0.1507]]],
[[[-0.0208, -0.0000, -0.0000],
[ 0.0000, -0.2395, -0.0704],
[ 0.0643, -0.1986, -0.1765]]],
[[[-0.1040, -0.1539, -0.0000],
[-0.3282, 0.0630, -0.3091],
[-0.0085, 0.2010, 0.0239]]],
[[[-0.2581, -0.0000, 0.2150],
[-0.0000, 0.2928, 0.2983],
[-0.0000, 0.0447, -0.0000]]],
[[[-0.0000, -0.0000, 0.0000],
[-0.0049, 0.0929, -0.2913],
[ 0.1029, 0.1205, -0.1133]]],
[[[-0.3326, 0.0000, 0.1331],
[-0.0000, 0.0032, -0.2707],
[-0.2712, -0.2373, 0.2067]]]], requires_grad=True)), ('bias', Parameter containing:
tensor([ 0.0172, -0.0222, 0.0000, 0.0000, 0.0000, 0.0145],
requires_grad=True))]
# 2. 缓存区查看
print(list(module.named_buffers()))
# 输出: [] 一个空列表
# 3. 属性查看
print(module.bias)
# 输出:
Parameter containing:
tensor([ 0.0172, -0.0222, 0.0000, 0.0000, 0.0000, 0.0145],
requires_grad=True)
# 4. 钩子查看
print(module._forward_pre_hooks)
# 输出:
OrderedDict()
看到这里应该就能知道prune.remove的作用了，就是将剪枝的模块重新复制更新，取消掩码钩子等操作，变成普通网络的训练参数。
4.6 剪枝对参数量与浮点计算量的变化
这里对yolov5的代码进行参考，测试剪枝前后模型参数的变化与浮点计算量的变化，需要主要的是这里遍历每一个模块之后是对剪枝进行remove更新的，也就是说剪枝后权重与偏置为0的参数是直接被赋值到模型中，这样就可以计算模型剪枝后的稀疏度。
参考代码：
# 功能: 测试模型参数
def model_parms(model):
from thop import profile
input = torch.randn(1, 1, 32, 32)
flops, params = profile(model, inputs=(input,))
print('flops:{}M'.format(flops / 1e6))
print('params:{}Kb'.format(params / 1e3))
# 功能：测试模型稀疏度
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 model_prune(model):
import torch.nn.utils.prune as prune
for name, module in model.named_modules():
# Prune all 2D convolutional layers by 30%
if isinstance(module, torch.nn.Conv2d):
print("torch.nn.Conv2d:", name)
prune.random_unstructured(module,name='weight', amount=0.3)
prune.random_unstructured(module,name='bias', amount=3)
prune.remove(module, 'weight')
prune.remove(module, 'bias')
# Prune all linear layers by 50%.
elif isinstance(module, torch.nn.Linear):
print("torch.nn.Linear:", name)
prune.random_unstructured(module,name='weight', amount=0.5)
prune.random_unstructured(module,name='bias', amount=3)
prune.remove(module, 'weight')
prune.remove(module, 'bias')
# 测试函数
if __name__ == '__main__':
model = LeNet5()
# 剪枝前参数测试
model_parms(model)
print(sparsity(model))
model_prune(model)
# 剪枝后参数测试
model_parms(model)
print(sparsity(model))
输出：
# 剪枝前
[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.
[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.
[WARN] Cannot find rule for <class '__main__.LeNet5'>. Treat it as zero Macs and zero Params.
flops:0.28276M
params:81.194Kb
tensor(0.)
# 剪枝后
torch.nn.Conv2d: conv1
torch.nn.Conv2d: conv2
torch.nn.Linear: fc1
torch.nn.Linear: fc2
torch.nn.Linear: fc3
[INFO] Register count_convNd() for <class 'torch.nn.modules.conv.Conv2d'>.
[INFO] Register count_linear() for <class 'torch.nn.modules.linear.Linear'>.
[WARN] Cannot find rule for <class '__main__.LeNet5'>. Treat it as zero Macs and zero Params.
flops:0.28276M
params:81.194Kb
tensor(0.4965)
可以看见，这里由于是对模型进行掩码的操作，所以模型的参数量是没有改变的，因为正如上诉所说权重并不是真的被剪掉了，只是掩码的作用让他失活不对模型更新产生作用。从稀疏度的变化就可以了解到了，当前的模型有很多神经元被失活置为0。
参考资料：
1. 模型剪枝简介
2. 模型剪枝
3. Pytorch Utils 总结
4. Pytorch袖珍手册之十四
5. Pytorch剪枝代码示例和注释
6. 学习笔记——神经网络压缩