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 层。它将包含两个 参数 weightbias ,并且暂时没有缓冲区。

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_maskbias_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=2dim=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_origweight_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.pruneglobal_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 秒)



回到顶部