摘要
本次剪枝实战是基于下面这篇论文去复现的,主要是实现对BN层的γ/gamma
进行剪枝操作,本文用到的代码和数据集都可以在我的资源中免费下载到。
相关论文:
Learning Efficient Convolutional Networks through Network Slimming (ICCV 2017)
剪枝的原理
在一个卷积-BN-激活模块中,BN层可以实现通道的缩放。如下:
BN层的具体操作有两部分:
在归一化后会进行线性变换,那么当系数gamma很小时候,对应的激活(Zout)会相应很小。这些响应很小的输出可以裁剪掉,这样就实现了bn层的通道剪枝。
剪枝的过程
第一步、使用VGGNet训练模型。保存训练结果,方便将来的比对!
第二步、在BN层网络中加入稀疏因子,训练模型。
第三步、剪枝操作。
第四步、fine-tune模型,提高模型的ACC。
接下来,我们一起实现对VGGNet的剪枝。
一、项目结构
Slimming_Demo
├─checkpoints
│ ├─vgg
│ ├─vgg_pruned
│ └─vgg_sp
├─data
│ ├─train
│ │ ├─Black-grass
│ │ ├─Charlock
│ │ ├─Cleavers
│ │ ├─Common Chickweed
│ │ ├─Common wheat
│ │ ├─Fat Hen
│ │ ├─Loose Silky-bent
│ │ ├─Maize
│ │ ├─Scentless Mayweed
│ │ ├─Shepherds Purse
│ │ ├─Small-flowered Cranesbill
│ │ └─Sugar beet
│ └─val
│ ├─Black-grass
│ ├─Charlock
│ ├─Cleavers
│ ├─Common Chickweed
│ ├─Common wheat
│ ├─Fat Hen
│ ├─Loose Silky-bent
│ ├─Maize
│ ├─Scentless Mayweed
│ └─Shepherds Purse
├─train_sp.py
└─train_prune.py
train.py
:训练脚本,训练VGGNet原始模型 vgg.py
:模型脚本 train_sp.py
:稀疏训练脚本。 prune.py
:模型剪枝脚本。 train_prune.py
:微调模型脚本。
二、正常训练
首先,我们进行正常训练出一个模型,用来和剪枝后的模型进行对比。原始模型有67.8M,体积非常大。
运行train.py,正常训练结果保存为checkpoints/vgg/best.pth
通过Tensorboard工具,我们观察到,此时的BN层的gamma系数并不稀疏,很难判断出各个通道的重要性,不利于剪枝!
三、稀疏训练
运行train_sp.py,稀疏化训练结果保存为checkpoints/vgg_sp/best.pth
使用tensorboard查看runs/路径下保存的log,查看稀疏化训练结果。
可以看到通过稀疏化训练,==此时的gamma系数已经呈现趋近0的分布,==我们可以对gamma接近0的channel进行剪枝!
核心代码