参数量降低之剪枝

一、论文

https://openaccess.thecvf.com/content_ICCV_2017/papers/Liu_Learning_Efficient_Convolutional_ICCV_2017_paper.pdf

二、参考代码

GitHub - foolwood/pytorch-slimming: Learning Efficient Convolutional Networks through Network Slimming, In ICCV 2017.

三、剪枝流程(训练->剪枝->在训练)

1、网络结构图

1.1方法 

通过BN归一化里面的γ缩放系数 +稀疏化的L1范数，可以理解为 通过γ系数得到特征图比重大小，然后加上L1范数，进行稀疏化，把重要的值放大，不重要的值弄小

(原因：通过卷积层后是线性相关的，分布杂乱，使用BN归一化后把哪些偏离的离谱的分布给弄到均值为0方差为1的标准分布，这样训练的会很快，但是在BN后，感觉把数值分布强制在了非线性函数的线性区域中。于是用到了BN里面的另外两个参数γ和β，把数值进行缩放和偏移放在非线性区域)

1.2 L1的稀疏化

从图中可以看出L1的导数即梯度 是（-1，1）有稀疏性质

L2的导数即梯度类似一个线性函数

三、常用剪枝方法

四、代码实现 

第一种思路方法：

训练的时候使用L1权重约束项给BN归一化的weight具有稀疏性，

 剪枝的时候：获取所有的BN归一化里面的weight列表，使其进行排序，获取保留weight里面的最后一个值作为阈值。 然后对每一个BN归一化的weight进行mask(值为1表示需要剪枝保留的，值为0表示剪枝不需要的)然后把BN归一化需要的层数保留下来，不需要的丢弃，重新定义一个网络进行训练

1、训练模型(记住给BN里面的权重加上L1约束项，使其有稀疏性)

s = 0.0001
def updateBN():
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.weight.grad.data.add_(s*torch.sign(m.weight.data))  # L1  梯度 大于0的为1，小于0的为-1

2、剪枝模块

percent = 0.9 # 全选择前百分之多少有用的特征图
total = 0   # 记录所有的BN层的特征图的个数
for m in model.modules():
    if isinstance(m, nn.BatchNorm2d):
        total += m.weight.data.shape[0]

bn = torch.zeros(total) # 创建所有BN层的特征图的γ分值的存储空间

index = 0
for m in model.modules():
    if isinstance(m, nn.BatchNorm2d):
        size = m.weight.data.shape[0]
        bn[index:(index+size)] = m.weight.data.abs().clone()
        index += size

y, i = torch.sort(bn)  # 给bn里面的γ进行排序
percent = 0.9 # 全选择前百分之多少有用的特征图
thre_index = int(total * percent)
thre = y[thre_index]   # 找到最后一个γ的值

pruned = 0 
cfg = []  # 记录保留的特征图的个数
cfg_mask = []
for k, m in enumerate(model.modules()):
    if isinstance(m, nn.BatchNorm2d):
        weight_copy = m.weight.data.clone()
        mask = weight_copy.abs().gt(thre).float().cuda() # .gt(thre) 指的是实际的值大于                    
                                                         # thre的值，返回list，里面是0或者1
        pruned = pruned + mask.shape[0] - torch.sum(mask)  # 用于记录剪枝剪了多少层
        m.weight.data.mul_(mask)
        m.bias.data.mul_(mask)
        cfg.append(int(torch.sum(mask)))  # 记录保留的特征图的个数
        cfg_mask.append(mask.clone())  # 所有的特征图 
        print('layer index: {:d} \t total channel: {:d} \t remaining channel: {:d}'.
            format(k, mask.shape[0], int(torch.sum(mask))))
    elif isinstance(m, nn.MaxPool2d):
        cfg.append('M')


pruned_ratio = pruned/total  # 剪枝比例，指得是剪了多少

print('Pre-processing Successful!')

print(cfg)

"""
构建新的网络模型后把 一开始大模型的权重值给新模型做一个初始化，方法 根据索引把每层对应位置的权重筛选出来然后赋值给新的模型
"""

layer_id_in_cfg = 0 # 为剪枝后的模型复制权重
start_mask = torch.ones(3)  # 输入
end_mask = cfg_mask[layer_id_in_cfg]  # 输出
for [m0, m1] in zip(model.modules(), newmodel.modules()):
    if isinstance(m0, nn.BatchNorm2d):
        idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))  # 把＞阈值的特 
                                                            # 征图的索引筛选出来
        m1.weight.data = m0.weight.data[idx1].clone()
        m1.bias.data = m0.bias.data[idx1].clone()
        m1.running_mean = m0.running_mean[idx1].clone()
        m1.running_var = m0.running_var[idx1].clone()
        layer_id_in_cfg += 1
        start_mask = end_mask.clone()
        if layer_id_in_cfg < len(cfg_mask):  # do not change in Final FC
            end_mask = cfg_mask[layer_id_in_cfg]
    elif isinstance(m0, nn.Conv2d):
        idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
        idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
        print('In shape: {:d} Out shape:{:d}'.format(idx0.shape[0], idx1.shape[0]))
        w = m0.weight.data[:, idx0, :, :].clone()
        w = w[idx1, :, :, :].clone()
        m1.weight.data = w.clone()
        # m1.bias.data = m0.bias.data[idx1].clone()
    elif isinstance(m0, nn.Linear):
        idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
        m1.weight.data = m0.weight.data[:, idx0].clone()


torch.save({'cfg': cfg, 'state_dict': newmodel.state_dict()}, args.save)

3、把剪枝后的模型进行微调训练