模型剪枝-使用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))

  • 2
    点赞
  • 7
    收藏
    觉得还不错? 一键收藏
  • 12
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 12
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值