注意
跳转到页面底部 下载完整示例代码。
剪枝教程#
创建日期: 2019年7月22日 | 最后更新: 2023年11月2日 | 最后验证: 2024年11月5日
作者: Michela Paganini
最先进的深度学习技术依赖于过度参数化的模型，这些模型难以部署。相反，生物神经网络以高效的稀疏连接而闻名。寻找通过减少参数数量来压缩模型的最佳技术非常重要，这样可以在不牺牲精度的前提下降低内存、电池和硬件消耗。这反过来又允许你在设备上部署轻量级模型，并通过私有的端侧计算保证隐私。在研究领域，剪枝被用于研究过度参数化网络和欠参数化网络之间学习动态的差异，研究幸运的稀疏子网络和初始化的作用（“彩票假设”），以及作为一种破坏性的神经架构搜索技术等。
在本教程中，你将学习如何使用 torch.nn.utils.prune 来稀疏化你的神经网络，以及如何扩展它来实现你自己的自定义剪枝技术。
要求#
"torch>=1.4.0a0+8e8a5e0"
import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F
创建模型#
在本教程中，我们使用 LeCun 等人 1998 年提出的 LeNet 架构。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
# 1 input image channel, 6 output channels, 5x5 square conv kernel
self.conv1 = nn.Conv2d(1, 6, 5)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 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
model = LeNet().to(device=device)
检查模块#
让我们检查 LeNet 模型中（未剪枝的）conv1 层。目前它包含两个参数 weight 和 bias，且没有缓存（buffers）。
module = model.conv1
print(list(module.named_parameters()))
[('weight', Parameter containing:
tensor([[[[ 1.4951e-01, -1.2675e-01, 6.8154e-02, 8.1491e-02, -5.0470e-02],
[ 9.2922e-02, -1.8541e-01, -1.0656e-01, 1.4048e-01, -2.9564e-02],
[ 1.1579e-01, 5.8812e-02, 3.2663e-02, 1.1807e-01, 5.2710e-02],
[ 1.1038e-02, 3.8310e-02, -1.9490e-02, 6.7628e-02, -6.8343e-02],
[-1.2097e-01, -1.2120e-01, -8.9609e-02, 8.5565e-02, -4.3848e-02]]],
[[[-9.8414e-02, -2.3538e-02, -1.6892e-01, -1.1870e-01, 1.5420e-01],
[ 1.2853e-01, -9.0719e-02, 4.2823e-02, 9.4836e-02, 1.9044e-01],
[-1.2375e-01, -4.1072e-02, 7.7668e-02, -4.0979e-02, 1.5920e-02],
[ 8.3536e-02, -1.5734e-01, 1.8694e-01, 1.9463e-02, -8.3581e-02],
[-9.3909e-02, -6.8686e-03, 1.7742e-01, -3.5721e-02, -5.2636e-02]]],
[[[ 8.2311e-02, 8.7024e-02, 1.8021e-01, -1.7326e-01, 6.9537e-02],
[ 9.2128e-02, 7.4754e-03, 6.0818e-02, 1.6462e-01, -5.8529e-03],
[-1.3501e-01, 5.5084e-03, 1.6291e-01, 1.1028e-02, -1.8746e-01],
[-8.8652e-02, 1.9266e-01, 7.8326e-02, -2.9557e-02, -1.5176e-01],
[ 4.4897e-02, -6.9426e-02, 1.9812e-02, -2.5902e-02, -1.5892e-01]]],
[[[-1.7176e-02, -2.9392e-02, -5.3362e-02, 2.3486e-02, 6.6255e-02],
[-1.9425e-01, -3.4734e-02, 1.1201e-01, 1.3619e-01, -4.1869e-02],
[ 1.0564e-01, -4.7907e-02, 1.5579e-01, 8.5055e-02, -5.6828e-02],
[ 7.6433e-02, -6.6623e-03, 7.2439e-02, 4.4009e-02, 1.3422e-01],
[ 7.8337e-02, -9.4261e-02, -3.6983e-02, 4.5135e-03, -8.3209e-03]]],
[[[ 8.7208e-02, 5.5556e-02, 1.9781e-01, 3.1709e-02, -1.7607e-01],
[ 1.0245e-01, -1.5157e-01, 7.8186e-02, 3.6932e-02, 1.1580e-01],
[ 1.1514e-01, 6.2580e-02, 1.6249e-01, -6.6736e-02, -1.1764e-04],
[ 7.5656e-02, -1.0178e-01, 4.9827e-02, -1.3515e-01, 2.1252e-02],
[-1.7640e-01, -4.0824e-02, 6.4694e-03, -1.7306e-01, -1.5568e-01]]],
[[[-6.5916e-03, -3.8075e-02, -1.6303e-01, 2.2056e-02, 4.0869e-02],
[-7.3824e-02, 1.2453e-01, 1.4930e-01, 1.0905e-01, 5.0146e-02],
[-8.1015e-02, -1.7211e-01, -1.8528e-01, -1.8798e-01, -8.1064e-02],
[ 1.5227e-01, 3.5658e-02, -1.8680e-01, -2.1098e-02, -4.6249e-02],
[-7.7353e-03, -1.1231e-01, -6.6389e-02, -7.9792e-02, 1.4359e-01]]]],
device='cuda:0', requires_grad=True)), ('bias', Parameter containing:
tensor([ 0.0583, -0.1311, -0.1846, 0.0018, -0.1497, 0.0005], device='cuda:0',
requires_grad=True))]
print(list(module.named_buffers()))
[]
剪枝模块#
要剪枝一个模块（在本例中为我们 LeNet 架构的 conv1 层），首先从 torch.nn.utils.prune 中提供的技术中选择一种剪枝技术（或者通过继承 BasePruningMethod 实现你自己的）。然后，指定要剪枝的模块以及该模块中参数的名称。最后，使用所选剪枝技术所需的适当关键字参数来指定剪枝参数。
在本例中，我们将随机剪枝 conv1 层中名为 weight 的参数 30% 的连接。模块作为第一个参数传递给函数；name 使用其字符串标识符识别模块内的参数；amount 表示要剪枝的连接百分比（如果是 0 到 1 之间的浮点数），或者是要剪枝的连接的绝对数量（如果是非负整数）。
prune.random_unstructured(module, name="weight", amount=0.3)
Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
剪枝的作用是从参数中移除 weight，并将其替换为一个名为 weight_orig 的新参数（即在初始参数 name 后附加 "_orig"）。weight_orig 存储张量的未剪枝版本。bias 没有被剪枝，因此保持不变。
print(list(module.named_parameters()))
[('bias', Parameter containing:
tensor([ 0.0583, -0.1311, -0.1846, 0.0018, -0.1497, 0.0005], device='cuda:0',
requires_grad=True)), ('weight_orig', Parameter containing:
tensor([[[[ 1.4951e-01, -1.2675e-01, 6.8154e-02, 8.1491e-02, -5.0470e-02],
[ 9.2922e-02, -1.8541e-01, -1.0656e-01, 1.4048e-01, -2.9564e-02],
[ 1.1579e-01, 5.8812e-02, 3.2663e-02, 1.1807e-01, 5.2710e-02],
[ 1.1038e-02, 3.8310e-02, -1.9490e-02, 6.7628e-02, -6.8343e-02],
[-1.2097e-01, -1.2120e-01, -8.9609e-02, 8.5565e-02, -4.3848e-02]]],
[[[-9.8414e-02, -2.3538e-02, -1.6892e-01, -1.1870e-01, 1.5420e-01],
[ 1.2853e-01, -9.0719e-02, 4.2823e-02, 9.4836e-02, 1.9044e-01],
[-1.2375e-01, -4.1072e-02, 7.7668e-02, -4.0979e-02, 1.5920e-02],
[ 8.3536e-02, -1.5734e-01, 1.8694e-01, 1.9463e-02, -8.3581e-02],
[-9.3909e-02, -6.8686e-03, 1.7742e-01, -3.5721e-02, -5.2636e-02]]],
[[[ 8.2311e-02, 8.7024e-02, 1.8021e-01, -1.7326e-01, 6.9537e-02],
[ 9.2128e-02, 7.4754e-03, 6.0818e-02, 1.6462e-01, -5.8529e-03],
[-1.3501e-01, 5.5084e-03, 1.6291e-01, 1.1028e-02, -1.8746e-01],
[-8.8652e-02, 1.9266e-01, 7.8326e-02, -2.9557e-02, -1.5176e-01],
[ 4.4897e-02, -6.9426e-02, 1.9812e-02, -2.5902e-02, -1.5892e-01]]],
[[[-1.7176e-02, -2.9392e-02, -5.3362e-02, 2.3486e-02, 6.6255e-02],
[-1.9425e-01, -3.4734e-02, 1.1201e-01, 1.3619e-01, -4.1869e-02],
[ 1.0564e-01, -4.7907e-02, 1.5579e-01, 8.5055e-02, -5.6828e-02],
[ 7.6433e-02, -6.6623e-03, 7.2439e-02, 4.4009e-02, 1.3422e-01],
[ 7.8337e-02, -9.4261e-02, -3.6983e-02, 4.5135e-03, -8.3209e-03]]],
[[[ 8.7208e-02, 5.5556e-02, 1.9781e-01, 3.1709e-02, -1.7607e-01],
[ 1.0245e-01, -1.5157e-01, 7.8186e-02, 3.6932e-02, 1.1580e-01],
[ 1.1514e-01, 6.2580e-02, 1.6249e-01, -6.6736e-02, -1.1764e-04],
[ 7.5656e-02, -1.0178e-01, 4.9827e-02, -1.3515e-01, 2.1252e-02],
[-1.7640e-01, -4.0824e-02, 6.4694e-03, -1.7306e-01, -1.5568e-01]]],
[[[-6.5916e-03, -3.8075e-02, -1.6303e-01, 2.2056e-02, 4.0869e-02],
[-7.3824e-02, 1.2453e-01, 1.4930e-01, 1.0905e-01, 5.0146e-02],
[-8.1015e-02, -1.7211e-01, -1.8528e-01, -1.8798e-01, -8.1064e-02],
[ 1.5227e-01, 3.5658e-02, -1.8680e-01, -2.1098e-02, -4.6249e-02],
[-7.7353e-03, -1.1231e-01, -6.6389e-02, -7.9792e-02, 1.4359e-01]]]],
device='cuda:0', requires_grad=True))]
上述剪枝技术生成的剪枝掩码（mask）被保存为名为 weight_mask 的模块缓存（即在初始参数 name 后附加 "_mask"）。
print(list(module.named_buffers()))
[('weight_mask', tensor([[[[1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 0., 0.],
[1., 0., 1., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[1., 1., 1., 1., 0.],
[0., 1., 1., 0., 0.],
[0., 1., 1., 1., 0.],
[0., 1., 0., 1., 1.],
[0., 0., 0., 1., 1.]]],
[[[0., 1., 0., 1., 1.],
[0., 1., 0., 1., 1.],
[1., 1., 0., 1., 1.],
[1., 1., 0., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[0., 1., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 1., 0., 1., 1.],
[0., 1., 1., 1., 1.]]],
[[[1., 1., 1., 0., 1.],
[0., 0., 1., 0., 0.],
[0., 1., 0., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[0., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 1.],
[1., 0., 0., 1., 0.],
[1., 0., 1., 1., 0.]]]], device='cuda:0'))]
为了使前向传播无需修改即可工作，weight 属性需要存在。torch.nn.utils.prune 中实现的剪枝技术会计算权重的剪枝版本（通过将掩码与原始参数结合），并将其存储在 weight 属性中。请注意，这不再是 module 的参数，它现在只是一个属性。
print(module.weight)
tensor([[[[ 0.1495, -0.1267, 0.0682, 0.0815, -0.0000],
[ 0.0929, -0.1854, -0.1066, 0.1405, -0.0296],
[ 0.0000, 0.0588, 0.0327, 0.0000, 0.0000],
[ 0.0110, 0.0000, -0.0195, 0.0676, -0.0683],
[-0.1210, -0.1212, -0.0896, 0.0856, -0.0438]]],
[[[-0.0984, -0.0235, -0.1689, -0.1187, 0.0000],
[ 0.0000, -0.0907, 0.0428, 0.0000, 0.0000],
[-0.0000, -0.0411, 0.0777, -0.0410, 0.0000],
[ 0.0000, -0.1573, 0.0000, 0.0195, -0.0836],
[-0.0000, -0.0000, 0.0000, -0.0357, -0.0526]]],
[[[ 0.0000, 0.0870, 0.0000, -0.1733, 0.0695],
[ 0.0000, 0.0075, 0.0000, 0.1646, -0.0059],
[-0.1350, 0.0055, 0.0000, 0.0110, -0.1875],
[-0.0887, 0.1927, 0.0000, -0.0296, -0.1518],
[ 0.0449, -0.0694, 0.0198, -0.0259, -0.1589]]],
[[[-0.0000, -0.0294, -0.0534, 0.0235, 0.0663],
[-0.0000, -0.0000, 0.1120, 0.1362, -0.0419],
[ 0.0000, -0.0000, 0.1558, 0.0851, -0.0568],
[ 0.0000, -0.0067, 0.0000, 0.0440, 0.1342],
[ 0.0000, -0.0943, -0.0370, 0.0045, -0.0083]]],
[[[ 0.0872, 0.0556, 0.1978, 0.0000, -0.1761],
[ 0.0000, -0.0000, 0.0782, 0.0000, 0.0000],
[ 0.0000, 0.0626, 0.0000, -0.0667, -0.0000],
[ 0.0757, -0.1018, 0.0498, -0.1351, 0.0213],
[-0.1764, -0.0408, 0.0065, -0.1731, -0.1557]]],
[[[-0.0000, -0.0381, -0.1630, 0.0221, 0.0409],
[-0.0738, 0.1245, 0.1493, 0.1090, 0.0501],
[-0.0000, -0.1721, -0.1853, -0.1880, -0.0811],
[ 0.1523, 0.0000, -0.0000, -0.0211, -0.0000],
[-0.0077, -0.0000, -0.0664, -0.0798, 0.0000]]]], device='cuda:0',
grad_fn=<MulBackward0>)
最后，剪枝是在每次前向传播之前使用 PyTorch 的 forward_pre_hooks 应用的。具体来说，当 module 被剪枝时，它将为每个被剪枝的关联参数获得一个 forward_pre_hook。在本例中，由于我们目前只剪枝了名为 weight 的原始参数，因此只会存在一个钩子。
print(module._forward_pre_hooks)
OrderedDict([(0, <torch.nn.utils.prune.RandomUnstructured object at 0x7ff8d3e7faf0>)])
为了完整起见，我们现在也可以剪枝 bias，看看模块的参数、缓存、钩子和属性是如何变化的。为了尝试另一种剪枝技术，这里我们根据 L1 范数剪枝偏置中 3 个最小的条目，正如 l1_unstructured 剪枝函数中所实现的那样。
prune.l1_unstructured(module, name="bias", amount=3)
Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
我们现在预计命名参数将包括 weight_orig（之前的）和 bias_orig。缓存将包括 weight_mask 和 bias_mask。这两个张量的剪枝版本将作为模块属性存在，并且模块现在将拥有两个 forward_pre_hooks。
print(list(module.named_parameters()))
[('weight_orig', Parameter containing:
tensor([[[[ 1.4951e-01, -1.2675e-01, 6.8154e-02, 8.1491e-02, -5.0470e-02],
[ 9.2922e-02, -1.8541e-01, -1.0656e-01, 1.4048e-01, -2.9564e-02],
[ 1.1579e-01, 5.8812e-02, 3.2663e-02, 1.1807e-01, 5.2710e-02],
[ 1.1038e-02, 3.8310e-02, -1.9490e-02, 6.7628e-02, -6.8343e-02],
[-1.2097e-01, -1.2120e-01, -8.9609e-02, 8.5565e-02, -4.3848e-02]]],
[[[-9.8414e-02, -2.3538e-02, -1.6892e-01, -1.1870e-01, 1.5420e-01],
[ 1.2853e-01, -9.0719e-02, 4.2823e-02, 9.4836e-02, 1.9044e-01],
[-1.2375e-01, -4.1072e-02, 7.7668e-02, -4.0979e-02, 1.5920e-02],
[ 8.3536e-02, -1.5734e-01, 1.8694e-01, 1.9463e-02, -8.3581e-02],
[-9.3909e-02, -6.8686e-03, 1.7742e-01, -3.5721e-02, -5.2636e-02]]],
[[[ 8.2311e-02, 8.7024e-02, 1.8021e-01, -1.7326e-01, 6.9537e-02],
[ 9.2128e-02, 7.4754e-03, 6.0818e-02, 1.6462e-01, -5.8529e-03],
[-1.3501e-01, 5.5084e-03, 1.6291e-01, 1.1028e-02, -1.8746e-01],
[-8.8652e-02, 1.9266e-01, 7.8326e-02, -2.9557e-02, -1.5176e-01],
[ 4.4897e-02, -6.9426e-02, 1.9812e-02, -2.5902e-02, -1.5892e-01]]],
[[[-1.7176e-02, -2.9392e-02, -5.3362e-02, 2.3486e-02, 6.6255e-02],
[-1.9425e-01, -3.4734e-02, 1.1201e-01, 1.3619e-01, -4.1869e-02],
[ 1.0564e-01, -4.7907e-02, 1.5579e-01, 8.5055e-02, -5.6828e-02],
[ 7.6433e-02, -6.6623e-03, 7.2439e-02, 4.4009e-02, 1.3422e-01],
[ 7.8337e-02, -9.4261e-02, -3.6983e-02, 4.5135e-03, -8.3209e-03]]],
[[[ 8.7208e-02, 5.5556e-02, 1.9781e-01, 3.1709e-02, -1.7607e-01],
[ 1.0245e-01, -1.5157e-01, 7.8186e-02, 3.6932e-02, 1.1580e-01],
[ 1.1514e-01, 6.2580e-02, 1.6249e-01, -6.6736e-02, -1.1764e-04],
[ 7.5656e-02, -1.0178e-01, 4.9827e-02, -1.3515e-01, 2.1252e-02],
[-1.7640e-01, -4.0824e-02, 6.4694e-03, -1.7306e-01, -1.5568e-01]]],
[[[-6.5916e-03, -3.8075e-02, -1.6303e-01, 2.2056e-02, 4.0869e-02],
[-7.3824e-02, 1.2453e-01, 1.4930e-01, 1.0905e-01, 5.0146e-02],
[-8.1015e-02, -1.7211e-01, -1.8528e-01, -1.8798e-01, -8.1064e-02],
[ 1.5227e-01, 3.5658e-02, -1.8680e-01, -2.1098e-02, -4.6249e-02],
[-7.7353e-03, -1.1231e-01, -6.6389e-02, -7.9792e-02, 1.4359e-01]]]],
device='cuda:0', requires_grad=True)), ('bias_orig', Parameter containing:
tensor([ 0.0583, -0.1311, -0.1846, 0.0018, -0.1497, 0.0005], device='cuda:0',
requires_grad=True))]
print(list(module.named_buffers()))
[('weight_mask', tensor([[[[1., 1., 1., 1., 0.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 0., 0.],
[1., 0., 1., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[1., 1., 1., 1., 0.],
[0., 1., 1., 0., 0.],
[0., 1., 1., 1., 0.],
[0., 1., 0., 1., 1.],
[0., 0., 0., 1., 1.]]],
[[[0., 1., 0., 1., 1.],
[0., 1., 0., 1., 1.],
[1., 1., 0., 1., 1.],
[1., 1., 0., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[0., 1., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 1., 0., 1., 1.],
[0., 1., 1., 1., 1.]]],
[[[1., 1., 1., 0., 1.],
[0., 0., 1., 0., 0.],
[0., 1., 0., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[0., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 1.],
[1., 0., 0., 1., 0.],
[1., 0., 1., 1., 0.]]]], device='cuda:0')), ('bias_mask', tensor([0., 1., 1., 0., 1., 0.], device='cuda:0'))]
print(module.bias)
tensor([ 0.0000, -0.1311, -0.1846, 0.0000, -0.1497, 0.0000], device='cuda:0',
grad_fn=<MulBackward0>)
print(module._forward_pre_hooks)
OrderedDict([(0, <torch.nn.utils.prune.RandomUnstructured object at 0x7ff8d3e7faf0>), (1, <torch.nn.utils.prune.L1Unstructured object at 0x7ff8d3e7d7e0>)])
迭代剪枝#
模块中的同一个参数可以被多次剪枝，各种剪枝调用的效果等同于串联应用各种掩码。新掩码与旧掩码的组合由 PruningContainer 的 compute_mask 方法处理。
例如，假设我们现在想进一步剪枝 module.weight，这次使用沿张量第 0 轴的结构化剪枝（第 0 轴对应于卷积层的输出通道，对于 conv1 来说维度为 6），基于通道的 L2 范数。这可以通过使用 ln_structured 函数来实现，其中 n=2 且 dim=0。
prune.ln_structured(module, name="weight", amount=0.5, n=2, dim=0)
# As we can verify, this will zero out all the connections corresponding to
# 50% (3 out of 6) of the channels, while preserving the action of the
# previous mask.
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.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.0870, 0.0000, -0.1733, 0.0695],
[ 0.0000, 0.0075, 0.0000, 0.1646, -0.0059],
[-0.1350, 0.0055, 0.0000, 0.0110, -0.1875],
[-0.0887, 0.1927, 0.0000, -0.0296, -0.1518],
[ 0.0449, -0.0694, 0.0198, -0.0259, -0.1589]]],
[[[-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.0872, 0.0556, 0.1978, 0.0000, -0.1761],
[ 0.0000, -0.0000, 0.0782, 0.0000, 0.0000],
[ 0.0000, 0.0626, 0.0000, -0.0667, -0.0000],
[ 0.0757, -0.1018, 0.0498, -0.1351, 0.0213],
[-0.1764, -0.0408, 0.0065, -0.1731, -0.1557]]],
[[[-0.0000, -0.0381, -0.1630, 0.0221, 0.0409],
[-0.0738, 0.1245, 0.1493, 0.1090, 0.0501],
[-0.0000, -0.1721, -0.1853, -0.1880, -0.0811],
[ 0.1523, 0.0000, -0.0000, -0.0211, -0.0000],
[-0.0077, -0.0000, -0.0664, -0.0798, 0.0000]]]], device='cuda:0',
grad_fn=<MulBackward0>)
相应的钩子现在将是 torch.nn.utils.prune.PruningContainer 类型，并将存储应用于 weight 参数的剪枝历史。
[<torch.nn.utils.prune.RandomUnstructured object at 0x7ff8d3e7faf0>, <torch.nn.utils.prune.LnStructured object at 0x7ff8d3e7fca0>]
序列化已剪枝模型#
所有相关张量，包括用于计算剪枝张量的掩码缓存和原始参数，都存储在模型的 state_dict 中，因此如果需要，可以轻松序列化和保存。
print(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 和 weight_mask 方面的重参数化，并移除 forward_pre_hook，我们可以使用 torch.nn.utils.prune 中的 remove 功能。请注意，这不会撤销剪枝，就像从未发生过一样。相反，它通过将参数 weight 以其剪枝版本重新分配给模型参数，从而使其永久化。
移除重参数化之前
print(list(module.named_parameters()))
[('weight_orig', Parameter containing:
tensor([[[[ 1.4951e-01, -1.2675e-01, 6.8154e-02, 8.1491e-02, -5.0470e-02],
[ 9.2922e-02, -1.8541e-01, -1.0656e-01, 1.4048e-01, -2.9564e-02],
[ 1.1579e-01, 5.8812e-02, 3.2663e-02, 1.1807e-01, 5.2710e-02],
[ 1.1038e-02, 3.8310e-02, -1.9490e-02, 6.7628e-02, -6.8343e-02],
[-1.2097e-01, -1.2120e-01, -8.9609e-02, 8.5565e-02, -4.3848e-02]]],
[[[-9.8414e-02, -2.3538e-02, -1.6892e-01, -1.1870e-01, 1.5420e-01],
[ 1.2853e-01, -9.0719e-02, 4.2823e-02, 9.4836e-02, 1.9044e-01],
[-1.2375e-01, -4.1072e-02, 7.7668e-02, -4.0979e-02, 1.5920e-02],
[ 8.3536e-02, -1.5734e-01, 1.8694e-01, 1.9463e-02, -8.3581e-02],
[-9.3909e-02, -6.8686e-03, 1.7742e-01, -3.5721e-02, -5.2636e-02]]],
[[[ 8.2311e-02, 8.7024e-02, 1.8021e-01, -1.7326e-01, 6.9537e-02],
[ 9.2128e-02, 7.4754e-03, 6.0818e-02, 1.6462e-01, -5.8529e-03],
[-1.3501e-01, 5.5084e-03, 1.6291e-01, 1.1028e-02, -1.8746e-01],
[-8.8652e-02, 1.9266e-01, 7.8326e-02, -2.9557e-02, -1.5176e-01],
[ 4.4897e-02, -6.9426e-02, 1.9812e-02, -2.5902e-02, -1.5892e-01]]],
[[[-1.7176e-02, -2.9392e-02, -5.3362e-02, 2.3486e-02, 6.6255e-02],
[-1.9425e-01, -3.4734e-02, 1.1201e-01, 1.3619e-01, -4.1869e-02],
[ 1.0564e-01, -4.7907e-02, 1.5579e-01, 8.5055e-02, -5.6828e-02],
[ 7.6433e-02, -6.6623e-03, 7.2439e-02, 4.4009e-02, 1.3422e-01],
[ 7.8337e-02, -9.4261e-02, -3.6983e-02, 4.5135e-03, -8.3209e-03]]],
[[[ 8.7208e-02, 5.5556e-02, 1.9781e-01, 3.1709e-02, -1.7607e-01],
[ 1.0245e-01, -1.5157e-01, 7.8186e-02, 3.6932e-02, 1.1580e-01],
[ 1.1514e-01, 6.2580e-02, 1.6249e-01, -6.6736e-02, -1.1764e-04],
[ 7.5656e-02, -1.0178e-01, 4.9827e-02, -1.3515e-01, 2.1252e-02],
[-1.7640e-01, -4.0824e-02, 6.4694e-03, -1.7306e-01, -1.5568e-01]]],
[[[-6.5916e-03, -3.8075e-02, -1.6303e-01, 2.2056e-02, 4.0869e-02],
[-7.3824e-02, 1.2453e-01, 1.4930e-01, 1.0905e-01, 5.0146e-02],
[-8.1015e-02, -1.7211e-01, -1.8528e-01, -1.8798e-01, -8.1064e-02],
[ 1.5227e-01, 3.5658e-02, -1.8680e-01, -2.1098e-02, -4.6249e-02],
[-7.7353e-03, -1.1231e-01, -6.6389e-02, -7.9792e-02, 1.4359e-01]]]],
device='cuda:0', requires_grad=True)), ('bias_orig', Parameter containing:
tensor([ 0.0583, -0.1311, -0.1846, 0.0018, -0.1497, 0.0005], device='cuda:0',
requires_grad=True))]
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.]]],
[[[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., 0., 1., 1.],
[0., 1., 0., 1., 1.],
[1., 1., 0., 1., 1.],
[1., 1., 0., 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., 0., 1.],
[0., 0., 1., 0., 0.],
[0., 1., 0., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.]]],
[[[0., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 1.],
[1., 0., 0., 1., 0.],
[1., 0., 1., 1., 0.]]]], device='cuda:0')), ('bias_mask', tensor([0., 1., 1., 0., 1., 0.], device='cuda:0'))]
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.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.0870, 0.0000, -0.1733, 0.0695],
[ 0.0000, 0.0075, 0.0000, 0.1646, -0.0059],
[-0.1350, 0.0055, 0.0000, 0.0110, -0.1875],
[-0.0887, 0.1927, 0.0000, -0.0296, -0.1518],
[ 0.0449, -0.0694, 0.0198, -0.0259, -0.1589]]],
[[[-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.0872, 0.0556, 0.1978, 0.0000, -0.1761],
[ 0.0000, -0.0000, 0.0782, 0.0000, 0.0000],
[ 0.0000, 0.0626, 0.0000, -0.0667, -0.0000],
[ 0.0757, -0.1018, 0.0498, -0.1351, 0.0213],
[-0.1764, -0.0408, 0.0065, -0.1731, -0.1557]]],
[[[-0.0000, -0.0381, -0.1630, 0.0221, 0.0409],
[-0.0738, 0.1245, 0.1493, 0.1090, 0.0501],
[-0.0000, -0.1721, -0.1853, -0.1880, -0.0811],
[ 0.1523, 0.0000, -0.0000, -0.0211, -0.0000],
[-0.0077, -0.0000, -0.0664, -0.0798, 0.0000]]]], device='cuda:0',
grad_fn=<MulBackward0>)
移除重参数化之后
prune.remove(module, 'weight')
print(list(module.named_parameters()))
[('bias_orig', Parameter containing:
tensor([ 0.0583, -0.1311, -0.1846, 0.0018, -0.1497, 0.0005], device='cuda:0',
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.0000, 0.0870, 0.0000, -0.1733, 0.0695],
[ 0.0000, 0.0075, 0.0000, 0.1646, -0.0059],
[-0.1350, 0.0055, 0.0000, 0.0110, -0.1875],
[-0.0887, 0.1927, 0.0000, -0.0296, -0.1518],
[ 0.0449, -0.0694, 0.0198, -0.0259, -0.1589]]],
[[[-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.0872, 0.0556, 0.1978, 0.0000, -0.1761],
[ 0.0000, -0.0000, 0.0782, 0.0000, 0.0000],
[ 0.0000, 0.0626, 0.0000, -0.0667, -0.0000],
[ 0.0757, -0.1018, 0.0498, -0.1351, 0.0213],
[-0.1764, -0.0408, 0.0065, -0.1731, -0.1557]]],
[[[-0.0000, -0.0381, -0.1630, 0.0221, 0.0409],
[-0.0738, 0.1245, 0.1493, 0.1090, 0.0501],
[-0.0000, -0.1721, -0.1853, -0.1880, -0.0811],
[ 0.1523, 0.0000, -0.0000, -0.0211, -0.0000],
[-0.0077, -0.0000, -0.0664, -0.0798, 0.0000]]]], device='cuda:0',
requires_grad=True))]
print(list(module.named_buffers()))
[('bias_mask', tensor([0., 1., 1., 0., 1., 0.], device='cuda:0'))]
剪枝模型中的多个参数#
通过指定所需的剪枝技术和参数，我们可以轻松地剪枝网络中的多个张量，也许根据它们的类型，正如我们将在本例中看到的。
new_model = LeNet()
for name, module in new_model.named_modules():
# prune 20% of connections in all 2D-conv layers
if isinstance(module, torch.nn.Conv2d):
prune.l1_unstructured(module, name='weight', amount=0.2)
# prune 40% of connections in all linear layers
elif isinstance(module, torch.nn.Linear):
prune.l1_unstructured(module, name='weight', amount=0.4)
print(dict(new_model.named_buffers()).keys()) # to verify that all masks exist
dict_keys(['conv1.weight_mask', 'conv2.weight_mask', 'fc1.weight_mask', 'fc2.weight_mask', 'fc3.weight_mask'])
全局剪枝#
到目前为止，我们只研究了通常所说的“局部”剪枝，即通过仅比较该张量中每个条目的统计数据（权重幅度、激活值、梯度等）来逐个剪枝模型中张量的做法。然而，一种常见且可能更强大的技术是一次性剪枝整个模型，通过移除（例如）整个模型中 20% 最低的连接，而不是移除每一层中 20% 最低的连接。这很可能会导致每层的剪枝百分比不同。让我们看看如何使用 torch.nn.utils.prune 中的 global_unstructured 来做到这一点。
model = LeNet()
parameters_to_prune = (
(model.conv1, 'weight'),
(model.conv2, 'weight'),
(model.fc1, 'weight'),
(model.fc2, 'weight'),
(model.fc3, 'weight'),
)
prune.global_unstructured(
parameters_to_prune,
pruning_method=prune.L1Unstructured,
amount=0.2,
)
现在我们可以检查每个被剪枝参数中引起的稀疏度，它在每一层中不会等于 20%。但是，全局稀疏度将（大约）为 20%。
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(
"Sparsity in fc3.weight: {:.2f}%".format(
100. * float(torch.sum(model.fc3.weight == 0))
/ float(model.fc3.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)
+ torch.sum(model.fc3.weight == 0)
)
/ float(
model.conv1.weight.nelement()
+ model.conv2.weight.nelement()
+ model.fc1.weight.nelement()
+ model.fc2.weight.nelement()
+ model.fc3.weight.nelement()
)
)
)
Sparsity in conv1.weight: 4.67%
Sparsity in conv2.weight: 13.67%
Sparsity in fc1.weight: 22.12%
Sparsity in fc2.weight: 12.48%
Sparsity in fc3.weight: 9.88%
Global sparsity: 20.00%
使用自定义剪枝函数扩展 torch.nn.utils.prune#
要实现你自己的剪枝函数，你可以通过继承 BasePruningMethod 基类来扩展 nn.utils.prune 模块，就像所有其他剪枝方法所做的那样。基类为你实现了以下方法：__call__, apply_mask, apply, prune 和 remove。除了某些特殊情况外，你不应该需要为你新的剪枝技术重新实现这些方法。但是，你需要实现 __init__（构造函数）和 compute_mask（关于如何根据你的剪枝技术逻辑为给定张量计算掩码的指令）。此外，你需要指定此技术实现哪种类型的剪枝（支持的选项有 global、structured 和 unstructured）。这是为了在迭代应用剪枝的情况下确定如何组合掩码所必需的。换句话说，当剪枝一个已经预剪枝的参数时，当前的剪枝技术预计会作用于参数的未剪枝部分。指定 PRUNING_TYPE 将使 PruningContainer（它处理剪枝掩码的迭代应用）能够正确识别要剪枝的参数切片。
例如，假设你想实现一种剪枝技术，该技术剪枝张量中的每隔一个条目（或者如果张量之前已被剪枝，则在张量剩余的未剪枝部分中）。这将是 PRUNING_TYPE='unstructured'，因为它作用于层中的单个连接，而不是整个单元/通道（'structured'），或者跨不同的参数（'global'）。
class FooBarPruningMethod(prune.BasePruningMethod):
"""Prune every other entry in a tensor
"""
PRUNING_TYPE = 'unstructured'
def compute_mask(self, t, default_mask):
mask = default_mask.clone()
mask.view(-1)[::2] = 0
return mask
现在，要将其应用于 nn.Module 中的参数，你还应该提供一个简单的函数来实例化该方法并应用它。
def foobar_unstructured(module, name):
"""Prunes tensor corresponding to parameter called `name` in `module`
by removing every other entry in the tensors.
Modifies module in place (and also return the modified module)
by:
1) adding a named buffer called `name+'_mask'` corresponding to the
binary mask applied to the parameter `name` by the pruning method.
The parameter `name` is replaced by its pruned version, while the
original (unpruned) parameter is stored in a new parameter named
`name+'_orig'`.
Args:
module (nn.Module): module containing the tensor to prune
name (string): parameter name within `module` on which pruning
will act.
Returns:
module (nn.Module): modified (i.e. pruned) version of the input
module
Examples:
>>> m = nn.Linear(3, 4)
>>> foobar_unstructured(m, name='bias')
"""
FooBarPruningMethod.apply(module, name)
return module
让我们试一试！
model = LeNet()
foobar_unstructured(model.fc3, name='bias')
print(model.fc3.bias_mask)
tensor([0., 1., 0., 1., 0., 1., 0., 1., 0., 1.])
脚本的总运行时间： (0 分钟 0.545 秒)