修剪教程

原文： https://pytorch.org/tutorials/intermediate/pruning_tutorial.html

注意

单击此处的下载完整的示例代码

作者： Michela Paganini

最新的深度学习技术依赖于难以部署的过度参数化模型。 相反，已知生物神经网络使用有效的稀疏连通性。 为了减少内存，电池和硬件消耗，同时又不牺牲精度，在设备上部署轻量级模型并通过私有设备内计算来确保私密性，确定通过减少模型中参数数量来压缩模型的最佳技术很重要。 在研究方面，修剪用于研究参数过度配置和参数不足网络之间学习动态的差异，以研究幸运稀疏子网络和初始化(“ 彩票”）作为破坏性对象的作用。 神经结构搜索技术等等。

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

要求

"torch>=1.4.0a0+8e8a5e0"

  1. import torch
  2. from torch import nn
  3. import torch.nn.utils.prune as prune
  4. import torch.nn.functional as F

建立模型

在本教程中，我们使用 LeCun 等人，1998 年的 LeNet 体系结构。

  1. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  3. class LeNet(nn.Module):
  4.     def __init__(self):
  5.         super(LeNet, self).__init__()
  6.         # 1 input image channel, 6 output channels, 3x3 square conv kernel
  7.         self.conv1 = nn.Conv2d(1, 6, 3)
  8.         self.conv2 = nn.Conv2d(6, 16, 3)
  9.         self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5x5 image dimension
  10.         self.fc2 = nn.Linear(120, 84)
  11.         self.fc3 = nn.Linear(84, 10)
  13.     def forward(self, x):
  14.         x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
  15.         x = F.max_pool2d(F.relu(self.conv2(x)), 2)
  16.         x = x.view(-1, int(x.nelement() / x.shape[0]))
  17.         x = F.relu(self.fc1(x))
  18.         x = F.relu(self.fc2(x))
  19.         x = self.fc3(x)
  20.         return x
  22. model = LeNet().to(device=device)

检查模块

让我们检查一下 LeNet 模型中的(未经修剪的）conv1层。 目前它将包含两个参数weight和bias，并且没有缓冲区。

  1. module = model.conv1
  2. print(list(module.named_parameters()))

出：

  1. [('weight', Parameter containing:
  2. tensor([[[[ 0.3161, -0.2212,  0.0417],
  3.           [ 0.2488,  0.2415,  0.2071],
  4.           [-0.2412, -0.2400, -0.2016]]],
  6.         [[[ 0.0419,  0.3322, -0.2106],
  7.           [ 0.1776, -0.1845, -0.3134],
  8.           [-0.0708,  0.1921,  0.3095]]],
  10.         [[[-0.2070,  0.0723,  0.2876],
  11.           [ 0.2209,  0.2077,  0.2369],
  12.           [ 0.2108,  0.0861, -0.2279]]],
  14.         [[[-0.2799, -0.1527, -0.0388],
  15.           [-0.2043,  0.1220,  0.1032],
  16.           [-0.0755,  0.1281,  0.1077]]],
  18.         [[[ 0.2035,  0.2245, -0.1129],
  19.           [ 0.3257, -0.0385, -0.0115],
  20.           [-0.3146, -0.2145, -0.1947]]],
  22.         [[[-0.1426,  0.2370, -0.1089],
  23.           [-0.2491,  0.1282,  0.1067],
  24.           [ 0.2159, -0.1725,  0.0723]]]], device='cuda:0', requires_grad=True)), ('bias', Parameter containing:
  25. tensor([-0.1214, -0.0749, -0.2656, -0.1519, -0.1021,  0.1425], device='cuda:0',
  26.        requires_grad=True))]
  1. print(list(module.named_buffers()))

Out:

  1. []

修剪模块

要修剪模块(在此示例中，为 LeNet 架构的conv1层），请首先从torch.nn.utils.prune中可用的那些技术中选择一种修剪技术(或通过子类化BasePruningMethod实现您自己的）。 然后，指定模块和该模块中要修剪的参数的名称。 最后，使用所选修剪技术所需的适当关键字参数，指定修剪参数。

在此示例中，我们将在conv1层中名为weight的参数中随机修剪 30％的连接。 模块作为第一个参数传递给函数； name使用其字符串标识符在该模块内标识参数； amount表示与修剪的连接百分比(如果它是介于 0 和 1 之间的浮点数），或者表示与修剪的连接的绝对数量(如果它是非负整数）。

  1. prune.random_unstructured(module, name="weight", amount=0.3)

修剪是通过从参数中删除weight并将其替换为名为weight_orig的新参数(即，将"_orig"附加到初始参数name）来进行的。 weight_orig存储未修剪的张量版本。 bias未修剪，因此它将保持完整。

  1. print(list(module.named_parameters()))

Out:

  1. [('bias', Parameter containing:
  2. tensor([-0.1214, -0.0749, -0.2656, -0.1519, -0.1021,  0.1425], device='cuda:0',
  3.        requires_grad=True)), ('weight_orig', Parameter containing:
  4. tensor([[[[ 0.3161, -0.2212,  0.0417],
  5.           [ 0.2488,  0.2415,  0.2071],
  6.           [-0.2412, -0.2400, -0.2016]]],
  8.         [[[ 0.0419,  0.3322, -0.2106],
  9.           [ 0.1776, -0.1845, -0.3134],
  10.           [-0.0708,  0.1921,  0.3095]]],
  12.         [[[-0.2070,  0.0723,  0.2876],
  13.           [ 0.2209,  0.2077,  0.2369],
  14.           [ 0.2108,  0.0861, -0.2279]]],
  16.         [[[-0.2799, -0.1527, -0.0388],
  17.           [-0.2043,  0.1220,  0.1032],
  18.           [-0.0755,  0.1281,  0.1077]]],
  20.         [[[ 0.2035,  0.2245, -0.1129],
  21.           [ 0.3257, -0.0385, -0.0115],
  22.           [-0.3146, -0.2145, -0.1947]]],
  24.         [[[-0.1426,  0.2370, -0.1089],
  25.           [-0.2491,  0.1282,  0.1067],
  26.           [ 0.2159, -0.1725,  0.0723]]]], device='cuda:0', requires_grad=True))]

通过以上选择的修剪技术生成的修剪掩码将保存为名为weight_mask的模块缓冲区(即，将"_mask"附加到初始参数name）。

  1. print(list(module.named_buffers()))

Out:

  1. [('weight_mask', tensor([[[[0., 1., 0.],
  2.           [1., 0., 0.],
  3.           [1., 1., 1.]]],
  5.         [[[1., 0., 1.],
  6.           [1., 1., 0.],
  7.           [1., 0., 1.]]],
  9.         [[[1., 0., 0.],
  10.           [0., 1., 1.],
  11.           [1., 1., 1.]]],
  13.         [[[1., 0., 0.],
  14.           [1., 1., 1.],
  15.           [1., 1., 1.]]],
  17.         [[[1., 0., 1.],
  18.           [1., 1., 1.],
  19.           [0., 1., 1.]]],
  21.         [[[1., 1., 1.],
  22.           [1., 1., 0.],
  23.           [1., 1., 0.]]]], device='cuda:0'))]

为了使前向通过不更改即可工作，需要存在weight属性。 torch.nn.utils.prune中实现的修剪技术计算权重的修剪版本(通过将掩码与原始参数组合）并将其存储在属性weight中。 注意，这不再是module的参数，现在只是一个属性。

  1. print(module.weight)

Out:

  1. tensor([[[[ 0.0000, -0.2212,  0.0000],
  2.           [ 0.2488,  0.0000,  0.0000],
  3.           [-0.2412, -0.2400, -0.2016]]],
  5.         [[[ 0.0419,  0.0000, -0.2106],
  6.           [ 0.1776, -0.1845, -0.0000],
  7.           [-0.0708,  0.0000,  0.3095]]],
  9.         [[[-0.2070,  0.0000,  0.0000],
  10.           [ 0.0000,  0.2077,  0.2369],
  11.           [ 0.2108,  0.0861, -0.2279]]],
  13.         [[[-0.2799, -0.0000, -0.0000],
  14.           [-0.2043,  0.1220,  0.1032],
  15.           [-0.0755,  0.1281,  0.1077]]],
  17.         [[[ 0.2035,  0.0000, -0.1129],
  18.           [ 0.3257, -0.0385, -0.0115],
  19.           [-0.0000, -0.2145, -0.1947]]],
  21.         [[[-0.1426,  0.2370, -0.1089],
  22.           [-0.2491,  0.1282,  0.0000],
  23.           [ 0.2159, -0.1725,  0.0000]]]], device='cuda:0',
  24.        grad_fn=<MulBackward0>)

最后，使用 PyTorch 的forward_pre_hooks在每次向前传递之前应用修剪。 具体来说，当修剪module时(如我们在此处所做的那样），它将为与之关联的每个参数获取forward_pre_hook进行修剪。 在这种情况下，由于到目前为止我们只修剪了名称为weight的原始参数，因此只会出现一个钩子。

  1. print(module._forward_pre_hooks)

Out:

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

为了完整起见，我们现在也可以修剪bias，以查看module的参数，缓冲区，挂钩和属性如何变化。 仅出于尝试另一种修剪技术的目的，在此我们按 L1 范数修剪偏差中的 3 个最小条目，如l1_unstructured修剪功能中所实现的。

  1. prune.l1_unstructured(module, name="bias", amount=3)

现在，我们希望命名的参数同时包含weight_orig(从前）和bias_orig。 缓冲区将包括weight_mask和bias_mask。 两个张量的修剪版本将作为模块属性存在，并且该模块现在将具有两个forward_pre_hooks。

  1. print(list(module.named_parameters()))

Out:

  1. [('weight_orig', Parameter containing:
  2. tensor([[[[ 0.3161, -0.2212,  0.0417],
  3.           [ 0.2488,  0.2415,  0.2071],
  4.           [-0.2412, -0.2400, -0.2016]]],
  6.         [[[ 0.0419,  0.3322, -0.2106],
  7.           [ 0.1776, -0.1845, -0.3134],
  8.           [-0.0708,  0.1921,  0.3095]]],
  10.         [[[-0.2070,  0.0723,  0.2876],
  11.           [ 0.2209,  0.2077,  0.2369],
  12.           [ 0.2108,  0.0861, -0.2279]]],
  14.         [[[-0.2799, -0.1527, -0.0388],
  15.           [-0.2043,  0.1220,  0.1032],
  16.           [-0.0755,  0.1281,  0.1077]]],
  18.         [[[ 0.2035,  0.2245, -0.1129],
  19.           [ 0.3257, -0.0385, -0.0115],
  20.           [-0.3146, -0.2145, -0.1947]]],
  22.         [[[-0.1426,  0.2370, -0.1089],
  23.           [-0.2491,  0.1282,  0.1067],
  24.           [ 0.2159, -0.1725,  0.0723]]]], device='cuda:0', requires_grad=True)), ('bias_orig', Parameter containing:
  25. tensor([-0.1214, -0.0749, -0.2656, -0.1519, -0.1021,  0.1425], device='cuda:0',
  26.        requires_grad=True))]
  1. print(list(module.named_buffers()))

Out:

  1. [('weight_mask', tensor([[[[0., 1., 0.],
  2.           [1., 0., 0.],
  3.           [1., 1., 1.]]],
  5.         [[[1., 0., 1.],
  6.           [1., 1., 0.],
  7.           [1., 0., 1.]]],
  9.         [[[1., 0., 0.],
  10.           [0., 1., 1.],
  11.           [1., 1., 1.]]],
  13.         [[[1., 0., 0.],
  14.           [1., 1., 1.],
  15.           [1., 1., 1.]]],
  17.         [[[1., 0., 1.],
  18.           [1., 1., 1.],
  19.           [0., 1., 1.]]],
  21.         [[[1., 1., 1.],
  22.           [1., 1., 0.],
  23.           [1., 1., 0.]]]], device='cuda:0')), ('bias_mask', tensor([0., 0., 1., 1., 0., 1.], device='cuda:0'))]
  1. print(module.bias)

Out:

  1. tensor([-0.0000, -0.0000, -0.2656, -0.1519, -0.0000,  0.1425], device='cuda:0',
  2.        grad_fn=<MulBackward0>)
  1. print(module._forward_pre_hooks)

Out:

  1. OrderedDict([(0, <torch.nn.utils.prune.RandomUnstructured object at 0x7f1e6c425400>), (1, <torch.nn.utils.prune.L1Unstructured object at 0x7f1e6c425550>)])

迭代修剪

一个模块中的同一参数可以被多次修剪，各种修剪调用的效果等于串联应用的各种蒙版的组合。 PruningContainer的compute_mask方法可处理新遮罩与旧遮罩的组合。

例如，假设我们现在要进一步修剪module.weight，这一次是使用沿着张量的第 0 轴的结构化修剪(第 0 轴对应于卷积层的输出通道，并且conv1的维数为 6） ，基于渠道的 L2 规范。 这可以通过ln_structured和n=2和dim=0功能来实现。

  1. prune.ln_structured(module, name="weight", amount=0.5, n=2, dim=0)
  3. # As we can verify, this will zero out all the connections corresponding to
  4. # 50% (3 out of 6) of the channels, while preserving the action of the
  5. # previous mask.
  6. print(module.weight)

Out:

  1. tensor([[[[ 0.0000, -0.2212,  0.0000],
  2.           [ 0.2488,  0.0000,  0.0000],
  3.           [-0.2412, -0.2400, -0.2016]]],
  5.         [[[ 0.0000,  0.0000, -0.0000],
  6.           [ 0.0000, -0.0000, -0.0000],
  7.           [-0.0000,  0.0000,  0.0000]]],
  9.         [[[-0.2070,  0.0000,  0.0000],
  10.           [ 0.0000,  0.2077,  0.2369],
  11.           [ 0.2108,  0.0861, -0.2279]]],
  13.         [[[-0.0000, -0.0000, -0.0000],
  14.           [-0.0000,  0.0000,  0.0000],
  15.           [-0.0000,  0.0000,  0.0000]]],
  17.         [[[ 0.2035,  0.0000, -0.1129],
  18.           [ 0.3257, -0.0385, -0.0115],
  19.           [-0.0000, -0.2145, -0.1947]]],
  21.         [[[-0.0000,  0.0000, -0.0000],
  22.           [-0.0000,  0.0000,  0.0000],
  23.           [ 0.0000, -0.0000,  0.0000]]]], device='cuda:0',
  24.        grad_fn=<MulBackward0>)

现在，对应的钩子将为torch.nn.utils.prune.PruningContainer类型，并将存储应用于weight参数的修剪历史。

  1. for hook in module._forward_pre_hooks.values():
  2.     if hook._tensor_name == "weight":  # select out the correct hook
  3.         break
  5. print(list(hook))  # pruning history in the container

Out:

  1. [<torch.nn.utils.prune.RandomUnstructured object at 0x7f1e6c425400>, <torch.nn.utils.prune.LnStructured object at 0x7f1e6c4259b0>]

序列化修剪的模型

所有相关的张量，包括掩码缓冲区和用于计算修剪的张量的原始参数，都存储在模型的state_dict中，因此可以根据需要轻松地序列化和保存。

  1. print(model.state_dict().keys())

Out:

  1. 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重新分配为模型参数(修剪后的版本）来使其永久不变。

删除重新参数化之前：

  1. print(list(module.named_parameters()))

Out:

  1. [('weight_orig', Parameter containing:
  2. tensor([[[[ 0.3161, -0.2212,  0.0417],
  3.           [ 0.2488,  0.2415,  0.2071],
  4.           [-0.2412, -0.2400, -0.2016]]],
  6.         [[[ 0.0419,  0.3322, -0.2106],
  7.           [ 0.1776, -0.1845, -0.3134],
  8.           [-0.0708,  0.1921,  0.3095]]],
  10.         [[[-0.2070,  0.0723,  0.2876],
  11.           [ 0.2209,  0.2077,  0.2369],
  12.           [ 0.2108,  0.0861, -0.2279]]],
  14.         [[[-0.2799, -0.1527, -0.0388],
  15.           [-0.2043,  0.1220,  0.1032],
  16.           [-0.0755,  0.1281,  0.1077]]],
  18.         [[[ 0.2035,  0.2245, -0.1129],
  19.           [ 0.3257, -0.0385, -0.0115],
  20.           [-0.3146, -0.2145, -0.1947]]],
  22.         [[[-0.1426,  0.2370, -0.1089],
  23.           [-0.2491,  0.1282,  0.1067],
  24.           [ 0.2159, -0.1725,  0.0723]]]], device='cuda:0', requires_grad=True)), ('bias_orig', Parameter containing:
  25. tensor([-0.1214, -0.0749, -0.2656, -0.1519, -0.1021,  0.1425], device='cuda:0',
  26.        requires_grad=True))]
  1. print(list(module.named_buffers()))

Out:

  1. [('weight_mask', tensor([[[[0., 1., 0.],
  2.           [1., 0., 0.],
  3.           [1., 1., 1.]]],
  5.         [[[0., 0., 0.],
  6.           [0., 0., 0.],
  7.           [0., 0., 0.]]],
  9.         [[[1., 0., 0.],
  10.           [0., 1., 1.],
  11.           [1., 1., 1.]]],
  13.         [[[0., 0., 0.],
  14.           [0., 0., 0.],
  15.           [0., 0., 0.]]],
  17.         [[[1., 0., 1.],
  18.           [1., 1., 1.],
  19.           [0., 1., 1.]]],
  21.         [[[0., 0., 0.],
  22.           [0., 0., 0.],
  23.           [0., 0., 0.]]]], device='cuda:0')), ('bias_mask', tensor([0., 0., 1., 1., 0., 1.], device='cuda:0'))]
  1. print(module.weight)

Out:

  1. tensor([[[[ 0.0000, -0.2212,  0.0000],
  2.           [ 0.2488,  0.0000,  0.0000],
  3.           [-0.2412, -0.2400, -0.2016]]],
  5.         [[[ 0.0000,  0.0000, -0.0000],
  6.           [ 0.0000, -0.0000, -0.0000],
  7.           [-0.0000,  0.0000,  0.0000]]],
  9.         [[[-0.2070,  0.0000,  0.0000],
  10.           [ 0.0000,  0.2077,  0.2369],
  11.           [ 0.2108,  0.0861, -0.2279]]],
  13.         [[[-0.0000, -0.0000, -0.0000],
  14.           [-0.0000,  0.0000,  0.0000],
  15.           [-0.0000,  0.0000,  0.0000]]],
  17.         [[[ 0.2035,  0.0000, -0.1129],
  18.           [ 0.3257, -0.0385, -0.0115],
  19.           [-0.0000, -0.2145, -0.1947]]],
  21.         [[[-0.0000,  0.0000, -0.0000],
  22.           [-0.0000,  0.0000,  0.0000],
  23.           [ 0.0000, -0.0000,  0.0000]]]], device='cuda:0',
  24.        grad_fn=<MulBackward0>)

删除重新参数化后：

  1. prune.remove(module, 'weight')
  2. print(list(module.named_parameters()))

Out:

  1. [('bias_orig', Parameter containing:
  2. tensor([-0.1214, -0.0749, -0.2656, -0.1519, -0.1021,  0.1425], device='cuda:0',
  3.        requires_grad=True)), ('weight', Parameter containing:
  4. tensor([[[[ 0.0000, -0.2212,  0.0000],
  5.           [ 0.2488,  0.0000,  0.0000],
  6.           [-0.2412, -0.2400, -0.2016]]],
  8.         [[[ 0.0000,  0.0000, -0.0000],
  9.           [ 0.0000, -0.0000, -0.0000],
  10.           [-0.0000,  0.0000,  0.0000]]],
  12.         [[[-0.2070,  0.0000,  0.0000],
  13.           [ 0.0000,  0.2077,  0.2369],
  14.           [ 0.2108,  0.0861, -0.2279]]],
  16.         [[[-0.0000, -0.0000, -0.0000],
  17.           [-0.0000,  0.0000,  0.0000],
  18.           [-0.0000,  0.0000,  0.0000]]],
  20.         [[[ 0.2035,  0.0000, -0.1129],
  21.           [ 0.3257, -0.0385, -0.0115],
  22.           [-0.0000, -0.2145, -0.1947]]],
  24.         [[[-0.0000,  0.0000, -0.0000],
  25.           [-0.0000,  0.0000,  0.0000],
  26.           [ 0.0000, -0.0000,  0.0000]]]], device='cuda:0', requires_grad=True))]
  1. print(list(module.named_buffers()))

Out:

  1. [('bias_mask', tensor([0., 0., 1., 1., 0., 1.], device='cuda:0'))]

修剪模型中的多个参数

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

  1. new_model = LeNet()
  2. for name, module in new_model.named_modules():
  3.     # prune 20% of connections in all 2D-conv layers
  4.     if isinstance(module, torch.nn.Conv2d):
  5.         prune.l1_unstructured(module, name='weight', amount=0.2)
  6.     # prune 40% of connections in all linear layers
  7.     elif isinstance(module, torch.nn.Linear):
  8.         prune.l1_unstructured(module, name='weight', amount=0.4)
  10. print(dict(new_model.named_buffers()).keys())  # to verify that all masks exist

Out:

  1. 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进行操作。

  1. model = LeNet()
  3. parameters_to_prune = (
  4.     (model.conv1, 'weight'),
  5.     (model.conv2, 'weight'),
  6.     (model.fc1, 'weight'),
  7.     (model.fc2, 'weight'),
  8.     (model.fc3, 'weight'),
  9. )
  11. prune.global_unstructured(
  12.     parameters_to_prune,
  13.     pruning_method=prune.L1Unstructured,
  14.     amount=0.2,
  15. )

现在，我们可以检查在每个修剪参数中引起的稀疏性，该稀疏性将不等于每层中的 20％。 但是，全球稀疏度将(大约）为 20％。

  1. print(
  2.     "Sparsity in conv1.weight: {:.2f}%".format(
  3.         100\. * float(torch.sum(model.conv1.weight == 0))
  4.         / float(model.conv1.weight.nelement())
  5.     )
  6. )
  7. print(
  8.     "Sparsity in conv2.weight: {:.2f}%".format(
  9.         100\. * float(torch.sum(model.conv2.weight == 0))
  10.         / float(model.conv2.weight.nelement())
  11.     )
  12. )
  13. print(
  14.     "Sparsity in fc1.weight: {:.2f}%".format(
  15.         100\. * float(torch.sum(model.fc1.weight == 0))
  16.         / float(model.fc1.weight.nelement())
  17.     )
  18. )
  19. print(
  20.     "Sparsity in fc2.weight: {:.2f}%".format(
  21.         100\. * float(torch.sum(model.fc2.weight == 0))
  22.         / float(model.fc2.weight.nelement())
  23.     )
  24. )
  25. print(
  26.     "Sparsity in fc3.weight: {:.2f}%".format(
  27.         100\. * float(torch.sum(model.fc3.weight == 0))
  28.         / float(model.fc3.weight.nelement())
  29.     )
  30. )
  31. print(
  32.     "Global sparsity: {:.2f}%".format(
  33.         100\. * float(
  34.             torch.sum(model.conv1.weight == 0)
  35.             + torch.sum(model.conv2.weight == 0)
  36.             + torch.sum(model.fc1.weight == 0)
  37.             + torch.sum(model.fc2.weight == 0)
  38.             + torch.sum(model.fc3.weight == 0)
  39.         )
  40.         / float(
  41.             model.conv1.weight.nelement()
  42.             + model.conv2.weight.nelement()
  43.             + model.fc1.weight.nelement()
  44.             + model.fc2.weight.nelement()
  45.             + model.fc3.weight.nelement()
  46.         )
  47.     )
  48. )

Out:

  1. Sparsity in conv1.weight: 7.41%
  2. Sparsity in conv2.weight: 9.49%
  3. Sparsity in fc1.weight: 22.00%
  4. Sparsity in fc2.weight: 12.28%
  5. Sparsity in fc3.weight: 9.76%
  6. Global sparsity: 20.00%

使用自定义修剪功能扩展torch.nn.utils.prune

要实现自己的修剪功能，您可以通过继承BasePruningMethod基类来扩展nn.utils.prune模块，这与所有其他修剪方法一样。 基类为您实现以下方法：__call__，apply_mask，apply，prune和remove。 除了某些特殊情况外，您不必为新的修剪技术重新实现这些方法。 但是，您将必须实现__init__(构造函数）和compute_mask(有关如何根据修剪技术的逻辑为给定张量计算掩码的说明）。 另外，您将必须指定此技术实现的修剪类型(支持的选项为global，structured和unstructured）。 需要确定在迭代应用修剪的情况下如何组合蒙版。 换句话说，当修剪预修剪的参数时，当前的修剪技术应作用于参数的未修剪部分。 指定PRUNING_TYPE将使PruningContainer(处理修剪蒙版的迭代应用）正确识别要修剪的参数。

例如，假设您要实施一种修剪技术，以修剪张量中的所有其他条目(或者-如果先前已修剪过张量，则在张量的其余未修剪部分中）。 这将是PRUNING_TYPE='unstructured'，因为它作用于层中的单个连接，而不作用于整个单元/通道('structured'），或作用于不同的参数('global'）。

  1. class FooBarPruningMethod(prune.BasePruningMethod):
  2.     """Prune every other entry in a tensor
  3.     """
  4.     PRUNING_TYPE = 'unstructured'
  6.     def compute_mask(self, t, default_mask):
  7.         mask = default_mask.clone()
  8.         mask.view(-1)[::2] = 0
  9.         return mask

现在，要将其应用于nn.Module中的参数，还应该提供一个简单的函数来实例化该方法并将其应用。

  1. def foobar_unstructured(module, name):
  2.     """Prunes tensor corresponding to parameter called `name` in `module`
  3.     by removing every other entry in the tensors.
  4.     Modifies module in place (and also return the modified module)
  5.     by:
  6.     1) adding a named buffer called `name+'_mask'` corresponding to the
  7.     binary mask applied to the parameter `name` by the pruning method.
  8.     The parameter `name` is replaced by its pruned version, while the
  9.     original (unpruned) parameter is stored in a new parameter named
  10.     `name+'_orig'`.
  12.     Args:
  13.         module (nn.Module): module containing the tensor to prune
  14.         name (string): parameter name within `module` on which pruning
  15.                 will act.
  17.     Returns:
  18.         module (nn.Module): modified (i.e. pruned) version of the input
  19.             module
  21.     Examples:
  22.         >>> m = nn.Linear(3, 4)
  23.         >>> foobar_unstructured(m, name='bias')
  24.     """
  25.     FooBarPruningMethod.apply(module, name)
  26.     return module

试试吧！

  1. model = LeNet()
  2. foobar_unstructured(model.fc3, name='bias')
  4. print(model.fc3.bias_mask)

Out:

  1. tensor([0., 1., 0., 1., 0., 1., 0., 1., 0., 1.])

脚本的总运行时间：(0 分钟 0.146 秒）

Download Python source code: pruning_tutorial.py Download Jupyter notebook: pruning_tutorial.ipynb

由狮身人面像画廊生成的画廊