注意
点击 这里 下载完整的示例代码
剪枝教程
创建于：2019年7月22日 | 最后更新：2023年11月2日 | 最后验证：2024年11月5日
作者: Michela Paganini
最先进的深度学习技术依赖于难以部署的过度参数化模型。相反，生物神经网络已知使用高效的稀疏连接。为了在不牺牲准确性的情况下减少内存、电池和硬件消耗，识别通过减少参数数量来压缩模型的最佳技术非常重要。这反过来又允许您在设备上部署轻量级模型，并通过设备上的私有计算保证隐私。在研究方面，剪枝用于研究过度参数化和欠参数化网络之间学习动态的差异，研究幸运稀疏子网络和初始化（“lottery tickets”）作为破坏性神经架构搜索技术的作用，以及更多。
在本教程中，您将学习如何使用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，并且没有缓冲区。
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.0922, -0.0565, -0.1203, 0.0189, -0.1975],
[ 0.1806, -0.1699, 0.1544, 0.0333, -0.0649],
[ 0.1236, 0.0312, 0.1616, 0.0219, -0.0631],
[ 0.0537, -0.0542, 0.0842, 0.1786, 0.1156],
[-0.0874, 0.1155, 0.0358, 0.1016, -0.1219]]],
[[[-0.1980, -0.0773, -0.1534, 0.1641, 0.0576],
[ 0.0828, 0.0633, -0.0035, 0.1565, -0.1421],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0667, 0.1925, -0.1651, -0.1984]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.1033, -0.1363, 0.1061, -0.0808, 0.1214],
[-0.0475, 0.1144, -0.1554, -0.1009, 0.0610],
[ 0.0423, -0.0510, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0675, -0.0709, -0.1935]]],
[[[-0.1145, 0.0500, -0.0264, -0.1452, 0.0047],
[-0.1366, -0.1697, -0.1101, -0.1750, -0.1273],
[ 0.1999, 0.0378, 0.0616, -0.1865, -0.1314],
[-0.0666, 0.0313, -0.1760, -0.0862, -0.1197],
[ 0.0006, -0.0744, -0.0139, -0.1355, -0.1373]]],
[[[-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))]
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.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.0922, -0.0565, -0.1203, 0.0189, -0.1975],
[ 0.1806, -0.1699, 0.1544, 0.0333, -0.0649],
[ 0.1236, 0.0312, 0.1616, 0.0219, -0.0631],
[ 0.0537, -0.0542, 0.0842, 0.1786, 0.1156],
[-0.0874, 0.1155, 0.0358, 0.1016, -0.1219]]],
[[[-0.1980, -0.0773, -0.1534, 0.1641, 0.0576],
[ 0.0828, 0.0633, -0.0035, 0.1565, -0.1421],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0667, 0.1925, -0.1651, -0.1984]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.1033, -0.1363, 0.1061, -0.0808, 0.1214],
[-0.0475, 0.1144, -0.1554, -0.1009, 0.0610],
[ 0.0423, -0.0510, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0675, -0.0709, -0.1935]]],
[[[-0.1145, 0.0500, -0.0264, -0.1452, 0.0047],
[-0.1366, -0.1697, -0.1101, -0.1750, -0.1273],
[ 0.1999, 0.0378, 0.0616, -0.1865, -0.1314],
[-0.0666, 0.0313, -0.1760, -0.0862, -0.1197],
[ 0.0006, -0.0744, -0.0139, -0.1355, -0.1373]]],
[[[-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))]
由上述选择的剪枝技术生成的剪枝掩码被保存为一个名为weight_mask的模块缓冲区（即在初始参数name后附加"_mask"）。
print(list(module.named_buffers()))
[('weight_mask', tensor([[[[1., 1., 1., 1., 1.],
[1., 0., 1., 1., 1.],
[1., 0., 0., 1., 1.],
[1., 0., 1., 1., 1.],
[1., 0., 0., 1., 1.]]],
[[[1., 1., 1., 0., 1.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 0.],
[1., 1., 0., 1., 0.],
[0., 1., 0., 1., 1.]]],
[[[1., 0., 0., 0., 1.],
[1., 0., 1., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[1., 0., 1., 1., 0.]]],
[[[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 0.],
[1., 1., 1., 0., 1.],
[0., 0., 1., 1., 1.],
[1., 1., 0., 1., 1.]]],
[[[1., 0., 1., 1., 1.],
[1., 1., 0., 0., 0.],
[1., 1., 0., 0., 0.],
[0., 1., 1., 0., 1.],
[1., 0., 0., 0., 1.]]],
[[[1., 0., 1., 0., 1.],
[0., 1., 1., 1., 1.],
[1., 1., 0., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 0., 0., 1., 1.]]]], device='cuda:0'))]
为了使前向传播无需修改即可工作，weight 属性需要存在。torch.nn.utils.prune 中实现的剪枝技术计算剪枝后的权重（通过将掩码与原始参数结合）并将其存储在 weight 属性中。请注意，这不再是 module 的参数，现在它只是一个属性。
print(module.weight)
tensor([[[[ 0.1529, 0.1660, -0.0469, 0.1837, -0.0438],
[ 0.0404, -0.0000, 0.1175, 0.1763, -0.1467],
[ 0.1738, 0.0000, 0.0000, 0.0271, 0.0964],
[-0.0282, 0.0000, 0.0296, -0.0934, 0.0510],
[-0.0921, -0.0000, -0.0000, 0.1327, -0.1579]]],
[[[-0.0922, -0.0565, -0.1203, 0.0000, -0.1975],
[ 0.1806, -0.1699, 0.1544, 0.0333, -0.0649],
[ 0.0000, 0.0312, 0.1616, 0.0219, -0.0000],
[ 0.0537, -0.0542, 0.0000, 0.1786, 0.0000],
[-0.0000, 0.1155, 0.0000, 0.1016, -0.1219]]],
[[[-0.1980, -0.0000, -0.0000, 0.0000, 0.0576],
[ 0.0828, 0.0000, -0.0035, 0.1565, -0.0000],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0000, 0.1925, -0.1651, -0.0000]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.0000, -0.1363, 0.1061, -0.0808, 0.0000],
[-0.0475, 0.1144, -0.1554, -0.0000, 0.0610],
[ 0.0000, -0.0000, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0000, -0.0709, -0.1935]]],
[[[-0.1145, 0.0000, -0.0264, -0.1452, 0.0047],
[-0.1366, -0.1697, -0.0000, -0.0000, -0.0000],
[ 0.1999, 0.0378, 0.0000, -0.0000, -0.0000],
[-0.0000, 0.0313, -0.1760, -0.0000, -0.1197],
[ 0.0006, -0.0000, -0.0000, -0.0000, -0.1373]]],
[[[-0.1167, -0.0000, -0.1579, 0.0000, -0.0397],
[ 0.0000, 0.0623, -0.1694, 0.1384, -0.0550],
[-0.0767, -0.1660, -0.0000, 0.0572, -0.0000],
[ 0.0779, -0.1641, 0.1485, -0.1468, -0.0345],
[ 0.0418, 0.0000, 0.0000, 0.1822, -0.1586]]]], 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 0x7fce471ea0b0>)])
为了完整性，我们现在也可以修剪bias，看看module的参数、缓冲区、钩子和属性如何变化。只是为了尝试另一种修剪技术，这里我们通过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([[[[ 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.0922, -0.0565, -0.1203, 0.0189, -0.1975],
[ 0.1806, -0.1699, 0.1544, 0.0333, -0.0649],
[ 0.1236, 0.0312, 0.1616, 0.0219, -0.0631],
[ 0.0537, -0.0542, 0.0842, 0.1786, 0.1156],
[-0.0874, 0.1155, 0.0358, 0.1016, -0.1219]]],
[[[-0.1980, -0.0773, -0.1534, 0.1641, 0.0576],
[ 0.0828, 0.0633, -0.0035, 0.1565, -0.1421],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0667, 0.1925, -0.1651, -0.1984]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.1033, -0.1363, 0.1061, -0.0808, 0.1214],
[-0.0475, 0.1144, -0.1554, -0.1009, 0.0610],
[ 0.0423, -0.0510, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0675, -0.0709, -0.1935]]],
[[[-0.1145, 0.0500, -0.0264, -0.1452, 0.0047],
[-0.1366, -0.1697, -0.1101, -0.1750, -0.1273],
[ 0.1999, 0.0378, 0.0616, -0.1865, -0.1314],
[-0.0666, 0.0313, -0.1760, -0.0862, -0.1197],
[ 0.0006, -0.0744, -0.0139, -0.1355, -0.1373]]],
[[[-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_orig', Parameter containing:
tensor([ 0.0503, -0.0860, -0.0219, -0.1497, 0.1822, -0.1468], device='cuda:0',
requires_grad=True))]
print(list(module.named_buffers()))
[('weight_mask', tensor([[[[1., 1., 1., 1., 1.],
[1., 0., 1., 1., 1.],
[1., 0., 0., 1., 1.],
[1., 0., 1., 1., 1.],
[1., 0., 0., 1., 1.]]],
[[[1., 1., 1., 0., 1.],
[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 0.],
[1., 1., 0., 1., 0.],
[0., 1., 0., 1., 1.]]],
[[[1., 0., 0., 0., 1.],
[1., 0., 1., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[1., 0., 1., 1., 0.]]],
[[[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 0.],
[1., 1., 1., 0., 1.],
[0., 0., 1., 1., 1.],
[1., 1., 0., 1., 1.]]],
[[[1., 0., 1., 1., 1.],
[1., 1., 0., 0., 0.],
[1., 1., 0., 0., 0.],
[0., 1., 1., 0., 1.],
[1., 0., 0., 0., 1.]]],
[[[1., 0., 1., 0., 1.],
[0., 1., 1., 1., 1.],
[1., 1., 0., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 0., 0., 1., 1.]]]], device='cuda:0')), ('bias_mask', tensor([0., 0., 0., 1., 1., 1.], device='cuda:0'))]
print(module.bias)
tensor([ 0.0000, -0.0000, -0.0000, -0.1497, 0.1822, -0.1468], device='cuda:0',
grad_fn=<MulBackward0>)
print(module._forward_pre_hooks)
OrderedDict([(0, <torch.nn.utils.prune.RandomUnstructured object at 0x7fce471ea0b0>), (1, <torch.nn.utils.prune.L1Unstructured object at 0x7fce471e98a0>)])
迭代剪枝
模块中的相同参数可以被多次修剪，各种修剪调用的效果等同于连续应用的各种掩码的组合。
新掩码与旧掩码的组合由PruningContainer的compute_mask方法处理。
例如，假设我们现在想要进一步修剪module.weight，这次基于通道的L2范数，沿着张量的第0轴（第0轴对应于卷积层的输出通道，对于conv1具有维度6）进行结构化修剪。这可以通过使用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.1980, -0.0000, -0.0000, 0.0000, 0.0576],
[ 0.0828, 0.0000, -0.0035, 0.1565, -0.0000],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0000, 0.1925, -0.1651, -0.0000]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.0000, -0.1363, 0.1061, -0.0808, 0.0000],
[-0.0475, 0.1144, -0.1554, -0.0000, 0.0610],
[ 0.0000, -0.0000, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0000, -0.0709, -0.1935]]],
[[[-0.0000, 0.0000, -0.0000, -0.0000, 0.0000],
[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000],
[ 0.0000, 0.0000, 0.0000, -0.0000, -0.0000],
[-0.0000, 0.0000, -0.0000, -0.0000, -0.0000],
[ 0.0000, -0.0000, -0.0000, -0.0000, -0.0000]]],
[[[-0.1167, -0.0000, -0.1579, 0.0000, -0.0397],
[ 0.0000, 0.0623, -0.1694, 0.1384, -0.0550],
[-0.0767, -0.1660, -0.0000, 0.0572, -0.0000],
[ 0.0779, -0.1641, 0.1485, -0.1468, -0.0345],
[ 0.0418, 0.0000, 0.0000, 0.1822, -0.1586]]]], device='cuda:0',
grad_fn=<MulBackward0>)
相应的钩子现在将是类型
torch.nn.utils.prune.PruningContainer，并将存储应用于
weight参数的修剪历史。
[<torch.nn.utils.prune.RandomUnstructured object at 0x7fce471ea0b0>, <torch.nn.utils.prune.LnStructured object at 0x7fce471e9a20>]
序列化一个剪枝模型
所有相关的张量，包括掩码缓冲区和用于计算修剪张量的原始参数，都存储在模型的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([[[[ 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.0922, -0.0565, -0.1203, 0.0189, -0.1975],
[ 0.1806, -0.1699, 0.1544, 0.0333, -0.0649],
[ 0.1236, 0.0312, 0.1616, 0.0219, -0.0631],
[ 0.0537, -0.0542, 0.0842, 0.1786, 0.1156],
[-0.0874, 0.1155, 0.0358, 0.1016, -0.1219]]],
[[[-0.1980, -0.0773, -0.1534, 0.1641, 0.0576],
[ 0.0828, 0.0633, -0.0035, 0.1565, -0.1421],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0667, 0.1925, -0.1651, -0.1984]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.1033, -0.1363, 0.1061, -0.0808, 0.1214],
[-0.0475, 0.1144, -0.1554, -0.1009, 0.0610],
[ 0.0423, -0.0510, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0675, -0.0709, -0.1935]]],
[[[-0.1145, 0.0500, -0.0264, -0.1452, 0.0047],
[-0.1366, -0.1697, -0.1101, -0.1750, -0.1273],
[ 0.1999, 0.0378, 0.0616, -0.1865, -0.1314],
[-0.0666, 0.0313, -0.1760, -0.0862, -0.1197],
[ 0.0006, -0.0744, -0.0139, -0.1355, -0.1373]]],
[[[-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_orig', Parameter containing:
tensor([ 0.0503, -0.0860, -0.0219, -0.1497, 0.1822, -0.1468], 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.]]],
[[[1., 0., 0., 0., 1.],
[1., 0., 1., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[1., 0., 1., 1., 0.]]],
[[[1., 1., 1., 1., 1.],
[0., 1., 1., 1., 0.],
[1., 1., 1., 0., 1.],
[0., 0., 1., 1., 1.],
[1., 1., 0., 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., 0., 1., 0., 1.],
[0., 1., 1., 1., 1.],
[1., 1., 0., 1., 0.],
[1., 1., 1., 1., 1.],
[1., 0., 0., 1., 1.]]]], device='cuda:0')), ('bias_mask', tensor([0., 0., 0., 1., 1., 1.], 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.1980, -0.0000, -0.0000, 0.0000, 0.0576],
[ 0.0828, 0.0000, -0.0035, 0.1565, -0.0000],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0000, 0.1925, -0.1651, -0.0000]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.0000, -0.1363, 0.1061, -0.0808, 0.0000],
[-0.0475, 0.1144, -0.1554, -0.0000, 0.0610],
[ 0.0000, -0.0000, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0000, -0.0709, -0.1935]]],
[[[-0.0000, 0.0000, -0.0000, -0.0000, 0.0000],
[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000],
[ 0.0000, 0.0000, 0.0000, -0.0000, -0.0000],
[-0.0000, 0.0000, -0.0000, -0.0000, -0.0000],
[ 0.0000, -0.0000, -0.0000, -0.0000, -0.0000]]],
[[[-0.1167, -0.0000, -0.1579, 0.0000, -0.0397],
[ 0.0000, 0.0623, -0.1694, 0.1384, -0.0550],
[-0.0767, -0.1660, -0.0000, 0.0572, -0.0000],
[ 0.0779, -0.1641, 0.1485, -0.1468, -0.0345],
[ 0.0418, 0.0000, 0.0000, 0.1822, -0.1586]]]], device='cuda:0',
grad_fn=<MulBackward0>)
移除重新参数化后：
prune.remove(module, 'weight')
print(list(module.named_parameters()))
[('bias_orig', Parameter containing:
tensor([ 0.0503, -0.0860, -0.0219, -0.1497, 0.1822, -0.1468], 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.1980, -0.0000, -0.0000, 0.0000, 0.0576],
[ 0.0828, 0.0000, -0.0035, 0.1565, -0.0000],
[ 0.0126, -0.1365, 0.0617, -0.0689, 0.0613],
[-0.0417, 0.1659, -0.1185, -0.1193, -0.1193],
[ 0.1799, 0.0000, 0.1925, -0.1651, -0.0000]]],
[[[-0.1565, -0.1345, 0.0810, 0.0716, 0.1662],
[-0.0000, -0.1363, 0.1061, -0.0808, 0.0000],
[-0.0475, 0.1144, -0.1554, -0.0000, 0.0610],
[ 0.0000, -0.0000, 0.1192, 0.1360, -0.1450],
[-0.1068, 0.1831, -0.0000, -0.0709, -0.1935]]],
[[[-0.0000, 0.0000, -0.0000, -0.0000, 0.0000],
[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000],
[ 0.0000, 0.0000, 0.0000, -0.0000, -0.0000],
[-0.0000, 0.0000, -0.0000, -0.0000, -0.0000],
[ 0.0000, -0.0000, -0.0000, -0.0000, -0.0000]]],
[[[-0.1167, -0.0000, -0.1579, 0.0000, -0.0397],
[ 0.0000, 0.0623, -0.1694, 0.1384, -0.0550],
[-0.0767, -0.1660, -0.0000, 0.0572, -0.0000],
[ 0.0779, -0.1641, 0.1485, -0.1468, -0.0345],
[ 0.0418, 0.0000, 0.0000, 0.1822, -0.1586]]]], device='cuda:0',
requires_grad=True))]
print(list(module.named_buffers()))
[('bias_mask', tensor([0., 0., 0., 1., 1., 1.], 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%连接。
这可能会导致每层的剪枝比例不同。让我们看看如何使用global_unstructured从
torch.nn.utils.prune来实现这一点。
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.92%
Sparsity in fc1.weight: 22.16%
Sparsity in fc2.weight: 12.10%
Sparsity in fc3.weight: 11.31%
Global sparsity: 20.00%
使用自定义剪枝函数扩展torch.nn.utils.prune
To implement your own pruning function, you can extend the
nn.utils.prune module by subclassing the BasePruningMethod
base class, the same way all other pruning methods do. The base class
implements the following methods for you: __call__, apply_mask,
apply, prune, and remove. Beyond some special cases, you shouldn’t
have to reimplement these methods for your new pruning technique.
You will, however, have to implement __init__ (the constructor),
and compute_mask (the instructions on how to compute the mask
for the given tensor according to the logic of your pruning
technique). In addition, you will have to specify which type of
pruning this technique implements (supported options are global,
structured, and unstructured). This is needed to determine
how to combine masks in the case in which pruning is applied
iteratively. In other words, when pruning a prepruned parameter,
the current pruning technique is expected to act on the unpruned
portion of the parameter. Specifying the PRUNING_TYPE will
enable the PruningContainer (which handles the iterative
application of pruning masks) to correctly identify the slice of the
parameter to prune.
例如，假设您想要实现一种修剪技术，该技术会修剪张量中的每隔一个条目（或者——如果张量之前已经被修剪过——在张量的剩余未修剪部分中）。这将是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.281 秒)
