剪枝論文一（Network Slimming）

本文介紹一種經典的模型壓縮方法

Network Slimming

，可以實作：

1. 減小模型大小
2. 減少運作時的記憶體占用
3. 在不影響精度的同時，降低計算操作數

論文中提供的示意圖如下，可以看到左側

BN

層中橙色的神經元權重較小，是以剪枝之後就去掉了這些層的連接配接。論文的思路即通過去掉

BN

層中權重較小的神經元來達到精簡網絡的目的。

實作去掉權重較小的神經元的流程如下：

1. sparsity regularization

論文提出，在訓練的時候增加稀疏正則化方法。

主要作用：令

BN

層中權值為

的神經元盡可能多，以達到更好的剪枝效果。

實作方式：添加一個新的

loss

，其大小為

BN

層中所有神經元的權值和，使用梯度下降法使這個新的

loss

盡可能小。

# sparsity-induced懲罰項的附加次梯度下降
def updateBN():
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.weight.grad.data.add_(args.s*torch.sign(m.weight.data))  # L1
           

2. 劃分保留層和剪枝層

步驟為：

1. 給定要保留層的比例，記下所有

   BN

   層大于該比例的權重。
2. 根據比例設定門檻值，根據門檻值建立一個

   mask

   ，大于門檻值的部分為

   1

   ，小于門檻值的部分為 。
3. 利用

   mask

   提取保留的神經元。

# 将這些權重排序
y, i = torch.sort(bn) # 這些權重排序
thre_index = int(total * args.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.abs().clone()
        # 小于權重thre的為0，大于的為1，即保留的部分
        mask = weight_copy.gt(thre).float().cuda()
        # 統計被剪枝的權重的總數
        pruned = pruned + mask.shape[0] - torch.sum(mask)
        # 權重和偏置分别對應相乘
        m.weight.data.mul_(mask) 
        m.bias.data.mul_(mask)
        # 記錄每個batchnorm保留的權重和權重數
        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')
           

3. 進行BN層的剪枝

進行

BN

層的剪枝，即丢棄小于門檻值的參數；

# ===================================正式剪枝====================================
# 層數
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()):

    # 對BN層進行剪枝
    if isinstance(m0, nn.BatchNorm2d):
        # 擷取大于0的所有資料的索引，使用squeeze變成向量
        idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
        if idx1.size == 1: # 隻有一個要變成數組的1個
            idx1 = np.resize(idx1,(1,))
        # 用經過剪枝後的層參數的替換原來的
        # x = (x - mean)/war*weight + data
        m1.weight.data = m0.weight.data[idx1.tolist()].clone()
        m1.bias.data = m0.bias.data[idx1.tolist()].clone()
        m1.running_mean = m0.running_mean[idx1.tolist()].clone()
        m1.running_var = m0.running_var[idx1.tolist()].clone()
        # 下一層
        layer_id_in_cfg += 1
        # 目前在處理的層的mask
        start_mask = end_mask.clone()
        # 全連接配接層不做處理
        if layer_id_in_cfg < len(cfg_mask):
            end_mask = cfg_mask[layer_id_in_cfg]
           

4. 進行卷積層剪枝

根據前後

BN

層的保留層，可以計算得到卷積層保留的卷積核大小（上層

BN

層輸出，下層

BN

層輸入），保留前後

BN

的對應保留的元素，其餘剪枝。

# 對卷積層進行剪枝
    elif isinstance(m0, nn.Conv2d):
        # 上一層BN層對應的輸出mask為start_mask
        idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
        # 下一層BN層對應的輸入mask為start_mask
        idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
        print('In shape: {:d}, Out shape {:d}.'.format(idx0.size, idx1.size))
        if idx0.size == 1:
            idx0 = np.resize(idx0, (1,))
        if idx1.size == 1:
            idx1 = np.resize(idx1, (1,))
        # 剪枝
        [c_out, c_int, k, k]
        w1 = m0.weight.data[:, idx0.tolist(), :, :].clone()
        w1 = w1[idx1.tolist(), :, :, :].clone()
        m1.weight.data = w1.clone()
           

5. 對FC層進行剪枝

由于

FC

層的輸出是固定的（分類類數），是以隻對

FC

層的輸入次元進行剪枝，也是根據上一層

BN

層的輸出，對應保留的元素，其餘剪枝。

# 對全連接配接層進行剪枝
    elif isinstance(m0, nn.Linear):
        idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
        if idx0.size == 1:
            idx0 = np.resize(idx0, (1,))
        # 剪枝
        m1.weight.data = m0.weight.data[:, idx0].clone()
        m1.bias.data = m0.bias.data.clone()
           

源代碼