快捷方式

修剪教程

作者: 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 层。它将包含两个参数 weightbias,以及目前没有缓冲区。

[('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 的参数,它现在只是一个属性。

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 0x7feca8fe3be0>)])

为了完整起见,我们现在也可以修剪 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_maskbias_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 0x7feca8fe3be0>), (1, <torch.nn.utils.prune.L1Unstructured object at 0x7feca8fe3430>)])

迭代修剪

模块中的同一个参数可以被多次修剪,各种修剪调用的效果等于按顺序应用的各种掩码的组合。PruningContainercompute_mask 方法处理新的掩码与旧掩码的组合。

例如,假设我们现在要进一步修剪 module.weight,这次使用结构化修剪沿着张量的第 0 轴(第 0 轴对应于卷积层的输出通道,对于 conv1 来说维度是 6),基于通道的 L2 范数。这可以使用 ln_structured 函数来实现,其中 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],
          [ 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 参数的修剪历史。

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 0x7feca8fe3be0>, <torch.nn.utils.prune.LnStructured object at 0x7feca8fe3340>]

序列化修剪模型

所有相关的张量,包括掩码缓冲区和用于计算修剪张量的原始参数,都存储在模型的 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 功能。请注意,这不会撤消修剪,就像它从未发生过一样。相反,它只是通过将参数 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'))]
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% 的连接。这很可能导致每层不同的修剪百分比。让我们看看如何使用 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.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

要实现自己的修剪函数,您可以通过继承 nn.utils.prune 模块中的 BasePruningMethod 基类来扩展它,就像所有其他修剪方法一样。基类为您实现了以下方法:__call__apply_maskapplypruneremove。除了某些特殊情况之外,您不应该为新的修剪技术重新实现这些方法。但是,您必须实现 __init__(构造函数)和 compute_mask(关于如何根据修剪技术的逻辑计算给定张量的掩码的说明)。此外,您必须指定此技术实现哪种类型的修剪(支持的选项是 globalstructuredunstructured)。这需要确定如何在迭代应用修剪的情况下组合掩码。换句话说,当修剪预修剪参数时,当前修剪技术应该作用于参数的未修剪部分。指定 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.228 秒)

Sphinx-Gallery 生成的图库

文档

访问 PyTorch 的综合开发者文档

查看文档

教程

获取适合初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得答案

查看资源