剪枝是一种常用的模型压缩策略。通过将模型中不重要的连接失效，实现模型瘦身的效果，并减少计算量。PyTorch中实现的剪枝方式有三种：
- 局部剪枝
- 全局剪枝
- 自定义剪枝
Pytorch中与剪枝有关的接口封装在torch.nn.utils.prune中。下面开始演示三种剪枝在LeNet网络中的应用效果，首先给出LeNet网络结构。
import torch from torch import nn class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() # 1: 图像的输入通道(1是黑白图像), 6: 输出通道, 3x3: 卷积核的尺寸 self.conv1 = nn.Conv2d(1, 6, 3) self.conv2 = nn.Conv2d(6, 16, 3) self.fc1 = nn.Linear(16 * 5 * 5, 120) # 5x5 是经历卷积操作后的图片尺寸 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
局部剪枝
局部剪枝实验，假定对模型的第一个卷积层中的权重进行剪枝
model_1 = LeNet() module = model_1.conv1 # 剪枝前 print(list(module.named_parameters())) print(list(module.named_buffers())) prune.random_unstructured(module, name="weight", amount=0.3) # 剪枝后 print(list(module.named_parameters())) print(list(module.named_buffers()))
运行结果
## 剪枝前 [('weight', Parameter containing: tensor([[[[ 0.1729, -0.0109, -0.1399], [ 0.1019, 0.1883, 0.0054], [-0.0790, -0.1790, -0.0792]]], ... [[[ 0.2465, 0.2114, 0.3208], [-0.2067, -0.2097, -0.0431], [ 0.3005, -0.2022, 0.1341]]]], requires_grad=True)), ('bias', Parameter containing: tensor([-0.1437, 0.0605, 0.1427, -0.3111, -0.2476, 0.1901], requires_grad=True))] [] ## 剪枝后 [('bias', Parameter containing: tensor([-0.1437, 0.0605, 0.1427, -0.3111, -0.2476, 0.1901], requires_grad=True)), ('weight_orig', Parameter containing: tensor([[[[ 0.1729, -0.0109, -0.1399], [ 0.1019, 0.1883, 0.0054], [-0.0790, -0.1790, -0.0792]]], ... [[[ 0.2465, 0.2114, 0.3208], [-0.2067, -0.2097, -0.0431], [ 0.3005, -0.2022, 0.1341]]]], requires_grad=True))] [('weight_mask', tensor([[[[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]]], [[[0., 1., 0.], [0., 1., 1.], [1., 0., 1.]]], [[[0., 1., 1.], [1., 0., 1.], [1., 0., 1.]]], [[[1., 1., 1.], [1., 0., 1.], [0., 1., 0.]]], [[[0., 0., 1.], [0., 1., 1.], [1., 1., 1.]]], [[[0., 1., 1.], [0., 1., 0.], [1., 1., 1.]]]]))]
模型经历剪枝操作后， 原始的权重矩阵weight参数不见了，变成了weight_orig。 并且剪枝前打印为空列表的module.named_buffers()，此时拥有了一个weight_mask参数。经过剪枝操作后的模型，原始的参数存放在了weight_orig中，对应的剪枝矩阵存放在weight_mask中, 而将weight_mask视作掩码张量，再和weight_orig相乘的结果就存放在了weight中。
全局剪枝
局部剪枝只能以部分网络模块为单位进行剪枝，更广泛的剪枝策略是采用全局剪枝(global pruning)，比如在整体网络的视角下剪枝掉20%的权重参数，而不是在每一层上都剪枝掉20%的权重参数。采用全局剪枝后，不同的层被剪掉的百分比不同。
model_2 = LeNet().to(device=device) # 首先打印初始化模型的状态字典 print(model_2.state_dict().keys()) # 构建参数集合, 决定哪些层, 哪些参数集合参与剪枝 parameters_to_prune = ( (model_2.conv1, 'weight'), (model_2.conv2, 'weight'), (model_2.fc1, 'weight'), (model_2.fc2, 'weight'), (model_2.fc3, 'weight')) # 调用prune中的全局剪枝函数global_unstructured执行剪枝操作, 此处针对整体模型中的20%参数量进行剪枝 prune.global_unstructured(parameters_to_prune, pruning_method=prune.L1Unstructured, amount=0.2) # 最后打印剪枝后的模型的状态字典 print(model_2.state_dict().keys())
输出结果
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.bias', 'fc3.weight_orig', 'fc3.weight_mask'])
当采用全局剪枝策略的时候(假定20%比例参数参与剪枝)，仅保证模型总体参数量的20%被剪枝掉，具体到每一层的情况则由模型的具体参数分布情况来定。
自定义剪枝
自定义剪枝可以自定义一个子类，用来实现具体的剪枝逻辑，比如对权重矩阵进行间隔性的剪枝
class my_pruning_method(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_pruning(module, name): my_pruning_method.apply(module, name) return module model_3 = LeNet() print(model_3)
在剪枝前查看网络结构
LeNet( (conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1)) (conv2): Conv2d(6, 16, kernel_size=(3, 3), stride=(1, 1)) (fc1): Linear(in_features=400, out_features=120, bias=True) (fc2): Linear(in_features=120, out_features=84, bias=True) (fc3): Linear(in_features=84, out_features=10, bias=True) )
采用自定义剪枝的方式对局部模块fc3进行剪枝
my_unstructured_pruning(model.fc3, name="bias") print(model.fc3.bias_mask)
输出结果
tensor([0., 1., 0., 1., 0., 1., 0., 1., 0., 1.])
最后的剪枝效果与实现的逻辑一致。
参考文档
