Pytorch剪枝api测试和结果
Pytorch 官方给出的prune接口
下面是基于prune的接口进行剪枝的方法步骤
1、首先prune接口在 torch.nn.utils.prune中，目前支持的剪枝方法有：
- RandomUnstructured
- L1Unstructured
- RandomStructured
- LnStructured
- CustomFromMask
ps：非结构性剪枝不会给剪枝后模型的速度带来提升。
2、选择一个方法，定义好一个model后，将要剪枝的模块，及模块剪枝的部分作为函数的参数传入剪枝参数
from torch.nn.utils import prune
prune.ln_structured(module, name="weight", amount=0.5, n=2, dim=0)
'''
module: 模型的模块名字，如 model.conv1、model.fc1 ，这些跟你在构建模型时有关，可以用 models.state_dict().keys() 查看
name:模块中要剪枝的部分，可以是、weight、bias
amount:指的是模型本次剪枝的概率
n：前面使用的是ln_structured 模型，n表示使用那种剪枝策略，L1、L2、L3
dim：表示对第几个维度进行剪枝，如卷积层可以是维度 0 ， 1， 2， 3
'''
3、剪枝完后会产生一个weight_mask的掩码，本身不会直接作用于模型，会产生一个weight的属性，这时候原module是不存在weight的parameter,仅仅是一个attribute
如果此时输出模型的model.state_dict().keys()
之前是 conv1.weight 变成了 conv1.weight_orig ，以及conv1.weight_mask
4、此时模型的参数仍然是没有发生变化的，需要对剪枝后的模型进行保存
prune.remove(module, 'weight')
print(list(module.named_parameters()))
5、此时模型保存的是剪枝之后的权重值，同时weight_orig已经被删除掉了
6、所以直接对每一层需要剪枝的地方选择一个剪枝方法后，直接进行剪枝就可以了，然后保存模型此时的状态参数。
对模型进行全局剪枝，prune只提供了一个全局剪枝的接口global_unstructured()
import torch.nn.utils.prune as pt_prune
pt_prune.global_unstructured(parameters_to_prune,pruning_method=pt_prune.L1Unstructured,amount=amount)
'''
parameters_to_prune:list 待剪枝模块的 名字
pruning_method：全局剪枝的方法
amount：剪枝率
'''
然后对剪枝后的模块进行remove操作即可
但是全局剪枝，只支持非结构性剪枝
prune全局非结构性剪枝测试结果
# 推理模型tiny-yolov4
def model_global_prune(amount: float):
detect_model = Darknet(
'/Users/wuzhensheng01/Documents/wzs/code/yolov4-tiny-model_pruning/cfg/yolov4-tiny.cfg') # TODO:改成相对路径
detect_model.load_weights(
"/Users/wuzhensheng01/Documents/wzs/code/yolov4-tiny-model_pruning/weight/yolov4-tiny.weights")
parameters_to_prune = list()
nums = 0
for i, modules in enumerate(detect_model.models):
if isinstance(modules, nn.Sequential):
for j, module in enumerate(modules):
if isinstance(detect_model.models[i][j], nn.Conv2d):
nums += 1
parameters_to_prune.append((detect_model.models[i][j], 'weight'))
elif isinstance(detect_model.models[i][j], nn.BatchNorm2d):
nums += 2
parameters_to_prune.append((detect_model.models[i][j], 'weight'))
parameters_to_prune.append((detect_model.models[i][j], 'bias'))
parameters_to_prune = tuple(parameters_to_prune)
assert (nums == len(parameters_to_prune))
pt_prune.global_unstructured(parameters_to_prune,
pruning_method=pt_prune.L1Unstructured,
amount=amount)
for i, modules in enumerate(detect_model.models):
if isinstance(modules, nn.Sequential):
for j, module in enumerate(modules):
if isinstance(detect_model.models[i][j], nn.Conv2d):
pt_prune.remove(detect_model.models[i][j], 'weight')
elif isinstance(detect_model.models[i][j], nn.BatchNorm2d):
pt_prune.remove(detect_model.models[i][j], 'weight')
pt_prune.remove(detect_model.models[i][j], 'bias')
return detect_model
base_line:
base model average time : 0.2082s
bicycle:0.605963
truck:0.814734
dog:0.870323
case 1: 剪枝率0.5 只剪卷积层.
model_pruned average time:0.2122
bicycle:0.597527
truck:0.825150
dog:0.592364
case2: 全局非结构性剪枝 剪枝率0.2
model_pruned average time:0.2078
bicycle:0.637542
truck:0.839107
dog:0.851859
case4:全局非结构性剪枝 剪枝率0.5 只剪bn层
‘’‘精度降为0’‘’
case4:全局非结构性剪枝 剪枝率0.2 只剪bn层
model_pruned average time : 0.2666
truck: 0.714715
truck: 0.594537
cat: 0.435578
case5:全局非结构性剪枝 剪枝率
model_pruned average time:0.2138
bicycle:0.636104
truck:0.840595
dog:0.850322
prune结构性剪枝测试结果
通过L2方法对模型的卷积层进行结构化剪枝（剪枝率0.5、0.4、0.2、0.1），剪枝完后模型的速度并没有变快，相反，模型的精度大幅度的下降，（模型精度下降的问题不知道是不是需要进行重新训练来提升，但是模型的速度并未得到提升）
结论：对于训练好的模型，prune接口只是提供了一种方法去“剪掉”模型每一层中最不重要的结构。而并没有稀疏训练这一步，导致在结构性剪枝中，模型的精度大幅度下降map趋近于0。同时剪枝方法只是使用简单的L1或L2对权重参数进行计算。
此外，接口中的“剪枝”只是找到模型中那些位置不重要参数，生成相应大小的掩膜，把不重要的位置置0，但是并没有删除与这些位置相连的前后层（只针对结构性剪枝而言），最后模型的权重大小并未发生改变，只是不重要的位置的参数大小变为了0，使得模型的速度并未提升。即使模型剪枝率达到95%，模型的速度仍与baseline保持一致。