• Slimming剪枝方法

    本文参考：5-剪枝后模型参数赋值_哔哩哔哩_bilibiliz

     https://github.com/foolwood/pytorch-slimming

    一、模型剪枝理论说明

    论文：Learning Efficient Convolutional Networks through Network Slimming

    （1）卷积后得到多个特征图（channel=64, 128, 256…），这些图不一定都重要，所以量化计算特征图的重要性

    （2）训练模型的时候需要加入一些策略，让权重参数有明显的大小之分，从而筛选重要的特征图

    Channel scaling factors里面的数值为特征图的打分，直观理解为分值大的特征图需要保留，分值小的特征图可以去掉。

    二、计算特征图重要性

    Network slimming ，利用BN层中的缩放因子Ƴ

    BN的理论支持：

    ，使得数据为（0，1）正态分布。

    整体感觉是一个归一化操作，但是BN中需要额外引入两个可训练的参数：Ƴ和β

    BatchNorm的本质：

    （1）BN要做的就是把越来越偏离的分布给拉回来

    （2）再重新规范化到均值为0方差为1的标准正态分布

    （3）这样能够使得激活函数在数值层面更敏感，训练更快。

    （4）产生的问题：经过BN之后，把数值分布强制在了非线性函数的线性区域中。

    针对第（3）点解释：

    在激活函数中，两边处于饱和区域不敏感，接近于0位置非饱和处于敏感区域。

    针对第（4）点解释：

    BN将数据强制压缩到中间红色区域的线性部分，F(x)只做仿射变化,F=sigmoid，多个仿射变化的叠加仍然是仿射变化，添加再多隐藏层与单层神经网络是等价的。

    所以，BN需要保证一些非线性，对规范后的结果再进行变化。

    添加两个参数后重新训练：

    ，这两个参数是网络训练过程中得到的，而不是超参给的。

    该公式相当于BN的逆变换，

    相当于对正态分布进行一些改变，拉动一下，变一下形状，做适当的还原。

    Ƴ值越大越重要，那么该特征图调整的幅度越大，说明该特征图越重要。

    三、让特征图重要度两极分化更明显

    使用L1正则化对参数进行稀疏操作。

    L1求导后为：sign(Θ)，相当于稳定前进，都为1，最后学成0了

    L2求导后为：Θ，相当于越来越慢，很多参数都接近0，平滑。

    论文核心：

    四、剪枝流程

    使用到的vgg模型架构：

    1. import torch
    2. import torch.nn as nn
    3. import math
    4. from torch.autograd import Variable
    5.  
    6. class vgg(nn.Module):
    7.     def __init__(self, dataset='cifar10', init_weights=True, cfg=None):
    8.         super(vgg, self).__init__()
    9.         if cfg is None:
    10.             cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512]
    11.         self.feature = self.make_layers(cfg, True)
    12.  
    13.         if dataset == 'cifar10':
    14.             num_classes = 10
    15.         elif dataset == 'cifar100':
    16.             num_classes = 100
    17.         self.classifier = nn.Linear(cfg[-1], num_classes)
    18.         if init_weights:
    19.             self._initialize_weights()
    20.  
    21.  
    22.     def make_layers(self, cfg, batch_norm=False):
    23.         layers = []
    24.         in_channels = 3
    25.         for v in cfg:
    26.             if v == 'M':
    27.                 layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
    28.             else:
    29.                 conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1, bias=False)
    30.                 if batch_norm:
    31.                     layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
    32.                 else:
    33.                     layers += [conv2d, nn.ReLU(inplace=True)]
    34.                 in_channels = v
    35.         return nn.Sequential(*layers)
    36.  
    37.     def forward(self, x):
    38.         x = self.feature(x)
    39.         x = nn.AvgPool2d(2)(x)
    40.         x = x.view(x.size(0), -1)
    41.         y = self.classifier(x)
    42.         return y
    43.  
    44.     def _initialize_weights(self):
    45.         for m in self.modules():
    46.             if isinstance(m, nn.Conv2d):
    47.                 n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
    48.                 m.weight.data.normal_(0, math.sqrt(2. / n))
    49.                 if m.bias is not None:
    50.                     m.bias.data.zero_()
    51.             elif isinstance(m, nn.BatchNorm2d):
    52.                 m.weight.data.fill_(0.5)
    53.                 m.bias.data.zero_()
    54.             elif isinstance(m, nn.Linear):
    55.                 m.weight.data.normal_(0, 0.01)
    56.                 m.bias.data.zero_()
    57.  
    58. if __name__ == '__main__':
    59.     net = vgg()
    60.     x = Variable(torch.FloatTensor(16, 3, 40, 40))
    61.     y = net(x)
    62.     print(y.data.shape)

    1、原始模型训练：

    （1）BN的L1稀疏正则化：使用次梯度下降法，对BN层的权重进行再调整

    （2）训练完成后主要保存原始模型的参数信息

    1. import torch
    2. import torch.nn as nn
    3. import torch.nn.functional as F
    4. import torch.optim as optim
    5. from torchvision import datasets, transforms
    6. from torch.autograd import Variable
    7. from vgg import vgg
    8. import shutil
    9. from tqdm import tqdm
    10.  
    11. learning_rate = 0.1
    12. momentum = 0.9
    13. weight_decay = 1e-4
    14. epochs = 3
    15. log_interval = 100
    16. batch_size = 100
    17. sparsity_regularization = True
    18. scale_sparse_rate = 0.0001
    19.  
    20. checkpoint_model_path = 'checkpoint,pth.tar'
    21. best_model_path = 'model_best.pth.tar'
    22.  
    23. train_loader = torch.utils.data.DataLoader(
    24.     datasets.CIFAR10('D:\\ai_data\\cifar10', train=True, download=True,
    25.                      transform=transforms.Compose([
    26.                          transforms.Pad(4),
    27.                          transforms.RandomCrop(32),
    28.                          transforms.RandomHorizontalFlip(),
    29.                          transforms.ToTensor(),
    30.                          transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    31.                      ])),
    32.     batch_size=batch_size, shuffle=True)
    33.  
    34. test_loader = torch.utils.data.DataLoader(
    35.     datasets.CIFAR10('D:\\ai_data\\cifar10', train=False, transform=transforms.Compose([
    36.         transforms.ToTensor(),
    37.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    38.     ])),
    39.     batch_size=batch_size, shuffle=True)
    40.  
    41. model = vgg()
    42. model.cuda()
    43. optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
    44.  
    45. def train(epoch):
    46.     model.train()
    47.     for batch_idx, (data, target) in enumerate(tqdm(train_loader)):
    48.         data, target = data.cuda(), target.cuda()
    49.         data, target = Variable(data), Variable(target)
    50.         optimizer.zero_grad()
    51.         output = model(data)
    52.         loss = F.cross_entropy(output, target)
    53.         loss.backward()
    54.         if sparsity_regularization:
    55.             updateBN()
    56.         optimizer.step()
    57.         if batch_idx % log_interval == 0:
    58.             print('Train Epoch: {} [{}/{} ({:.1f}%)]\tLoss: {:.6f}'.format(
    59.                 epoch, batch_idx * len(data), len(train_loader.dataset),
    60.                        100. * batch_idx / len(train_loader), loss.item()))
    61.  
    62. def test():
    63.     model.eval()
    64.     test_loss = 0
    65.     correct = 0
    66.     for data, target in tqdm(test_loader):
    67.         data , target = data.cuda(), target.cuda()
    68.         data, target = Variable(data), Variable(target)
    69.         output = model(data)
    70.         test_loss += F.cross_entropy(output, target, size_average=False).item()
    71.         pred = output.data.max(1, keepdim=True)[1]
    72.         correct += pred.eq(target.data.view_as(pred)).cpu().sum()
    73.  
    74.     test_loss /= len(test_loader.dataset)
    75.     print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.1f}%)\n'.format(
    76.         test_loss, correct, len(test_loader.dataset),
    77.         100. * correct / len(test_loader.dataset)))
    78.     return correct / float(len(test_loader.dataset))
    79.  
    80. def save_checkpoint(state, is_best, filename=checkpoint_model_path):
    81.     torch.save(state, filename)
    82.     if is_best:
    83.         shutil.copyfile(filename, best_model_path)
    84.  
    85. def updateBN():
    86.     for m in model.modules():
    87.         if isinstance(m, nn.BatchNorm2d):
    88.             m.weight.grad.data.add_(scale_sparse_rate * torch.sign(m.weight.data)) # L1，使用次梯度下降
    89.  
    90. best_prec = 0
    91. for epoch in range(epochs):
    92.     train(epoch)
    93.     prec = test()
    94.     is_best = prec > best_prec
    95.     best_prec = max(prec, best_prec)
    96.     save_checkpoint({
    97.         'epoch': epoch + 1,
    98.         'state_dict': model.state_dict(),
    99.         'best_prec': best_prec,
    100.         'optimizer': optimizer.state_dict()
    101.     }, is_best)

    2、模型剪枝

    （1）剪枝过程主要分为两部分：第一部分是计算mask，第二部分是根据mask调整各层的shape

    （2）BN层通道数：Conv -> BN -> ReLU -> MaxPool--à Linear，所以BN的输入维度对应Conv的输出通道数

    （3）BN层总通道数：将所有BN层的通道数进行汇总

    （4）BN层剪枝百分位：取总通道数的百分位得到具体的float值，大于该值的通道对应的mask置为1，否则对应的mask置为0

    （5）改变权重weight：BN层抽取mask为1的通道数的值，该操作会改变BN的shape，从而上下游操作中的Conv和Linear也需要被动做出调整，对Maxpool和ReLu的通道数无影响

    （6）Conv层的参数为[out_channels, in_channels, kernel_size1, kernel_size2]，所以需要调整两次，先对in_channels进行调整，再对out_channels进行调整。Conv初始输入为RGB的3通道。

    假如计算出的保留通道数信息为：

    [48, 60, 115, 118, 175, 163, 141, 130, 259, 267, 258, 249, 225, 212, 234, 97]

    Conv的输入输出变为：

    In shape: 3 Out shape:48

    In shape: 48 Out shape:60

    In shape: 60 Out shape:115

    In shape: 115 Out shape:118

    ……

    In shape: 234 Out shape:97

    （7）保存模型时，一方面把有用的参数信息保存了下来，同时剪枝后的最新的模型结构参数也保存了，方便后续再训练时构建新的模型结构

    1. import os
    2. import torch
    3. import torch.nn as nn
    4. from torch.autograd import Variable
    5. from torchvision import datasets, transforms
    6. from vgg import vgg
    7. import numpy as np
    8. from tqdm import tqdm
    9.  
    10. percent = 0.5
    11. batch_size = 100
    12. raw_model_path = 'model_best.pth.tar'
    13. save_model_path = 'prune_model.pth.tar'
    14.  
    15. model = vgg()
    16. model.cuda()
    17. if os.path.isfile(raw_model_path):
    18.     print("==> loading checkpoint '{}'".format(raw_model_path))
    19.     checkpoint = torch.load(raw_model_path)
    20.     start_epoch = checkpoint['epoch']
    21.     best_prec = checkpoint['best_prec']
    22.     model.load_state_dict(checkpoint['state_dict'])
    23.     print("==> loaded checkpoint '{}'(epoch {}) Prec:{:f}".format(raw_model_path, start_epoch, best_prec) )
    24. print(model)
    25.  
    26. total = 0
    27. for m in model.modules():
    28.     if isinstance(m, nn.BatchNorm2d):
    29.         total += m.weight.data.shape[0]
    30.  
    31. bn = torch.zeros(total)
    32. index = 0
    33. for m in model.modules():
    34.     if isinstance(m, nn.BatchNorm2d):
    35.         size = m.weight.data.shape[0]
    36.         bn[index : index + size] = m.weight.data.abs().clone()
    37.         index += size
    38.  
    39. y, i = torch.sort(bn)
    40. thre_index = int(total * percent)
    41. thre = y[thre_index]
    42.  
    43. pruned = 0
    44. cfg = []
    45. cfg_mask = []
    46. for k, m in enumerate(model.modules()):
    47.     if isinstance(m, nn.BatchNorm2d):
    48.         weight_copy = m.weight.data.clone()
    49.         mask = weight_copy.abs().gt(thre).float().cuda()
    50.         pruned += mask.shape[0] - torch.sum(mask)
    51.         m.weight.data.mul_(mask)
    52.         m.bias.data.mul_(mask)
    53.         cfg.append(int(torch.sum(mask)))
    54.         cfg_mask.append(mask.clone())
    55.         print('layer index: {:d} \t total channel: {:d} \t remaining channel: {:d}'.format(k, mask.shape[0], int(torch.sum(mask))))
    56.     elif isinstance(m, nn.MaxPool2d):
    57.         cfg.append('M')
    58. pruned_ratio = pruned / total
    59.  
    60. print('pruned_ratio: {},Pre-processing Successful!'.format(pruned_ratio))
    61.  
    62. # simple test model after Pre-processing prune(simple set BN scales to zeros)
    63. def test():
    64.     test_loader = torch.utils.data.DataLoader(
    65.         datasets.CIFAR10('D:\\ai_data\\cifar10', train=False, transform=transforms.Compose([
    66.             transforms.ToTensor(),
    67.             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    68.         ])),
    69.         batch_size=batch_size, shuffle=True)
    70.     model.eval()
    71.  
    72.     correct = 0
    73.     for data, target in tqdm(test_loader):
    74.         data, target = data.cuda(), target.cuda()
    75.         data, target = Variable(data), Variable(target)
    76.         output = model(data)
    77.         pred = output.data.max(1, keepdim=True)[1]
    78.         correct += pred.eq(target.data.view_as(pred)).cpu().sum()
    79.  
    80.     print('\nTest set: Accuracy: {}/{} ({:.1f}%)\n'.format(
    81.         correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)))
    82.  
    83. test()
    84.  
    85. # make real prune
    86. print(cfg)
    87. new_model = vgg(cfg=cfg)
    88. new_model.cuda()
    89.  
    90. layer_id_in_cfg = 0  # cfg中的层数索引
    91. start_mask = torch.ones(3)
    92. end_mask = cfg_mask[layer_id_in_cfg]
    93. for [m0,  m1] in zip(model.modules(), new_model.modules()):
    94.     if isinstance(m0, nn.BatchNorm2d):
    95.         idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
    96.         m1.weight.data = m0.weight.data[idx1].clone()
    97.         m1.bias.data = m0.bias.data[idx1].clone()
    98.         m1.running_mean = m0.running_mean[idx1].clone()
    99.         m1.running_var = m0.running_var[idx1].clone()
    100.         layer_id_in_cfg += 1
    101.         start_mask = end_mask.clone()
    102.         if layer_id_in_cfg < len(cfg_mask):  # do not change in Final FC
    103.             end_mask = cfg_mask[layer_id_in_cfg]
    104.     elif isinstance(m0, nn.Conv2d):
    105.         idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
    106.         idx1 = np.squeeze(np.argwhere(np.asarray(end_mask.cpu().numpy())))
    107.         print('In shape: {:d} Out shape:{:d}'.format(idx0.shape[0], idx1.shape[0]))
    108.         w = m0.weight.data[:, idx0, :, :].clone()
    109.         w = w[idx1, :, :, :].clone()
    110.         m1.weight.data = w.clone()
    111.     elif isinstance(m0, nn.Linear):
    112.         idx0 = np.squeeze(np.argwhere(np.asarray(start_mask.cpu().numpy())))
    113.         m1.weight.data = m0.weight.data[:, idx0].clone()
    114.  
    115. torch.save({'cfg': cfg, 'state_dict': new_model.state_dict()}, save_model_path)
    116. print(new_model)
    117. model = new_model
    118. test()
    119.  
    120.

    3、再训练

    剪枝后保存的模型参数相当于训练过程中的一个checkpoint，根据新的模型结构，在此checkpoint的基础上再进行训练，直到得到满意的指标。

    1. import torch
    2. import torch.nn.functional as F
    3. import torch.optim as optim
    4. from torchvision import datasets, transforms
    5. from torch.autograd import Variable
    6. from vgg import vgg
    7. import shutil
    8. from tqdm import tqdm
    9.  
    10. learning_rate = 0.1
    11. momentum = 0.9
    12. weight_decay = 1e-4
    13. epochs = 3
    14. log_interval = 100
    15. batch_size = 100
    16. sparsity_regularization = True
    17. scale_sparse_rate = 0.0001
    18.  
    19. prune_model_path = 'prune_model.pth.tar'
    20. prune_checkpoint_path = 'pruned_checkpoint.pth.tar'
    21. prune_best_model_path = 'pruned_model_best.pth.tar'
    22.  
    23.  
    24. train_loader = torch.utils.data.DataLoader(
    25.     datasets.CIFAR10('D:\\ai_data\\cifar10', train=True, download=True,
    26.                      transform=transforms.Compose([
    27.                          transforms.Pad(4),
    28.                          transforms.RandomCrop(32),
    29.                          transforms.RandomHorizontalFlip(),
    30.                          transforms.ToTensor(),
    31.                          transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    32.                      ])),
    33.     batch_size=batch_size, shuffle=True)
    34.  
    35. test_loader = torch.utils.data.DataLoader(
    36.     datasets.CIFAR10('D:\\ai_data\\cifar10', train=False, transform=transforms.Compose([
    37.         transforms.ToTensor(),
    38.         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    39.     ])),
    40.     batch_size=batch_size, shuffle=True)
    41.  
    42. checkpoint = torch.load(prune_model_path)
    43. model = vgg(cfg=checkpoint['cfg'])
    44. model.cuda()
    45. model.load_state_dict(checkpoint['state_dict'])
    46.  
    47. optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=weight_decay)
    48.  
    49. def train(epoch):
    50.     model.train()
    51.     for batch_idx, (data, target) in enumerate(tqdm(train_loader)):
    52.         data, target = data.cuda(), target.cuda()
    53.         data, target = Variable(data), Variable(target)
    54.         optimizer.zero_grad()
    55.         output = model(data)
    56.         loss = F.cross_entropy(output, target)
    57.         loss.backward()
    58.         optimizer.step()
    59.         if batch_idx % log_interval == 0:
    60.             print('Train Epoch: {} [{}/{} ({:.1f}%)]\tLoss: {:.6f}'.format(
    61.                 epoch, batch_idx * len(data), len(train_loader.dataset),
    62.                        100. * batch_idx / len(train_loader), loss.item()))
    63.  
    64. def test():
    65.     model.eval()
    66.     test_loss = 0
    67.     correct = 0
    68.     for data, target in tqdm(test_loader):
    69.         data , target = data.cuda(), target.cuda()
    70.         data, target = Variable(data), Variable(target)
    71.         output = model(data)
    72.         test_loss += F.cross_entropy(output, target, size_average=False).item()
    73.         pred = output.data.max(1, keepdim=True)[1]
    74.         correct += pred.eq(target.data.view_as(pred)).cpu().sum()
    75.  
    76.     test_loss /= len(test_loader.dataset)
    77.     print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.1f}%)\n'.format(
    78.         test_loss, correct, len(test_loader.dataset),
    79.         100. * correct / len(test_loader.dataset)))
    80.     return correct / float(len(test_loader.dataset))
    81.  
    82. def save_checkpoint(state, is_best, filename=prune_checkpoint_path):
    83.     torch.save(state, filename)
    84.     if is_best:
    85.         shutil.copyfile(filename, prune_best_model_path)
    86.  
    87. best_prec = 0
    88. for epoch in range(epochs):
    89.     train(epoch)
    90.     prec = test()
    91.     is_best = prec > best_prec
    92.     best_prec = max(prec, best_prec)
    93.     save_checkpoint({
    94.         'epoch': epoch + 1,
    95.         'state_dict': model.state_dict(),
    96.         'best_prec': best_prec,
    97.         'optimizer': optimizer.state_dict()
    98.     }, is_best)
    99.

    4、原始模型和剪枝后模型比较：

    在cifar10上通过vgg模型分别迭代3次。

    原始模型为156M，准确率为70%左右

    剪枝后模型为36M，准确率为76%左右

    备注：最好是原始模型达到顶峰时再剪枝，此时再比较剪枝前后的准确率影响。

     

  • 相关阅读:
    【Python大数据笔记_day05_Hive基础操作】
    刷题1：数组篇
    R语言将向量数据按照行方式转化为矩阵数据（设置参数byrow为TRUE）
    精品基于springboot的线上跳蚤市场平台
    给网站添加春节灯笼效果：引入即用，附源码！
    Bra12同态加密方案初步学习
    MybatisPlus的基础使用
    大话STL第三期——deque（双端动态数组）
    Spring Cloud 学习笔记（2 3）
    C++刷题 日期差值
  • 原文地址：https://blog.csdn.net/benben044/article/details/127886261