Skip to content

修剪教程

译者：片刻小哥哥

项目地址：https://pytorch.apachecn.org/2.0/tutorials/intermediate/pruning_tutorial

原始地址：https://pytorch.org/tutorials/intermediate/pruning_tutorial.html

作者 : 米歇拉·帕格尼尼

最先进的深度学习技术依赖于难以部署的过度参数化模型。相反，生物神经网络已知使用有效的稀疏连接。为了在不牺牲准确性的情况下减少内存、电池和硬件消耗， 通过减少模型中参数的数量来识别 最佳压缩模型的技术非常重要。这反过来又允许您在设备上部署轻量级模型，并通过私有设备上计算来保证 隐私。在研究前沿，剪枝被 用来研究参数化网络和参数化不足网络之间学习动态的差异，研究幸运 s稀疏子网络和初始化的作用 (“ 彩票 ”) 作为一种破坏性的 神经架构搜索技术，等等。

在本教程中，您将学习如何使用 torch.nn.utils.prune 来解析您的神经网络，以及如何扩展它以实现 您自己的自定义修剪技术。

要求

"火炬>=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 架构。\ n

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)

检查模块

让’s 检查 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]，\ n [ 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],\ n [-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% 连接。 该模块作为第一个参数传递给function; 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 的新参数(即附加 "_orig"\ n 到 初始参数 name )。 weight_orig 存储tensor的未剪枝版本。 偏差 未被修剪，因此它将保持不变。

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]，\ n [ 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],\ n [-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 的模块缓冲区(即附加 "_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]，\ n [ 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],\ n [-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’s forward_pre_hooks 进行修剪。具体来说，当“模块”被修剪时，正如我们在此处所做的那样，它将为与其相关的每个被修剪的参数获取一个“forward_pre_hook”。在这种情况下，由于我们到目前为止 只修剪了名为 weight 的原始参数，因此 仅存在一个钩子。

print(module._forward_pre_hooks)

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

为了完整起见，我们现在也可以修剪 bias ，以查看 module 的 参数、缓冲区、挂钩和属性如何变化。 只是为了方便起见为了尝试另一种修剪技术，这里我们按照 L1 范数修剪 偏差中的 3 个最小条目，如 l1_unstructed 修剪函数中实现的那样。

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 。两个tensor的修剪版本将作为 模块属性存在，并且该模块现在将有两个 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]，\ n [ 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],\ n [-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 0x7f4d3635f700>), (1, <torch.nn.utils.prune.L1Unstructured object at 0x7f4d3635fa90>)])

迭代修剪

模块中的同一参数可以多次修剪， 各种修剪调用的效果等于 串联应用的各种掩码的组合。 处理新掩码与旧掩码的组合通过 PruningContainer ’s compute_mask 方法。

举例来说，我们现在想要进一步剪枝 module.weight ，这次 沿tensor的第 0 个轴使用结构化剪枝(第 0 个轴 对应于卷积神经网络的输出通道)层，并且基于通道’ L2 范数， conv1 的维数为 6。这可以使用 ln_structed 函数以及 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]，\ n [ 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],\ n [-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 参数的剪枝历史记录。

for hook in module._forward_pre_hooks.values():
    if hook._tensor_name == "weight":  # select out the correct hook
        break

print(list(hook))  # pruning history in the container

[<torch.nn.utils.prune.RandomUnstructured object at 0x7f4d3635f700>, <torch.nn.utils.prune.LnStructured object at 0x7f4d3635c0d0>]

序列化修剪后的模型

所有相关tensor，包括掩码缓冲区和用于计算修剪tensor的原始参数 都存储在模型’s 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 功能。 请注意，这不会’ 撤消修剪，因为如果它从未发生过。它只是 通过将参数 “权重” 重新分配给 模型参数(在其修剪版本中)来使其永久化。

在删除重新参数化之前：

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]，\ n [ 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],\ n [-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]，\ n [ 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],\ n [-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]，\ n [ 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],\ n [-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'))]

修剪模型中的多个参数

通过指定所需的修剪技术和参数，我们可以轻松 修剪网络中的多个tensor，也许根据它们的类型，如我们 将在本示例中看到的那样。

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'])

全局修剪

到目前为止，我们只研究了通常所说的 “local” 剪枝， i.e.通过将每个条目的统计数据(权重大小、激活、梯度等)与该tensor中的其他条目进行专门比较，逐一修剪模型中的tensor。然而，一种常见且可能更强大的技术是一次性修剪模型，例如删除整个模型中最低 20% 的连接，而不是删除每个模型中最低 20% 的连接。 n层。这可能会导致每层的修剪百分比不同。 让’s 看看如何使用 来自 torch.nn.utils.prune的 global_unstructed 来做到这一点 .

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

要实现您自己的修剪功能，您可以通过子类化 BasePruningMethod 基类来扩展 nn.utils.prune 模块，就像所有其他修剪方法一样。基类 为您实现以下方法: __call__ , apply_mask , apply ,\ n 修剪 和 删除 。除了一些特殊情况之外，您应该’t 必须为新的修剪技术重新实现这些方法。 但是，您必须实现 __init__ (构造函数)， 和 compute_mask (有关如何根据剪枝技术的逻辑计算给定tensor的掩码的说明)。此外，您必须指定此技术实现的剪枝类型(支持的选项有“全局”、“结构化”和“非结构化”)。这需要确定 在迭代应用修剪的情况下如何组合掩码。换句话说，当剪枝预剪枝参数时，当前剪枝技术预计作用于参数的未剪枝部分。指定 PRUNING_TYPE 将启用

PruningContainer(处理修剪掩码的迭代应用) 以正确识别 要修剪的参数切片。

让\xe2\x80\x99s 假设，例如，您想要实现一种修剪 技术，修剪tensor中的每个其他条目(或 xe2\x80\x93，如果 tensor之前已被修剪\xe2\tensor剩余未剪枝部分中的 x80\x93)。这将是 PRUNING_TYPE='非结构化' 因为它作用于层中的各个连接，而不是作用于整个 单元/通道( '结构化' )，或跨不同的连接参数 ( '全局' )。

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

让’s 尝试一下！

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.366 秒)

回到顶部