恒达娱乐 > SEO算法 > 专栏 深度学习算法优化系列二 基于Pytorch的模恒达登录型剪枝代码实战

专栏 深度学习算法优化系列二 基于Pytorch的模恒达登录型剪枝代码实战

admin SEO算法 2020年02月11日

  昨天讲了一篇ICLR 2017《Pruning Filters for Efficient ConvNets》 ,相信大家对模型剪枝有一定的了解了。今天我就剪一个简单的网络,体会一下模型剪枝的魅力。本文的代码均放在我的github工程,我是克隆了一个原始的pytorch模型压缩工程,然后我最近会公开一些在这个基础上新增的自测结果,一些经典的网络压缩benchmark,一些有趣的实验。欢迎关注,github地址见文后。最后申明一下,本人处于初学阶段,肯定了解的知识很浅并且会犯很多错误,有错误之处欢迎大家指出并和我交流讨论。

  剪枝代码在 prune/normal_regular_prune.py 中。通道剪枝的方法多种多样,恒达注册登录这个工程所用的剪枝方法是统计每个卷积层后面接的BN层的 weight 的绝对值,也就是BN层的 gamma 参数。BN层的公式可以表示为:

  那么 beta 就是BN层的 bias 参数,剪枝的时候将BN层的每个缩放系数即 scale 当成每一个通道的重要程度即可。然后,根据我们预先设置的剪枝比例 percent 和网络中所有BN层的 weight 参数组成的数组确定剪枝的权重阈值 thre_0 。有了这个阈值就可以自行预剪枝和剪枝操作了。

  首先确定剪枝的全局阈值,然后根据阈值得到剪枝后的网络每层的通道数 cfg_mask ,这个 cfg_mask 就可以确定我们剪枝后的模型的结构了,注意这个过程只是确定每一层那一些索引的通道要被剪枝掉并获得 cfg_mask ,还没有真正的执行剪枝操作。我给代码加了部分注释,应该不难懂。

  在预剪枝之后我们获得了每一个特征图需要剪掉哪些通道数的索引列表,接下来我们就可以按照这个列表执行剪枝操作了。注意一下,在预剪枝阶段是通过BN层的 scale 参数获取的需要剪枝的通道索引,在剪枝阶段不仅仅需要剪掉BN层的对应通道,还要剪掉BN层前的卷积层的对应通道。剪枝的完整代码如下:

  # 注意卷积核Tensor维度为[n, c, w, h],两个卷积层连接,下一层的输入维度n就等于当前层的c

  这是深度学习算法优化系列的第二篇文章,之后会陆续学习和更新这个系列,恒达娱乐不介意可以给公众号点个关注哦。

  AI 研习社专栏旨在满足有持续创作能力的的知名学者、技术专家、优秀博主、讲师和人工智能人才文字创作和打造专属个人IP的需求。AI 研习社希望能够帮助大咖见证求知,分享洞见~

标签: