模型剪枝-使用NNI框架进行推理测试

        剪枝是一种常用的压缩神经网络模型的技术。剪枝方法探索模型权重（参数）中的冗余，并尝试删除/修剪冗余和不重要的权重。 从模型中删除冗余元素，将它们的值归零，确保它们不参与反向传播过程。 消除了过度参数化神经网络中的冗余，在几乎不损害性能的情况下实现了存储和计算节省。

根据移除组件的粒度，剪枝可以分为两类：

    （1）非结构化剪枝。基于启发式方法（如权重大小排序），将不重要的参数归零。 由于不规则的稀疏性而难以加速。

    （2）结构化剪枝。 丢弃子结构的层/通道数，基于优化的重要性分数的方法。 减少神经元/通道的数量，以实现经验鲁棒性和网络加速。

框架介绍：NNI（Neural Network Intelligence）。用于超参数优化、神经架构搜索、模型压缩和特征工程的开源AutoML工具包。

使用教程：

1.安装：pip install nni。同时，需要安装torch、timm、tqdm和thop库。

2.剪枝模型。

3.推理模型测试。

模型剪枝代码，main.py。

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm as tm
from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import L1NormPruner
from nni.compression.pytorch.speedup import ModelSpeedup

# 定义待裁剪模型，在这里使用预训练模型
model = tm.create_model("efficientnet_b4", pretrained=True, num_classes = 2)

# 加载模型参数
pretrained = torch.load("./checkpoint/epoch12_acc0.9667.pth")
pretrained_state_dict = pretrained['model_state_dict']
model_state_dict = model.state_dict()
for key in pretrained_state_dict:model_state_dict[key] = pretrained_state_dict[key]
model.load_state_dict(model_state_dict, strict = False)  
print(model)
torch.save(model.state_dict(), "./checkpoint/baseline.pth")

# 模型剪枝
# 定义配置文件，sparsity表示稀疏度，0~1之间，数值越大表明剪枝的越多，模型越精简；
config_list = [{
    'sparsity': 0.3,
    'op_types': ['Conv2d']
}
# eff-b4中blocks.0.1.conv_dw层剪枝有问题，当前跳过改成
, {
    'exclude': True,
    'op_names': ['blocks.0.1.conv_dw']
}
]

# L1NormPruner
pruner = L1NormPruner(model, config_list)
# 基于剪枝规则获得待剪枝的通道，以masks定义；
_, masks = pruner.compress()
pruner._unwrap_model()
# 网络重构-实现加速的关键
ModelSpeedup(model, dummy_input = torch.ones((1, 3, 512, 512)), masks_file=masks).speedup_model()
print(model)
# 只保存整个网络，无需保存具体参数
torch.save(model, "./checkpoint/model_pruner.pth")

剪枝模型与原始模型推理测试，infer.py。

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm as tm
import os
import numpy as np
import tqdm
from thop import profile
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
USE_GPU = True
input_image = [torch.ones(B, 3, 512, 512) for B in [1, 4, 8, 16]]

# 无剪枝模型与剪枝模型
model_baseline = tm.create_model("efficientnet_b4", pretrained=True, num_classes = 2)
model_pruner = torch.load("./checkpoint/model_pruner.pth")

if USE_GPU:
    input_image = [_image.cuda() for _image in input_image]
    model_baseline = model_baseline.cuda()
    model_pruner = model_pruner.cuda()

# 预热, GPU 平时可能为了节能而处于休眠状态, 因此需要预热
print('warm up ...\n')
with torch.no_grad():
    for _ in range(10): _ = model_baseline(input_image[0])    
torch.cuda.synchronize()

# 设置用于测量时间的 cuda Event, 这是PyTorch 官方推荐的接口,理论上应该最靠谱
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
# 初始化一个时间容器
repetitions = 300
timings = np.zeros((repetitions, 1))

model_name = ["model_baseline", "model_pruner"]
for index, model in enumerate([model_baseline, model_pruner]):
    model.eval()
    print('testing ', model_name[index])
    # input_image ：遍历不同输入大小[x, 3, 512, 512], x=1,4,8,16 送入网络测试时间
    for _input_image in input_image:
        print("Batch size:", _input_image.size(0))
        with torch.no_grad():
            for rep in tqdm.tqdm(range(repetitions)):
                starter.record()
                _ = model(_input_image)
                ender.record()
                torch.cuda.synchronize()
                curr_time = starter.elapsed_time(ender) # 从 starter 到 ender 之间用时,单位为毫秒
                timings[rep] = curr_time
        avg = timings.sum()/repetitions
        print('\navg={}  ms\n'.format(avg))