pytorch自带网络_Pytorch网络压缩系列教程一:Prune你的模型

Pytorch网络压缩系列教程一:Prune你的模型

本文由林大佬原创，转载请注明出处，来自腾讯、阿里等一线AI算法工程师组成的QQ交流群欢迎你的加入: 1037662480

深度学习模型取得了前所未有的巨大成功, 并且引领着新的科技潮流. 但学术界流传的深度学习模型, 通常都是超参数化的, 在更低功耗, 便携式的设备上部署时, 会遇到很多问题. 为了将我们的模型进一步压缩, 同时保证模型具有和以前一样的精确度, 一个新的领域开始备受瞩目: 网络压缩与加速. 总的来说, 这里面分为一个不同的领域, 但是实际上做的是同一件事:

• 剪枝(Prune);
• 稀疏训练;
• 蒸馏

对于Prune来说, 这是有很好的范例和理论依据来做这件事的, 现在越来越多的公司也在这个领域发力, 不断的将深度学习运算的极限push到更高的值, 让低功耗设备下的深度学习模型运算成为可能.

正式由于现有的深度学习模型都存在超参数化的现象, 对于一些对于整体网络贡献度不高的层, 或者参数, 我们可以将其裁剪掉, 进而使得网络轻量化. 如果你正在研究, 学习这个方向, 这篇教程可以让你足够入门.

我们将结合pytorch和实际的例子来教大家如何来prune你的模型, 这篇文章不仅仅教你找到你要剪掉的层, 同时还会告诉你如何保存你减掉的模型, 并且这个模型你可以拿来推理, 可以导出简化模型的onnx并部署.

01. Prune Lenet

我们先从一个十分简单的例子开始, 比如LeNet. 这是一个只有两个卷积层和三个全连接层的极简网络. 它的结构如下:

LeNet(
  (conv1): Conv2d(1, 6, kernel_size=(3, 3), stride=(1, 1))
  (conv2): Conv2d(6, 16, kernel_size=(3, 3), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)

即便是这么一个网络, 我们也是可以对它进行压缩的. 事实上, 我们会先从pytorch自带的Prune模块入手, 教大家如何去寻找网络的稀疏度, 但如果你想要更高层次的东西, 比如Distiller里面那种通过修改每一层再做evaluation, 测试每一层的影响, 并生成剪枝表的更复杂操作, 我们会在后面的教程跟大家讲.

先来看看如何用pytorch来做.

实际上在pytorch里面, prune通过对权重进行掩码来完成. 这个如何理解?

首先, 我们打印一下原始的conv1的权重看看:

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

这里我列举了后面可能会用到三个方法, 这个可以查看当前的module到底是一个啥情况.

然后我们着重观察 named_parameters 因为参数都保存在这里, 打印完了之后你可以看到:

[('weight', Parameter containing:
tensor([[[[-0.2312,  0.2133, -0.1313],
          [-0.2980, -0.1838, -0.2902],
          [-0.3006,  0.1338, -0.0980]]],
        [[[-0.1239,  0.1060,  0.3271],
          [-0.0301, -0.0245,  0.0493],
          [-0.0160,  0.0397, -0.1242]]]], device='cuda:0', requires_grad=True)), ('bias', Parameter containing:
tensor([-0.2593, -0.0520,  0.0303,  0.0382, -0.0468, -0.1053], device='cuda:0',
       requires_grad=True))]
​

它有一个weight和一个bias, 这没错, 合乎常理. 我们甚至可以看看weights的尺寸是多少.

for a in module.named_parameters():
    print(a[1].shape)

输出:

torch.Size([6, 1, 3, 3])
torch.Size([6])

这个其实就是卷积的维度了, 6指的是channel, 1值得还是stride, 3指的是kernel size.

然后重点来了, 我们要开始做prune了, 在pytorch里面操作也很简单, 只需要一行代码:

import torch.nn.utils.prune as prune
​
prune.random_unstructured(module, name='weight', amount=0.3)
​

这个可以从众多的剪枝方法中, 选择一个很好的手段来完成同样的目的. 然后我们再打印一下named_parameters:

[('bias', Parameter containing:
tensor([-0.2281,  0.3085,  0.0937, -0.0540,  0.3295,  0.1107], device='cuda:0',
       requires_grad=True)), ('weight_orig', Parameter containing:
tensor([[[[ 0.1934, -0.0172, -0.1957],
          [ 0.1655,  0.1669, -0.2448],
          [-0.2250, -0.0963, -0.0195]]],
​
        [[[-0.3154,  0.1868,  0.0103],
          [-0.2245,  0.1548,  0.2567],
          [ 0.0713,  0.1262,  0.1547]]]], device='cuda:0', requires_grad=True))]
​

唯一的变化就是 weights 变成了 weights_orig, prune之后通过掩码的方式存放在了 named_buffers里面:

print(list(module.named_buffers()))

可以看到:

[('weight_mask', tensor([[[[1., 1., 1.],
          [0., 1., 0.],
          [0., 1., 1.]]],

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

那么问题来了, 你只是把权重进行了掩码, 那么我要知道你剪掉了哪几个channel怎么办? 而且你这个是剪的权重, 结构呢? 我怎么把这个结构找出来??

所以说这只是第一步, 接下来我们来看看结构化修剪. 结构化修剪讲道理你可以知道你修剪了哪些结构.

一个比较好的结构化修建的例子是通过沿着Tensor的某个维度进行裁剪, 这样你可以直接看到维度的变化. 关于结构化剪枝和非结构的区别, 大家可以看看我们的知乎这个回答:

https://www.zhihu.com/question/391195715/answer/1421770287

我们现在开始修剪模块, 比如上面的LeNet的conv1层, 首先我们可以从prune层里面拿一个我们喜欢的技术, 比如基于 ln范数 的评判标准来进行结构话的裁剪.

prune.ln_structured(module, name='weight', amount=0.5, n=2, dim=0)

这个操作之后, 我们得到的将是一个新的权重, 和上面的非结构化的不同的地方在于, 这里是整个矩阵的一行为零, 上面我们用的dim=0, 那么就是channel这一个维度, 会有50%为零.

讲到这里, 你们很容易复现, 也很容易理解, 我在编写这篇教程的时候与参考了很多博客, 文章, 我发现即便是你看完了你还是无法这些读者遇到的问题:

简单来说, 大家其实关注的本身不在于你怎么去做的prune, 而是prune完了之后我要怎么用.

很可惜网上大部分教程到此戛然而止, 没有人告诉你接下来要怎么做, 网络如何保存? 如何训练? 如何部署? 如何把新的结构提取出来?等等.

这就是你要关注本专栏或者本公众号的根本原因: 我们做后20%的事情.

02. 保存裁剪后的模型

接下来是我们的重点. 如何将这些模型保存下来. 换句话说, 你剪枝完了之后, 你得到了一个mask的掩码, 不管是结构化还是非结构化, 这个新的模型你要怎么样才能导出呢? 试想一下, 如果只是得到这么一个掩码, 你在推理的时候, 只是这一部分权重为零而已, 你推理的长宽高啥的都么有改变, so, 速度会变得更快吗?

事实上善于谷歌的我们, 可以在stackoverflow轻松找到一位老哥试图复现的代码:

stackoverflow

这位老哥发现, 他尝试不管是结构化剪枝还是非结构化剪枝, 速度和模型的大小都没有变化. 然后另外一个老哥就告诉他真像了:

重要的是要了解非结构化修剪和结构化修剪之间的区别。
结构化修剪：通过删除张量的整行/列来减小重量张量的尺寸。这转化为去除神经元及其所有传入和传出连接（在密集层中）或整个卷积过滤器（在卷积层中）。
非结构化修剪：单个权重可以“去除”（归零），而不受最终张量形状的限制。这意味着删除神经元之间的单个连接（在密集层中）或删除卷积过滤器的单个权重（在卷积层中）。
请注意，生成的重量张量可以稀疏，但保持其原始形状。目前，torch.nn.utils.prune仅支持非结构化修剪，这几乎无法帮助降低推理成本，因为GPU并未针对稀疏矩阵乘法进行优化。虽然您可能希望减小权重张量的尺寸以减少浮点运算的数量，但非结构化修剪会产生带有许多零的权重张量，但不会自动减小此类张量的大小。仅当去除许多负担时，非结构化修剪才能帮助提高性能。在这种情况下，您可以依靠PyTorch稀疏操作，也可以尝试查找包含全零的行/列，因此可以将其删除。相反，如果要研究结构化修剪，可以看一下TorchPruner，这是我为研究目的而自行开发的一个库，它提供了实用程序来查找最不重要的神经元并相应地对重量张量进行切片。

所以, 这个部分我们也应该知道, 事实上pytorch的prune里面做的事情, 只是找到哪些层是可以置零的而已, 请注意, 这里的层指的不是某个layer, 而是一个卷积层矩阵里面的某一整行或者某一整列.

那么问题来了, 听起来我们要保存模型, 要得到速度加快, 模型变小的模型, 还需要再每一个prune的层里面找到某一被全置零的链接, 然后再拿来做推理.

可是这个怎么听起来这么复杂呢? 事实上在实际应用中应该怎么做呢?

这就引入了我们的下一篇教程: 使用Distiller的工业应用模型剪枝教程.

更多

http://weixin.qq.com/r/TjgIEHTEzymwrTBz921c (二维码自动识别)

如果你想学习人工智能，对前沿的AI技术比较感兴趣，可以加入我们的知识星球，获取第一时间资讯，前沿学术动态，业界新闻等等！你的支持将会鼓励我们更频繁的创作，我们也会帮助你开启更深入的深度学习之旅！

往期文章

https://zhuanlan.zhihu.com/p/165009477

https://zhuanlan.zhihu.com/p/149398749

https://zhuanlan.zhihu.com/p/147622974

https://zhuanlan.zhihu.com/p/144727162