剪枝是一种常用的压缩神经网络模型的技术。剪枝方法探索模型权重(参数)中的冗余,并尝试删除/修剪冗余和不重要的权重。
从模型中删除冗余元素,将它们的值归零,确保它们不参与反向传播过程。
消除了过度参数化神经网络中的冗余,在几乎不损害性能的情况下实现了存储和计算节省。
根据移除组件的粒度,剪枝可以分为两类:
(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))