• 模型压缩（一）通道剪枝-BN层

    论文：https://arxiv.org/pdf/1708.06519.pdf

    BN层中缩放因子γ与卷积层中的每个通道关联起来。在训练过程中对这些比例因子进行稀疏正则化，以自动识别不重要的通道。缩放因子值较小的通道(橙色)将被修剪(左侧)。剪枝后，获得了紧凑的模型(右侧)，然后对其进行微调，以达到与正常训练的全网络相当(甚至更高)的精度。

    BN层原理：

     归一化化后，BN层服从正态分布，当γ,β趋于0时，经过阈值分离，输出为0，与之连接的卷积层输入为0。

    剪枝流程：

     

    剪枝原理：

    在BN层网络中加入稀疏因子，训练使得BN层稀疏化，对稀疏训练的后的模型中所有BN层权重进行统计排序，获取指定保留BN层数量即取得排序后权重阈值thres。遍历模型中的BN层权重，制作各层mask（权重>thres值为1，权重

    如下实现一个简单的网络剪枝。

    1、自定义一个网络

    对网络进行

    1. import torch
    2. import torch.nn as nn
    3. import numpy as np
    6. class net(nn.Module):
    7.     def __init__(self,cfg=None):
    8.         super(net, self).__init__()
    9.         if cfg:
    10.             self.features=self.make_layer(cfg)
    11.             self.linear = nn.Linear(cfg[2], 2)
    12.         else:
    13.             layers=[]
    14.             layers+=[nn.Conv2d(3,64,7,2,1,bias=False),
    15.                 nn.BatchNorm2d(64),
    16.                 nn.ReLU(inplace=True)]
    18.             layers += [
    19.                 nn.Conv2d(64,128,3,2,1,bias=False),
    20.                 nn.BatchNorm2d(128),
    21.                 nn.ReLU(inplace=True)
    22.             ]
    23.             layers += [
    24.                 nn.Conv2d(128, 256, 3, 2, 1,bias=False),
    25.                 nn.BatchNorm2d(256),
    26.                 nn.ReLU(inplace=True)
    27.             ]
    28.             layers += [nn.AvgPool2d(2)]
    29.             self.features=nn.Sequential(*layers)
    30.             self.linear=nn.Linear(256,2)
    32.     def make_layer(self,cfg):
    33.         layers=[]
    34.         layers += [nn.Conv2d(3, cfg[0], 7, 2, 1, bias=False),
    35.                    nn.BatchNorm2d(cfg[0]),
    36.                    nn.ReLU(inplace=True)]
    38.         layers += [
    39.             nn.Conv2d(cfg[0], cfg[1], 3, 2, 1, bias=False),
    40.             nn.BatchNorm2d(cfg[1]),
    41.             nn.ReLU(inplace=True)
    42.         ]
    43.         layers += [
    44.             nn.Conv2d(cfg[1], cfg[2], 3, 2, 1, bias=False),
    45.             nn.BatchNorm2d(cfg[2]),
    46.             nn.ReLU(inplace=True)
    47.         ]
    48.         layers += [nn.AvgPool2d(2)]
    49.         return nn.Sequential(*layers)
    51.     def forward(self,x):
    52.         x=self.features(x)
    53.         # print(x.shape)
    54.         x=x.view(x.size(0),-1)
    55.         x=self.linear(x)
    56.         return x

    网络参数信息：

    1. ----------------------------------------------------------------
    2.         Layer (type)               Output Shape         Param #
    3. ================================================================
    4.             Conv2d-1              [1, 64, 8, 8]           9,408
    5.        BatchNorm2d-2              [1, 64, 8, 8]             128
    6.               ReLU-3              [1, 64, 8, 8]               0
    7.             Conv2d-4             [1, 128, 4, 4]          73,728
    8.        BatchNorm2d-5             [1, 128, 4, 4]             256
    9.               ReLU-6             [1, 128, 4, 4]               0
    10.             Conv2d-7             [1, 256, 2, 2]         294,912
    11.        BatchNorm2d-8             [1, 256, 2, 2]             512
    12.               ReLU-9             [1, 256, 2, 2]               0
    13.         AvgPool2d-10             [1, 256, 1, 1]               0
    14.            Linear-11                     [1, 2]             514
    15. ================================================================
    16. Total params: 379,458
    17. Trainable params: 379,458
    18. Non-trainable params: 0
    19. ----------------------------------------------------------------
    20. Input size (MB): 0.00
    21. Forward/backward pass size (MB): 0.17
    22. Params size (MB): 1.45
    23. Estimated Total Size (MB): 1.62

    2、稀疏训练

    在BN层中各权重加入稀疏因子。

    1. def updateBN(model,s=0.0001):
    2.     for m in model.modules():
    3.         if isinstance(m,nn.BatchNorm2d):
    4.             m.weight.grad.data.add_(s*torch.sign(m.weight.data))
    6. if __name__=="__main__":
    7.     model=net()
    8.     # from torchsummary import summary
    9.     # print(summary(model,(3,20,20),1))
    10.     # x = torch.rand((1, 3, 20, 20))
    11.     # print(model(x))
    12.     optimer=torch.optim.Adam(model.parameters())
    13.     loss_fn=torch.nn.CrossEntropyLoss()
    14.     for e in range(100):
    15.         x = torch.rand((1, 3, 20, 20))
    16.         y=torch.tensor(np.random.randint(0,2,(1))).long()
    17.         out=model(x)
    18.         loss=loss_fn(out,y)
    19.         optimer.zero_grad()
    20.         loss.backward()
    21.         #BN权重稀疏化
    22.         updateBN(model)
    23.         optimer.step()
    24.     torch.save(model.state_dict(),"net.pth")

    3、剪枝

    稀疏训练后的模型，解析。

    1. import net
    2. import torch
    3. import torch.nn as nn
    4. import numpy as np
    5.  
    6. model = net.net()
    7. #加载稀疏训练的模型
    8. model.load_state_dict(torch.load("net.pth"))
    9. total = 0  # 统计所有BN层的参数量
    10. for m in model.modules():
    11.     if isinstance(m, nn.BatchNorm2d):
    12.         # print(m.weight.data.shape[0])  # 每个BN层权重w参数量：64/128/256
    13.         # print(m.weight.data)
    14.         total += m.weight.data.shape[0]
    15.  
    16. print("所有BN层总weight数量：",total)
    17.  
    18. bn_data=torch.zeros(total)
    19. index=0
    20. for m in model.modules():
    21.     #将各个BN层的参数值拷贝到bn中
    22.     if isinstance(m,nn.BatchNorm2d):
    23.         size=m.weight.data.shape[0]
    24.         bn_data[index:(index+size)]=m.weight.data.abs().clone()
    25.         index=size
    26. #对bn中的weight值排序
    27. data,id=torch.sort(bn_data)
    28. percent=0.7#保留70%的BN层通道数
    29. thresh_index=int(total*percent)
    30. thresh=data[thresh_index]#取bn排序后的第thresh_index索引值为bn权重的截断阈值
    31.  
    32. #制作mask
    33. pruned_num=0#统计BN层剪枝通道数
    34. cfg=[]#统计保存通道数
    35. cfg_mask=[]#BN层权重矩阵，剪枝的通道记为0，未剪枝通道记为1
    36.  
    37. for k,m in enumerate(model.modules()):
    38.     if isinstance(m,nn.BatchNorm2d):
    39.         weight_copy=m.weight.data.abs().clone()
    40.         # print(weight_copy)
    41.         mask=weight_copy.gt(thresh).float()#阈值分离权重
    42.         # print(mask)
    43.         # exit()
    44.         pruned_num+=mask.shape[0]-torch.sum(mask)#
    45.         # print(pruned_num)
    46.         m.weight.data.mul_(mask)#更新BN层的权重，剪枝通道的权重值为0
    47.         m.bias.data.mul_(mask)
    48.  
    49.         cfg.append(int(torch.sum(mask)))#记录未被剪枝的通道数量
    50.         cfg_mask.append(mask.clone())
    51.         print("layer index:{:d}\t total channel:{:d}\t remaining channel:{:d}".format(k,mask.shape[0],int(torch.sum(mask))))
    52.     elif isinstance(m,nn.AvgPool2d):
    53.         cfg.append("A")
    54.  
    55.  
    56. pruned_ratio=pruned_num/total
    57. print("剪枝通道占比：",pruned_ratio)
    58. print(cfg)
    59. newmodel=net.net(cfg)
    60. # print(newmodel)
    61. # from torchsummary import summary
    62. # print(summary(newmodel,(3,20,20),1))
    63.  
    64. layer_id_in_cfg=0#层
    65. start_mask=torch.ones(3)
    66. end_mask=cfg_mask[layer_id_in_cfg]#第一个BN层对应的mask
    67. # print(cfg_mask)
    68. # print(end_mask)
    69.  
    70. for(m0,m1)in zip(model.modules(),newmodel.modules()):#以最少的为准
    71.     if isinstance(m0,nn.BatchNorm2d):
    72.         # idx1=np.squeeze(np.argwhere(np.asarray(end_mask.numpy())))#获得mask中非零索引即未被减掉的序号
    73.         # print(idx1)
    74.         # exit()
    75.         # idx1=np.array([1])
    76.         # # print(idx1)
    77.         if idx1.size==1:
    78.             idx1=np.resize(idx1,(1,))
    79.             # print(idx1)
    80.         # exit()
    81.         #将旧模型的参数值拷贝到新模型中
    82.         m1.weight.data=m0.weight.data[idx1.tolist()].clone()
    83.         m1.bias.data=m0.bias.data[idx1.tolist()].clone()
    84.         m1.running_mean=m0.running_mean[idx1.tolist()].clone()
    85.         m1.running_var = m0.running_var[idx1.tolist()].clone()
    86.  
    87.         layer_id_in_cfg+=1#下一个mask
    88.         start_mask=end_mask.clone()
    89.         if layer_id_in_cfg<len(cfg_mask):
    90.             end_mask=cfg_mask[layer_id_in_cfg]
    91.     elif isinstance(m0,nn.Conv2d):#输入
    92.         idx0=np.squeeze(np.argwhere(np.asarray(start_mask.numpy())))#输入非0索引
    93.         idx1=np.squeeze(np.argwhere(np.asarray(end_mask.numpy())))#输出非0索引
    94.  
    95.         if idx0.size==1:
    96.             idx0=np.resize(idx0,(1,))
    97.         if idx1.size==1:
    98.             idx1=np.resize(idx1,(1,))
    99.  
    100.         w1=m0.weight.data[:,idx0.tolist(),:,:].clone()
    101.         w1=w1[idx1.tolist(),:,:,:].clone()
    102.         m1.weight.data=w1.clone()
    103.     elif isinstance(m0,nn.Linear):
    104.         idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.numpy())))  # 输入非0索引
    105.         if idx0.size==1:
    106.             idx0=np.resize(idx0,(1,))
    107.  
    108.         m1.weight.data=m0.weight.data[:,idx0].clone()
    109.         m1.bias.data=m0.bias.data.clone()
    110.  
    111. torch.save(newmodel.state_dict(),"prune_net.pth")
    112. print(newmodel)

    新模型结构：

    1. 所有BN层总weight数量： 448
    2. layer index:3	 total channel:64	 remaining channel:29
    3. layer index:6	 total channel:128	 remaining channel:56
    4. layer index:9	 total channel:256	 remaining channel:75
    5. 剪枝通道占比： tensor(0.6429)
    6. [29, 56, 75, 'A']
    7. net(
    8.   (features): Sequential(
    9.     (0): Conv2d(3, 29, kernel_size=(7, 7), stride=(2, 2), padding=(1, 1), bias=False)
    10.     (1): BatchNorm2d(29, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    11.     (2): ReLU(inplace=True)
    12.     (3): Conv2d(29, 56, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    13.     (4): BatchNorm2d(56, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    14.     (5): ReLU(inplace=True)
    15.     (6): Conv2d(56, 75, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    16.     (7): BatchNorm2d(75, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    17.     (8): ReLU(inplace=True)
    18.     (9): AvgPool2d(kernel_size=2, stride=2, padding=0)
    19.   )
    20.   (linear): Linear(in_features=75, out_features=2, bias=True)
    21. )
    22. ----------------------------------------------------------------
    23.         Layer (type)               Output Shape         Param #
    24. ================================================================
    25.             Conv2d-1              [1, 29, 8, 8]           4,263
    26.        BatchNorm2d-2              [1, 29, 8, 8]              58
    27.               ReLU-3              [1, 29, 8, 8]               0
    28.             Conv2d-4              [1, 56, 4, 4]          14,616
    29.        BatchNorm2d-5              [1, 56, 4, 4]             112
    30.               ReLU-6              [1, 56, 4, 4]               0
    31.             Conv2d-7              [1, 75, 2, 2]          37,800
    32.        BatchNorm2d-8              [1, 75, 2, 2]             150
    33.               ReLU-9              [1, 75, 2, 2]               0
    34.         AvgPool2d-10              [1, 75, 1, 1]               0
    35.            Linear-11                     [1, 2]             152
    36. ================================================================
    37. Total params: 57,151
    38. Trainable params: 57,151
    39. Non-trainable params: 0
    40. ----------------------------------------------------------------
    41. Input size (MB): 0.00
    42. Forward/backward pass size (MB): 0.07
    43. Params size (MB): 0.22
    44. Estimated Total Size (MB): 0.29
    45. ----------------------------------------------------------------

    模型大小由1.45m压缩到230k，压缩率：84%

    4、fine-tune训练

    1. newmodel.load_state_dict(torch.load("prune_net.pth"))
    2. #
    3. optimer=torch.optim.Adam(model.parameters())
    4. loss_fn=torch.nn.CrossEntropyLoss()
    5. for e in range(100):
    6.     x = torch.rand((1, 3, 20, 20))
    7.     y=torch.tensor(np.random.randint(0,2,(1))).long()
    8.     out=newmodel(x)
    9.     loss=loss_fn(out,y)
    10.     optimer.zero_grad()
    11.     loss.backward()
    12.     optimer.step()
    13. torch.save(newmodel.state_dict(),"prune_net.pth")

    以上过程仅供参考。

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

    Network Slimming——有效的通道剪枝方法（Channel Pruning）_Law-Yao的博客-CSDN博客_通道剪枝算法

  • 相关阅读:
    【黑马-SpringCloud技术栈】【02】服务拆分及远程调用_服务提供者与消费者
    （Note）C++数值标识符
    基于Qt QList和QMap容器类示例
    理性分析不同模型的性能指标
    嵌入式（Linux内核的安装与加载）
    leetcode298周赛记录
    STL 应用 —— set / multiset
    (02)Cartographer源码无死角解析-(10) 配置文件加载
    linux清理缓存垃圾命令和方法介绍
    QT总结汇总
  • 原文地址：https://blog.csdn.net/m0_37264397/article/details/126157647