模型压缩（2）——模型剪枝

1. 使用pytorch的prune工具进行剪枝

使用pytorch自带的prune函数进行剪枝，剪枝后被剪掉的参数为0，应将为0的参数剔除运算，否则为虚假的剪枝，速度甚至更慢

参考代码如下，model为训练后的模型，经循环得到剪枝模型，需微调回复精度，感觉没太大作用，不做具体演示。

from torch.nn.utils import prune
# 使用named_modules可以得到每一个最小层，使用named_children仅能得到较大的块
for n,module in model.named_modules():
    # 对线性层剪枝
    if isinstance(module,torch.nn.Linear):
        # 可以选择多种裁剪方式，此处选择了随机裁剪；
        # 其中name代表是对哪个参数进行裁剪，如果对偏置进行裁剪则该参数值为'bias'；
        # amount是指裁剪比例
        prune.random_unstructured(module,name = 'weight', amount = 0.3)
        # 此时model.weight被替换为model.weight_orig和model.weight_mask 
        # 使用list(module.named_buffers())可以查看
        prune.remove(module,'weight')

2. 使用微软的nni工具进行剪枝

需要安装nni库，从nni.algorithms.compression.pytorch.pruning中选择想要的剪枝方法

教程链接：Pruning — An open source AutoML toolkit for neural architecture search, model compression and hyper-parameter tuning (NNI v2.6.1)https://nni.readthedocs.io/en/stable/Compression/pruning.html

演示demo：https://github.com/microsoft/nni/blob/70706eba4e6723b5647cbd20f02c218568bbbcf8/examples/model_compress/pruning/basic_pruners_torch.pyhttps://github.com/microsoft/nni/blob/70706eba4e6723b5647cbd20f02c218568bbbcf8/examples/model_compress/pruning/basic_pruners_torch.py

这里给一个简单的实现过程，有些打印可能冗余了，不要介意

import torch
from torch import nn
from torchvision import models
from torchvision import datasets
from torch import optim
import torchvision.transforms as transforms
from torchsummary import summary
import numpy as np
import random
import nni
from nni.algorithms.compression.pytorch.pruning import LevelPruner
from nni.algorithms.compression.pytorch.pruning import L2FilterPruner
from nni.compression.pytorch import ModelSpeedup

epochs = 10
device = torch.device("cpu")
batch_size = 64
train_loader = torch.utils.data.DataLoader(
    datasets.CIFAR10('./data', train=True, download=True,
                     transform=transforms.Compose([
                         transforms.ToTensor()
                     ])),
    batch_size=batch_size, shuffle=True)

test_loader = torch.utils.data.DataLoader(
    datasets.CIFAR10('./data', train=False,
                     transform=transforms.Compose([
                         transforms.ToTensor()
                     ])),
    batch_size=batch_size, shuffle=True)


class simple_model(nn.Module):
    def __init__(self):
        super(simple_model, self).__init__()
        self.feature = nn.Sequential(
            nn.Conv2d(3, 16, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 32, 3),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, 3),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((5, 5))
        )
        self.classifer = nn.Sequential(
            nn.Linear(64*5*5, 10),
        )

    def forward(self, x):
        x = self.feature(x)
        x = torch.flatten(x, 1)
        x = self.classifer(x)
        return x


def train(model, mode="train"):
    model = model.to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criteon = nn.CrossEntropyLoss()
    for epoch in range(epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            if mode == "train":
                logits = model(data.to(device))
                loss = criteon(logits, target.to(device))
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                if batch_idx % 1000 == 0:
                    print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        epoch, batch_idx * len(data), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader), loss.item()))

        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in test_loader:
                logits = model(data.to(device))
                test_loss += criteon(logits, target.to(device)).item()
                pred = logits.data.max(1)[1]
                target = target.cpu()
                pred = pred.cpu()
                correct += pred.eq(target.data).sum()
        test_loss /= len(test_loader.dataset)
        print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))
    if mode == "train":
        torch.save(model.state_dict(), "./model/origin_model.pth")

if __name__ == '__main__':
    model = simple_model()
    # 先训练一个模型
    train(model, mode="train")
    model.load_state_dict(torch.load("./model/origin_model.pth"))
    # 打印一下看看
    summary(model, (3, 64, 64), device="cpu")
    # 输出一下精度
    print("----------------原模型精度-----------------")
    train(model, mode="val")
    # 定义剪枝配置
    config_list = [{'sparsity': 0.8, 'op_types': ['Conv2d']}]
    # 生成剪枝后的模型以及掩膜
    # 有很多种剪枝方法，可以自己选
    pruner = L2FilterPruner(model, config_list)
    model = pruner.compress()
    pruner.export_model(model_path="./model/prune.pth", mask_path="./model/mask.pth")
    # 压缩模型
    pruner._unwrap_model()
    m_Speedup = ModelSpeedup(model, torch.randn([64, 3, 64, 64]), "./model/prune_mask.pth", "cpu")
    m_Speedup.speedup_model()
    # 打印一下模型
    summary(model, (3, 64, 64), device="cpu")
    # 打印一下模型精度
    print("---------------剪枝模型精度------------------")
    train(model, mode="val")
    # 再次训练微调模型
    train(model, mode="train")
    # 打印一下精度
    print("---------------剪枝微调精度------------------")
    train(model, mode="val")
    # 保存模型
    torch.save(model.state_dict(), "./model/prune_model.pth")

精度对比如下

----------------原模型精度-----------------

Test set: Average loss: 0.0145, Accuracy: 6798/10000 (68%)

---------------剪枝模型精度------------------

Test set: Average loss: 0.0422, Accuracy: 1114/10000 (11%)

---------------剪枝微调精度------------------

Test set: Average loss: 0.0168, Accuracy: 6225/10000 (62%)

模型对比如下

原模型大小为158k，剪枝后仅21k

 想要对自己模型进行剪枝的朋友可以参考上述代码实现简单的剪枝

 3. 参考链接

1. nni/basic_pruners_torch.py at 70706eba4e6723b5647cbd20f02c218568bbbcf8 · microsoft/nni (github.com)https://github.com/microsoft/nni/blob/70706eba4e6723b5647cbd20f02c218568bbbcf8/examples/model_compress/pruning/basic_pruners_torch.py