当前位置：首页 > news >正文

datawhale11月组队学习模型压缩技术2：PyTorch模型剪枝教程

news 2025/11/4 8:09:42

文章目录

- 一、 prune模块简介
- - 1.1 常用方法
  - 1.2 剪枝效果
  - 1.3 二、三、四章剪枝测试总结
- 二、局部剪枝（Local Pruning）
- - 2.1 结构化剪枝
  - - 2.1.1 对weight进行随机结构化剪枝（random_structured）
    - 2.1.2 对weight进行迭代剪枝（范数结构化剪枝，ln_structured）
  - 2.2 非结构化剪枝
  - - 2.2.1 对bias进行随机非结构化剪枝
    - 2.2.2 对多层网络进行范数非结构化剪枝（l1_unstructured）
  - 2.3 永久化剪枝（remove）
- 三、全局剪枝(GLobal pruning)
- 四、自定义剪枝(Custom pruning)

《datawhale2411组队学习之模型压缩技术1：模型剪枝（上）》：介绍模型压缩的几种技术；模型剪枝基本概念、分类方式、剪枝标准、剪枝频次、剪枝后微调等内容
《datawhale11月组队学习模型压缩技术2：PyTorch模型剪枝教程》：介绍PyTorch的prune模块具体用法
《datawhale11月组队学习模型压缩技术3：2：4结构稀疏化BERT模型》：介绍基于模式的剪枝——2：4结构稀疏化及其在BERT模型上的测试效果

项目地址awesome-compression、在线阅读

一、 prune模块简介

PyTorch教程《Pruning Tutorial》、torch.nn.utils.prune文档

1.1 常用方法

Pytorch在1.4.0版本开始，加入了剪枝操作，在torch.nn.utils.prune模块中，主要有以下剪枝方法：

剪枝类型	子类型	剪枝方法
局部剪枝	结构化剪枝	随机结构化剪枝 (`random_structured`)
		范数结构化剪枝 (`ln_structured`)
	非结构化剪枝	随机非结构化剪枝 (`random_unstructured`)
		范数非结构化剪枝 (`ln_unstructured`)
全局剪枝	非结构化剪枝	全局非结构化剪枝 (`global_unstructured`)
自定义剪枝		自定义剪枝 (`Custom Pruning`)

除此之外，模块中还有一些其它方法：

方法	描述
`prune.remove(module, name)`	剪枝永久化
`prune.apply`	使用指定的剪枝方法对模块进行剪枝。
`prune.is_pruned(module)`	检查给定模块的某个参数是否已被剪枝。
`prune.custom_from_mask(module, name, mask)`	基于自定义的掩码进行剪枝，用于定义更加细粒度的剪枝策略。

1.2 剪枝效果

参数变化：
- 剪枝前，weight 是模型的一个参数，意味着它是模型训练时优化的对象，可以通过梯度更新（通过 optimizer.step() 来更新它的值）。
- 剪枝过程中，原始权重被保存到新的变量 weight_orig中，便于后续访问原始权重。
- 剪枝后，weight是剪枝后的权重值（通过原始权重和剪枝掩码计算得出），但此时不再是参数，而是模型的属性（一个普通的变量）。
掩码存储：生成一个名为 weight_mask的剪枝掩码，会被保存为模块的一个缓冲区（buffer）。
前向传递：PyTorch 使用 forward_pre_hooks 来确保每次前向传递时都会应用剪枝处理。每个被剪枝的参数都会在模块中添加一个钩子来实现这一操作。

1.3 二、三、四章剪枝测试总结

对weight进行剪枝，效果见1.2 章节。
对weight进行迭代剪枝，相当于把多个剪枝核（mask）序列化成一个剪枝核，最终只有一个weight_orig和weight_mask，hook也被更新。
对weight剪枝后，再对bias进行剪枝，weight_orig和weight_mask不变，新增bias_orig和bias_mask，新增bias hook。
可以对多个模块同时进行剪枝，最后使用remove进行剪枝永久化
使用remove函数后， weight_orig 和 bias_orig 被移除，剪枝后的weight 和 bias 成为标准的模型参数。经过 remove 操作后，剪枝永久化生效。此时，剪枝掩码weight_mask 和 hook不再需要，named_buffers和_forward_pre_hooks 都被清空。
局部剪枝需要根据自己的经验来决定对某一层网络进行剪枝，需要对模型有深入了解，所以全局剪枝（跨不同参数）更通用，即从整体网络的角度进行剪枝。采用全局剪枝时，不同的层被剪掉的百分比可能不同。

parameters_to_prune = ((model.conv1, 'weight'),(model.conv2, 'weight'),(model.fc1, 'weight'),(model.fc2, 'weight'))# 应用20%全局剪枝
prune.global_unstructured(parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2)

最终各层剪枝比例为（随机的）：

Sparsity in conv1.weight: 5.33%
Sparsity in conv2.weight: 17.25%
Sparsity in fc1.weight: 22.03%
Sparsity in fc2.weight: 14.67%
Global sparsity: 20.00%

自定义剪枝需要通过继承class BasePruningMethod()来定义,，其内部有若干方法: call, apply_mask, apply, prune, remove。其中，必须实现__init__和compute_mask两个函数才能完成自定义的剪枝规则设定。此外，您必须指定要实现的修剪类型（ global, structured, and unstructured）。

二、局部剪枝（Local Pruning）

局部剪枝，指的是对网络的单个层或局部范围内进行剪枝。其中，非结构化剪枝会随机地将一些权重参数变为0，结构化剪枝则将某个维度某些通道的权重变成0。
总结一下2.1和2.2的效果：

import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F
from torchsummary import summary# 1.定义一个经典的LeNet网络
class LeNet(nn.Module):def __init__(self, num_classes=10):super(LeNet, self).__init__()self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)self.fc1 = nn.Linear(in_features=16 * 4 * 4, out_features=120)self.fc2 = nn.Linear(in_features=120, out_features=84)self.fc3 = nn.Linear(in_features=84, out_features=num_classes)def forward(self, x):x = self.maxpool(F.relu(self.conv1(x)))x = self.maxpool(F.relu(self.conv2(x)))x = x.view(x.size()[0], -1)x = F.relu(self.fc1(x))x = F.relu(self.fc2(x))x = self.fc3(x)return x
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet().to(device=device)# 2.打印模型结构
summary(model, input_size=(1, 28, 28))

----------------------------------------------------------------Layer (type)               Output Shape         Param #
================================================================Conv2d-1            [-1, 6, 24, 24]             156MaxPool2d-2            [-1, 6, 12, 12]               0Conv2d-3             [-1, 16, 8, 8]           2,416MaxPool2d-4             [-1, 16, 4, 4]               0Linear-5                  [-1, 120]          30,840Linear-6                   [-1, 84]          10,164Linear-7                   [-1, 10]             850
================================================================
Total params: 44,426
Trainable params: 44,426
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.00
Forward/backward pass size (MB): 0.04
Params size (MB): 0.17
Estimated Total Size (MB): 0.22
----------------------------------------------------------------

# 3.打印模型的状态字典，状态字典里包含了所有的参数
print(model.state_dict().keys())

odict_keys(['conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])

# 4.打印第一个卷积层的参数
module = model.conv1
print(list(module.named_parameters()))

[('weight', Parameter containing:
tensor([[[[ 0.1529,  0.1660, -0.0469,  0.1837, -0.0438],[ 0.0404, -0.0974,  0.1175,  0.1763, -0.1467],[ 0.1738,  0.0374,  0.1478,  0.0271,  0.0964],[-0.0282,  0.1542,  0.0296, -0.0934,  0.0510],[-0.0921, -0.0235, -0.0812,  0.1327, -0.1579]]],......[[[-0.1167, -0.0685, -0.1579,  0.1677, -0.0397],[ 0.1721,  0.0623, -0.1694,  0.1384, -0.0550],[-0.0767, -0.1660, -0.1988,  0.0572, -0.0437],[ 0.0779, -0.1641,  0.1485, -0.1468, -0.0345],[ 0.0418,  0.1033,  0.1615,  0.1822, -0.1586]]]], device='cuda:0',requires_grad=True)), ('bias', Parameter containing:
tensor([ 0.0503, -0.0860, -0.0219, -0.1497,  0.1822, -0.1468], device='cuda:0',requires_grad=True))]

# 5.打印module中属性张量named_buffers，此时为空列表
print(list(module.named_buffers()))

[]

2.1 结构化剪枝

2.1.1 对weight进行随机结构化剪枝（random_structured）

对LeNet的conv1层的weight参数进行随机结构化剪枝，其中 amount是一个介于0.0-1.0的float数值,代表比例, 或者一个正整数，代表剪裁掉多少个参数.

prune.random_structured(module, name="weight", amount=2, dim=0)

# 1.再次打印模型的状态字典，发现conv1层多了weight_orig和weight_mask
print(model.state_dict().keys())

odict_keys(['conv1.bias', 'conv1.weight_orig', 'conv1.weight_mask', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])

Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))

# 2. 剪枝后，原始的weight变成了weight_orig，并存放在named_parameters中
print(list(module.named_parameters()))

[('bias', Parameter containing:
tensor([ 0.0503, -0.0860, -0.0219, -0.1497,  0.1822, -0.1468], device='cuda:0',requires_grad=True)), ('weight_orig', Parameter containing:
tensor([[[[ 0.1529,  0.1660, -0.0469,  0.1837, -0.0438],[ 0.0404, -0.0974,  0.1175,  0.1763, -0.1467],[ 0.1738,  0.0374,  0.1478,  0.0271,  0.0964],[-0.0282,  0.1542,  0.0296, -0.0934,  0.0510],[-0.0921, -0.0235, -0.0812,  0.1327, -0.1579]]],......[[[-0.1167, -0.0685, -0.1579,  0.1677, -0.0397],[ 0.1721,  0.0623, -0.1694,  0.1384, -0.0550],[-0.0767, -0.1660, -0.1988,  0.0572, -0.0437],[ 0.0779, -0.1641,  0.1485, -0.1468, -0.0345],[ 0.0418,  0.1033,  0.1615,  0.1822, -0.1586]]]], device='cuda:0',requires_grad=True))]

# 3. 剪枝掩码矩阵weight_mask存放在模块的buffer中
print(list(module.named_buffers()))

[('weight_mask', tensor([[[[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.]]],[[[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.]]],[[[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.]]],[[[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.]]],[[[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.]]],[[[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.]]]]))]

# 4. 剪枝操作后的weight已经不再是module的参数, 而只是module的一个属性.
print(module.weight)

tensor([[[[ 0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[ 0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[ 0.0000,  0.0000, -0.0000, -0.0000, -0.0000],[-0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[ 0.0000,  0.0000, -0.0000, -0.0000, -0.0000]]],[[[-0.0540, -0.1928, -0.0355, -0.0075, -0.1481],[ 0.0135,  0.0192,  0.0082, -0.0120, -0.0164],[-0.0435, -0.1488,  0.1092, -0.0041,  0.1960],[-0.1045, -0.0136,  0.0398, -0.1286,  0.0617],[-0.0091,  0.0466,  0.1827,  0.1655,  0.0727]]],[[[ 0.1216, -0.0833, -0.1491, -0.1143,  0.0113],[ 0.0452,  0.1662, -0.0425, -0.0904, -0.1235],[ 0.0565,  0.0933, -0.0721,  0.0909,  0.1837],[-0.1739,  0.0263,  0.1339,  0.0648, -0.0382],[-0.1667,  0.1478,  0.0448, -0.0892,  0.0815]]],[[[ 0.0000,  0.0000,  0.0000, -0.0000,  0.0000],[-0.0000,  0.0000,  0.0000,  0.0000, -0.0000],[-0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[-0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[ 0.0000, -0.0000,  0.0000, -0.0000, -0.0000]]],[[[ 0.1278,  0.1037, -0.0323, -0.1504,  0.1080],[ 0.0266, -0.0996,  0.1499, -0.0845,  0.0609],[-0.0662, -0.1405, -0.0586, -0.0615, -0.0462],[-0.1118, -0.0961, -0.1325, -0.0417, -0.0741],[ 0.1842, -0.1040, -0.1786, -0.0593,  0.0186]]],[[[-0.0889, -0.0737, -0.1655, -0.1708, -0.0988],[-0.1787,  0.1127,  0.0706, -0.0352,  0.1238],[-0.0985, -0.1929, -0.0062,  0.0488, -0.1152],[-0.1659, -0.0448,  0.0821, -0.0956, -0.0262],[ 0.1928,  0.1767, -0.1792, -0.1364,  0.0507]]]],grad_fn=<MulBackward0>)

对于每一次剪枝操作，PyTorch 会为剪枝的参数（如 weight）添加一个 forward_pre_hook。这个钩子会在每次进行前向传递计算之前，自动应用剪枝掩码（即将某些权重置为零），这保证了剪枝后的权重在模型计算时被正确地使用。

# 5.打印_forward_pre_hooks
print(module._forward_pre_hooks)

OrderedDict([(0, <torch.nn.utils.prune.RandomStructured object at 0x7f04012f8ca0>)])

简单总结就是：

weight 不再是参数，它变成了一个属性，表示剪枝后的权重。
weight_orig 保存原始未剪枝的权重。
weight_mask 是一个掩码，表示哪些权重被剪去了（即哪些位置变为零）。
钩子会保证每次前向传递时，weight 会根据 weight_mask 来计算出剪枝后的版本。

2.1.2 对weight进行迭代剪枝（范数结构化剪枝，ln_structured）

一个模型的参数可以执行多次剪枝操作，这种操作被称为迭代剪枝（Iterative Pruning）。上述步骤已经对conv1进行了随机结构化剪枝，接下来对其再进行范数结构化剪枝，看看会发生什么？

# n代表范数，这里n=2表示l2范数
prune.ln_structured(module, name="weight", amount=0.5, n=2, dim=0)# 再次打印模型参数
print(" model state_dict keys:")
print(model.state_dict().keys())
print('*'*50)print(" module named_parameters:")
print(list(module.named_parameters()))
print('*'*50)print(" module named_buffers:")
print(list(module.named_buffers()))
print('*'*50)print(" module weight:")
print(module.weight)
print('*'*50)print(" module _forward_pre_hooks:")
print(module._forward_pre_hooks)

model state_dict keys:
odict_keys(['conv1.bias', 'conv1.weight_orig', 'conv1.weight_mask', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])
**************************************************
module named_parameters:	# 原始参数weight_orig不变
...
...
module named_buffers:
[('weight_mask', tensor([[[[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.]]],[[[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.]]],[[[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.]]],[[[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.]]],[[[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.],[0., 0., 0., 0., 0.]]],[[[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.],[1., 1., 1., 1., 1.]]]]))]
**************************************************module weight:......
module _forward_pre_hooks:
OrderedDict([(1, <torch.nn.utils.prune.PruningContainer object at 0x7f04c86756d0>)])

可见迭代剪枝相当于把多个剪枝核序列化成一个剪枝核，新的 mask 矩阵与旧的 mask 矩阵的结合由PruningContainer的compute_mask方法处理，最后只有一个weight_orig和weight_mask。

module._forward_pre_hooks是一个用于在模型的前向传播之前执行自定义操作的机制，这里记录了执行过的剪枝方法：

# 打印剪枝历史
for hook in module._forward_pre_hooks.values():if hook._tensor_name == "weight":  breakprint(list(hook))

[<torch.nn.utils.prune.RandomStructured object at 0x7f04012f8ca0>, <torch.nn.utils.prune.LnStructured object at 0x7f04c8675b80>]

2.2 非结构化剪枝

2.2.1 对bias进行随机非结构化剪枝

此时，我们也可以继续对偏置bias进行剪枝，看看module的参数、缓冲区、钩子和属性是如何变化的。

prune.random_unstructured(module, name="bias", amount=1)
# 再次打印模型参数
print(" model state_dict keys:")
print(model.state_dict().keys())
print('*'*50)print(" module named_parameters:")
print(list(module.named_parameters()))
print('*'*50)print(" module named_buffers:")
print(list(module.named_buffers()))
print('*'*50)print(" module bias:")
print(module.bias)
print('*'*50)print(" module _forward_pre_hooks:")
print(module._forward_pre_hooks)

model state_dict keys:
odict_keys(['conv1.weight_orig', 'conv1.bias_orig', 'conv1.weight_mask', 'conv1.bias_mask', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])
**************************************************
# weight_orig不变，添加了bias_origmodule named_parameters:  
[('weight_orig', Parameter containing:...
, requires_grad=True)), ('bias_orig', Parameter containing:
tensor([-0.0893, -0.1464, -0.1101, -0.0076,  0.1493, -0.0418],requires_grad=True))]
**************************************************
# weight_mask不变，添加了bias_maskmodule named_buffers:
[('weight_mask', 
...('bias_mask', tensor([1., 1., 0., 1., 1., 1.]))]
**************************************************module bias:
tensor([-0.0893, -0.1464, -0.0000, -0.0076,  0.1493, -0.0418],grad_fn=<MulBackward0>)
**************************************************module _forward_pre_hooks:
OrderedDict([(1, <torch.nn.utils.prune.PruningContainer object at 0x7f04c86756d0>), (2, <torch.nn.utils.prune.RandomUnstructured object at 0x7f04013a7d30>)])

对bias进行剪枝后，会发现state_dict和named_parameters中不仅仅有了weight_orig，也有了bias_orig。在named_buffers中, 也同时出现了weight_mask和bias_mask。最后，因为我们在两种参数上进行剪枝，因此会生成两个钩子。

2.2.2 对多层网络进行范数非结构化剪枝（l1_unstructured）

前面介绍了对指定的conv1层的weight和bias进行了不同方法的剪枝，那么能不能支持同时对多层网络的特定参数进行剪枝呢？

# 对于模型多个模块进行bias剪枝
for n, m in model.named_modules():# 对模型中所有的卷积层执行l1_unstructured剪枝操作, 选取20%的参数剪枝if isinstance(m, torch.nn.Conv2d):prune.l1_unstructured(m, name="bias", amount=0.2)# 对模型中所有全连接层执行ln_structured剪枝操作, 选取40%的参数剪枝# elif isinstance(module, torch.nn.Linear):#     prune.random_structured(module, name="weight", amount=0.4,dim=0)# 再次打印模型参数
print(" model state_dict keys:")
print(model.state_dict().keys())
print('*'*50)print(" module named_parameters:")
print(list(module.named_parameters()))
print('*'*50)print(" module named_buffers:")
print(list(module.named_buffers()))
print('*'*50)print(" module weight:")
print(module.weight)
print('*'*50)print(" module bias:")
print(module.bias)
print('*'*50)print(" module _forward_pre_hooks:")
print(module._forward_pre_hooks)

model state_dict keys:
odict_keys(['conv1.weight_orig', 'conv1.bias_orig', 'conv1.weight_mask', 'conv1.bias_mask', 'conv2.weight', 'conv2.bias_orig', 'conv2.bias_mask', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])
**************************************************module named_parameters:[('weight_orig', Parameter containing:...('bias_orig', Parameter containing:...
**************************************************
# # weight_mask不变，bias_mask更新
module named_buffers:
[('weight_mask', ...
('bias_mask', tensor([1., 1., 0., 0., 1., 1.]))]
**************************************************
# module weight不变
module weight:...
**************************************************
module bias:
tensor([-0.0893, -0.1464, -0.0000, -0.0000,  0.1493, -0.0418],grad_fn=<MulBackward0>)
**************************************************
module _forward_pre_hooks:
OrderedDict([(1, <torch.nn.utils.prune.PruningContainer object at 0x7f04c86756d0>), (3, <torch.nn.utils.prune.PruningContainer object at 0x7f04010c1100>)])

2.3 永久化剪枝（remove）

接下来对模型的weight和bias参数进行永久化剪枝操作prune.remove。

# 对module的weight执行剪枝永久化操作remove
for n, m in model.named_modules():if isinstance(m, torch.nn.Conv2d):prune.remove(m, 'bias')# 对conv1的weight执行剪枝永久化操作remove
prune.remove(module, 'weight')
print('*'*50)# 将剪枝后的模型的状态字典打印出来
print(" model state_dict keys:")
print(model.state_dict().keys())
print('*'*50)# 再次打印模型参数
print(" model named_parameters:")
print(list(module.named_parameters()))
print('*'*50)# 再次打印模型mask buffers参数
print(" model named_buffers:")
print(list(module.named_buffers()))
print('*'*50)# 再次打印模型的_forward_pre_hooks
print(" model forward_pre_hooks:")
print(module._forward_pre_hooks)

**************************************************model state_dict keys:
odict_keys(['conv1.bias', 'conv1.weight', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])
**************************************************model named_parameters:
[('bias', Parameter containing:
tensor([-0.0893, -0.1464, -0.0000, -0.0000,  0.1493, -0.0418],requires_grad=True)), ('weight', Parameter containing:
tensor([[[[ 0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[ 0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[ 0.0000,  0.0000, -0.0000, -0.0000, -0.0000],[-0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[ 0.0000,  0.0000, -0.0000, -0.0000, -0.0000]]],[[[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000],[ 0.0000,  0.0000,  0.0000, -0.0000, -0.0000],[-0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[-0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[-0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],[[[ 0.1216, -0.0833, -0.1491, -0.1143,  0.0113],[ 0.0452,  0.1662, -0.0425, -0.0904, -0.1235],[ 0.0565,  0.0933, -0.0721,  0.0909,  0.1837],[-0.1739,  0.0263,  0.1339,  0.0648, -0.0382],[-0.1667,  0.1478,  0.0448, -0.0892,  0.0815]]],[[[ 0.0000,  0.0000,  0.0000, -0.0000,  0.0000],[-0.0000,  0.0000,  0.0000,  0.0000, -0.0000],[-0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[-0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[ 0.0000, -0.0000,  0.0000, -0.0000, -0.0000]]],[[[ 0.0000,  0.0000, -0.0000, -0.0000,  0.0000],[ 0.0000, -0.0000,  0.0000, -0.0000,  0.0000],[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000],[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000],[ 0.0000, -0.0000, -0.0000, -0.0000,  0.0000]]],[[[-0.0889, -0.0737, -0.1655, -0.1708, -0.0988],[-0.1787,  0.1127,  0.0706, -0.0352,  0.1238],[-0.0985, -0.1929, -0.0062,  0.0488, -0.1152],[-0.1659, -0.0448,  0.0821, -0.0956, -0.0262],[ 0.1928,  0.1767, -0.1792, -0.1364,  0.0507]]]], requires_grad=True))]
**************************************************model named_buffers:
[]
**************************************************model forward_pre_hooks:
OrderedDict()

可见，执行remove操作后：

weight_orig 和 bias_orig 被移除，剪枝后的weight 和 bias 成为标准的模型参数。经过 remove 操作后，剪枝永久化生效。
剪枝掩码weight_mask 和 bias_mask不再需要，named_buffers被清空
_forward_pre_hooks 也被清空（由于剪枝后的权重和偏置将直接反映在最终模型中，所以无须再借助外部的掩码或钩子函数来维护剪枝过程）。

三、全局剪枝(GLobal pruning)

前面已经介绍了局部剪枝的四种方法，但这很大程度上需要根据自己的经验来决定对某一层网络进行剪枝。更通用的剪枝策略是采用全局剪枝，即从整体网络的角度进行剪枝。采用全局剪枝时，不同的层被剪掉的百分比可能不同。

model = LeNet().to(device=device)# 1.打印初始化模型的状态字典
print(model.state_dict().keys())
print('*'*50)# 2.构建参数集合, 决定哪些层, 哪些参数集合参与剪枝
parameters_to_prune = ((model.conv1, 'weight'),(model.conv2, 'weight'),(model.fc1, 'weight'),(model.fc2, 'weight'))# 3. 全局剪枝
prune.global_unstructured(parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2)# 4. 打印剪枝后模型的状态字典
print(model.state_dict().keys())

odict_keys(['conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])
**************************************************
odict_keys(['conv1.bias', 'conv1.weight_orig', 'conv1.weight_mask', 'conv2.bias', 'conv2.weight_orig', 'conv2.weight_mask', 'fc1.bias', 'fc1.weight_orig', 'fc1.weight_mask', 'fc2.bias', 'fc2.weight_orig', 'fc2.weight_mask', 'fc3.weight', 'fc3.bias'])

打印一下各层被剪枝的比例：

print("Sparsity in conv1.weight: {:.2f}%".format(100. * float(torch.sum(model.conv1.weight == 0))/ float(model.conv1.weight.nelement())))print("Sparsity in conv2.weight: {:.2f}%".format(100. * float(torch.sum(model.conv2.weight == 0))/ float(model.conv2.weight.nelement())))print("Sparsity in fc1.weight: {:.2f}%".format(100. * float(torch.sum(model.fc1.weight == 0))/ float(model.fc1.weight.nelement())))print("Sparsity in fc2.weight: {:.2f}%".format(100. * float(torch.sum(model.fc2.weight == 0))/ float(model.fc2.weight.nelement())))print("Global sparsity: {:.2f}%".format(100. * float(torch.sum(model.conv1.weight == 0)+ torch.sum(model.conv2.weight == 0)+ torch.sum(model.fc1.weight == 0)+ torch.sum(model.fc2.weight == 0))/ float(model.conv1.weight.nelement()+ model.conv2.weight.nelement()+ model.fc1.weight.nelement()+ model.fc2.weight.nelement())))

Sparsity in conv1.weight: 5.33%
Sparsity in conv2.weight: 17.25%
Sparsity in fc1.weight: 22.03%
Sparsity in fc2.weight: 14.67%
Global sparsity: 20.00%

四、自定义剪枝(Custom pruning)

剪枝模型通过继承class BasePruningMethod()来执行剪枝, 内部有若干方法: call, apply_mask, apply, prune, remove等等。其中，必须实现__init__构造函数和compute_mask两个函数才能完成自定义的剪枝规则设定。此外，您必须指定要实现的修剪类型（ global, structured, and unstructured）。

# 自定义剪枝方法的类, 一定要继承prune.BasePruningMethod
class custom_prune(prune.BasePruningMethod):# 指定此技术实现的修剪类型（支持的选项为global、 structured和unstructured）PRUNING_TYPE = "unstructured"# 内部实现compute_mask函数, 定义剪枝规则, 本质上就是如何去mask掉权重参数def compute_mask(self, t, default_mask):mask = default_mask.clone()# 此处定义的规则是每隔一个参数就遮掩掉一个, 最终参与剪枝的参数的50%被mask掉mask.view(-1)[::2] = 0return mask# 自定义剪枝方法的函数, 内部直接调用剪枝类的方法apply
def custome_unstructured_pruning(module, name):custom_prune.apply(module, name)return module

import time
# 实例化模型类
model = LeNet().to(device=device)start = time.time()
# 调用自定义剪枝方法的函数, 对model中的第1个全连接层fc1中的偏置bias执行自定义剪枝
custome_unstructured_pruning(model.fc1, name="bias")# 剪枝成功的最大标志, 就是拥有了bias_mask参数
print(model.fc1.bias_mask)# 打印一下自定义剪枝的耗时
duration = time.time() - start
print(duration * 1000, 'ms')

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.])
5.576610565185547 ms