pytorch prune剪支教程

剪枝教程

作者: 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

c:\Users\tangzeyu\AppData\Local\Programs\Python\Python310\lib\site-packages\tqdm\auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
创建模型
在本教程中，我们使用了来自LeCun等人的LeNet架构，该架构于1998年提出。

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, 3x3 square conv kernel
        self.conv1 = nn.Conv2d(1, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5x5 image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, int(x.nelement() / x.shape[0]))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = LeNet().to(device=device)

检查模块
让我们检查我们的LeNet模型中（未经剪枝的）conv1层。它将包含两个参数weight和bias，目前还没有缓冲区。 在PyTorch中，模型的缓冲区（buffers）是一种特殊的张量，用于存储模型的非参数状态。与模型的参数不同，缓冲区的值不会被更新，而是用于存储模型中的其他状态信息，例如运行统计数据、固定的权重等。在剪枝过程中，我们通常不需要修改或操作缓冲区。

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

[('weight', Parameter containing:
tensor([[[[-0.1755, -0.1171, -0.1945],
      [ 0.0642,  0.0352, -0.0277],
      [-0.0194,  0.2641,  0.2108]]],


    [[[ 0.2196, -0.1947, -0.1448],
      [-0.2786,  0.0057, -0.0353],
      [ 0.3258, -0.2772,  0.2517]]],


    [[[-0.0011, -0.2958,  0.1050],
      [ 0.3092, -0.2679,  0.2803],
      [-0.2357, -0.1246, -0.1471]]],


    [[[-0.1714,  0.0742, -0.1299],
      [ 0.1917, -0.1561,  0.2330],
      [ 0.1270,  0.1263,  0.2720]]],


    [[[ 0.2880,  0.2715,  0.3185],
      [-0.2470, -0.0679,  0.2465],
      [-0.2129, -0.0406, -0.1415]]],


    [[[ 0.3107, -0.1880, -0.0499],
      [ 0.0915, -0.2996, -0.2588],
      [-0.0128, -0.1073,  0.1996]]]], requires_grad=True)), ('bias', Parameter containing:
tensor([ 0.2663,  0.0798, -0.1957,  0.2318,  0.2336, -0.1380],
   requires_grad=True))]

print(list(module.named_buffers()))

[]

剪枝一个模块
要剪枝一个模块（在