【剪枝实战】使用VGGNet训练、稀疏训练、剪枝、微调等，剪枝出只有3M的模型

摘要

本次剪枝实战是基于下面这篇论文去复现的，主要是实现对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
├─vgg.py
├─train.py
├─train_sp.py
├─prune.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进行剪枝！

核心代码

四、剪枝

有了稀疏化训练的模型，我们便可以通过读取checkpoints/vgg_sp/best.pth，根据gamma系数来判断网络中各个通道的重要性，剪去不重要的通道。

运行prune.py，剪枝的结果保存为checkpoints/vgg_pruned/best.pth

查看剪枝后的VGG模型，只有3.58M，体积非常小，剪枝瘦身效果非常明显！

五、微调

对模型剪枝后，为了恢复模型的性能，我们还需要对剪枝的模型进行Finetune！

运行train_prune.py

对checkpoints/vgg_pruned/best.pth进行微调，结果生成在checkpoints/vgg_pruned_finetune/best.pth

六、总结

我们通过复现Learning Efficient Convolutional Networks through Network Slimming (ICCV 2017)论文里提到的网络大瘦身剪枝算法，完成对VGG模型进行剪枝。