PyTorch 模型剪枝实例教程二、结构化剪枝

目录

        • 1.导包&定义一个简单的网络
        • 2.获取网络需要剪枝的模块
        • 3.模块结构化剪枝(核心)
        • 4.总结

目前大部分最先进的(SOTA)深度学习技术虽然效果好,但由于其模型参数量和计算量过高,难以用于实际部署。而众所周知,生物神经网络使用高效的稀疏连接(生物大脑神经网络balabala啥的都是稀疏连接的),考虑到这一点,为了减少内存、容量和硬件消耗,同时又不牺牲模型预测的精度,在设备上部署轻量级模型,并通过私有的设备上计算以保证隐私,通过减少参数数量来压缩模型的最佳技术非常重要

稀疏神经网络在预测精度方面可以达到密集神经网络的水平,但由于模型参数量小,理论上来讲推理速度也会快很多。而模型剪枝是一种将密集神经网络训练成稀疏神经网络的方法。

本文将通过学习官方示例教程,介绍如何通过一个简单的实例教程来进行模型剪枝,实践深度学习模型压缩加速。

相关链接

深度学习模型压缩与加速技术(一):参数剪枝

PyTorch模型剪枝实例教程一、非结构化剪枝

PyTorch模型剪枝实例教程二、结构化剪枝

PyTorch模型剪枝实例教程三、多参数与全局剪枝

1.导包&定义一个简单的网络

import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

'''搭建类LeNet网络'''
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        # 单通道图像输入,5×5核尺寸
        self.conv1 = nn.Conv2d(1, 3, 5)
        self.conv2 = nn.Conv2d(3, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  
        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

2.获取网络需要剪枝的模块

model = LeNet().to(device=device)
module = model.conv1
print(list(module.named_parameters()))      # 6×5×5的weight + 6×1的bias 的参数量

输出:

[('weight', Parameter containing:
tensor([[[[ 0.1473,  0.1251,  0.0492, -0.1375, -0.0781],
          [ 0.0446, -0.1328,  0.0227,  0.0141, -0.1751],
          [ 0.0253,  0.0313,  0.0391,  0.1607, -0.0716],
          [-0.1125, -0.1641,  0.1691,  0.1583,  0.0449],
          [-0.0094, -0.1916,  0.1701,  0.0704,  0.0407]]],


        [[[-0.1945,  0.0709,  0.1071,  0.0038, -0.0686],
          [ 0.0187,  0.0710, -0.0955, -0.0778,  0.1927],
          [ 0.1643,  0.0791,  0.1235,  0.0241, -0.0021],
          [-0.1124,  0.0246, -0.0349, -0.1561,  0.0178],
          [-0.1779,  0.1216,  0.1086, -0.1837,  0.1789]]],


        [[[-0.0051, -0.1969, -0.0155,  0.1890,  0.1977],
          [-0.0654,  0.1219,  0.0849, -0.1937, -0.0933],
          [-0.0409,  0.1344,  0.1688,  0.1917, -0.1727],
          [ 0.1380, -0.1413, -0.1483, -0.0711, -0.0648],
          [-0.1571,  0.0570,  0.1783, -0.0786,  0.1367]]]], requires_grad=True)), ('bias', Parameter containing:
tensor([ 0.0346, -0.1446,  0.0633], requires_grad=True))]

3.模块结构化剪枝(核心)

剪枝一个模块,需要三步:

  • step1.在torch.nn.utils.prune中选定一个剪枝方案,或者自定义(通过子类BasePruningMethod)
  • step2.指定需要剪枝的模块和对应的名称
  • step3.输入对应函数需要的参数

这里,我们根据通道的L2范数,沿着张量的第0轴(第0轴对应卷积层的输出通道,conv1的维数为3×5×5)对weight参数进行结构化剪枝,使用ln_structured()方法。剪枝比例为33%,dim为0,基于L2范数(n=2)

prune.ln_structured(module, name="weight", amount=0.33, n=2, dim=0)
print(module.weight)

输出:

tensor([[[[ 0.1327,  0.0812, -0.0225, -0.0809,  0.1461],
          [ 0.1335, -0.1709,  0.0575, -0.1608, -0.0677],
          [-0.0397, -0.0982,  0.0654, -0.1030, -0.1656],
          [-0.0570,  0.1940,  0.0085,  0.1896,  0.1979],
          [-0.0673,  0.0910, -0.0177, -0.1748,  0.1667]]],

        [[[-0.0000,  0.0000,  0.0000,  0.0000, -0.0000],
          [-0.0000,  0.0000,  0.0000, -0.0000,  0.0000],
          [ 0.0000, -0.0000,  0.0000,  0.0000, -0.0000],
          [ 0.0000,  0.0000, -0.0000, -0.0000, -0.0000],
          [-0.0000, -0.0000, -0.0000, -0.0000,  0.0000]]],

        [[[ 0.1668,  0.1742, -0.1581, -0.1208,  0.0745],
          [ 0.0459, -0.0275, -0.1190, -0.1631, -0.1956],
          [-0.0480, -0.1716, -0.0168,  0.0089,  0.0876],
          [-0.0129, -0.1616, -0.1164, -0.1869, -0.1782],
          [ 0.0411, -0.0278, -0.1266,  0.1329, -0.1240]]]],
       grad_fn=)

所有相关的张量,包括mask缓冲区和用于计算剪枝张量的原始参数都存储在模型的state_dict中,因此,如果需要,可以很容易地序列化和保存。

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

输出:

odict_keys(['conv1.bias', 'conv1.weight_orig', 'conv1.weight_mask', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])

要使修剪永久存在,可以删除weight_trans(orig)和weight_mask重新参数化,并删除forward_pre_hook,可以使用torch.nn.util.prune中的remove函数。注意,这并没有取消修剪,就像它从未发生过一样。它只是简单地使它永久存在,相反,在它的剪枝版本中,通过将参数的权重重新分配给模型参数。

print(list(module.named_parameters()))
print(list(module.named_buffers()))
print(model.state_dict().keys())

输出:

[('bias', Parameter containing:
tensor([0.1673, 0.0794, 0.0110], requires_grad=True)), ('weight_orig', Parameter containing:
tensor([[[[ 1.2122e-01, -9.2122e-02,  6.4242e-02,  2.0936e-02,  7.7185e-03],
          [ 1.6201e-01, -1.2338e-01, -1.2014e-01,  3.0895e-03, -3.8402e-02],
          [ 7.5407e-03,  1.9274e-01, -3.0035e-02,  1.9638e-02, -5.5985e-03],
          [-1.2915e-01, -7.7561e-02,  6.8224e-02, -1.8743e-01, -1.6051e-01],
          [ 1.4066e-01,  1.1038e-01, -1.8010e-01,  9.4039e-02, -1.2981e-01]]],

        [[[-1.3836e-01, -1.8937e-01,  3.2540e-02, -6.2541e-02,  1.6695e-01],
          [ 1.3803e-01,  1.0196e-01,  8.2551e-02, -1.2815e-06, -1.4024e-02],
          [-3.7121e-02, -1.8625e-01,  4.1115e-02, -1.5329e-01,  3.8362e-02],
          [-5.7373e-02,  9.3459e-02,  5.9365e-02, -9.4975e-02,  1.7842e-01],
          [ 2.2319e-02, -5.2064e-02, -1.9440e-01, -1.7895e-03,  8.3709e-02]]],

        [[[ 1.4024e-01,  6.4016e-02,  1.6549e-01,  9.6163e-02,  1.8803e-01],
          [-5.8840e-02, -1.8487e-01,  1.8037e-01,  7.3717e-02,  1.9991e-01],
          [ 7.9629e-02, -1.1025e-01,  1.2504e-01,  4.6581e-02,  2.2388e-04],
          [-3.6367e-02,  9.8296e-02,  6.5209e-02,  1.7801e-01,  1.3420e-01],
          [ 1.4725e-01, -1.9269e-01,  1.9282e-02, -1.3924e-01, -6.2607e-02]]]],
       requires_grad=True))]
[('weight_mask', tensor([[[[1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 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., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.],
          [1., 1., 1., 1., 1.]]]]))]
odict_keys(['conv1.bias', 'conv1.weight_orig', 'conv1.weight_mask', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])

使用remove函数后

prune.remove(module, 'weight')
print(list(module.named_parameters()))
print(list(module.named_buffers()))
print(model.state_dict().keys())

输出:

[('bias', Parameter containing:
tensor([ 0.1144, -0.1641,  0.0962], requires_grad=True)), ('weight', Parameter containing:
tensor([[[[ 0.0142,  0.1698, -0.0730, -0.0358, -0.1309],
          [ 0.1520,  0.1900, -0.0843,  0.0950,  0.1674],
          [-0.1724,  0.1453, -0.1764,  0.0345, -0.1767],
          [ 0.0727,  0.1170,  0.1585, -0.0713, -0.0158],
          [ 0.1485, -0.0270, -0.0164,  0.0889,  0.1170]]],

        [[[ 0.0000, -0.0000, -0.0000,  0.0000, -0.0000],
          [ 0.0000,  0.0000, -0.0000,  0.0000, -0.0000],
          [ 0.0000,  0.0000, -0.0000, -0.0000,  0.0000],
          [ 0.0000,  0.0000, -0.0000, -0.0000, -0.0000],
          [-0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],

        [[[-0.1681,  0.1801, -0.0567, -0.0366,  0.0085],
          [ 0.0495,  0.0320, -0.0127, -0.1761, -0.0948],
          [ 0.1340,  0.1103,  0.1332, -0.1911, -0.1225],
          [ 0.0781, -0.0920, -0.1759,  0.0977,  0.0030],
          [-0.0436, -0.1694, -0.0094, -0.0553, -0.0591]]]], requires_grad=True))]
[]
odict_keys(['conv1.bias', 'conv1.weight', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias'])

可以发现,缓冲区中保存的mask参数没了

4.总结

本示例首先搭建了一个类LeNet网络模型,为了进行结构化剪枝,我们选取了LeNet的conv1模块,该模块参数包含为3×5×5的weight卷积核参数和3×1的bias参数,通过示例,我们利用torch.nn.prune中的ln_structured剪枝方法,实现了对weight的3个通道中其中一个通道进行了L2 norm结构化剪枝。

本文用到的核心函数方法:

  • model.state_dict().keys(),模型所有的张量参数,还包括了mask缓冲区的参数
  • prune.ln_structured(),Lx Norm结构化剪枝
  • prune.remove(),从模块中移除修剪重参数化。已修剪的名为name的参数将保持永久修剪,而名为name+_trans (orig)的参数将从参数列表中删除。类似地,名为name+_mask的缓冲区将从缓冲区中删除。

参考:

Torch官方剪枝教程

你可能感兴趣的:(模型压缩与加速,pytorch,pytorch,深度学习,人工智能)