如何在不影响0个权重神经元的情况下计算FLOP和参数？

问题描述

我的Prune代码如下所示，运行此代码后，我将得到一个名为“ pruned_model.pth”的文件。

import torch
from torch import nn
import torch.nn.utils.prune as prune
import torch.nn.functional as F
from cnn import net

ori_model = '/content/drive/My Drive/ECG_weight_prune/checkpoint_dir/model.pth'
save_path = '/content/drive/My Drive/ECG_weight_prune/checkpoint_dir/pruned_model.pth'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = net().to(device)
model.load_state_dict(torch.load(ori_model))  

module = model.conv1
print(list(module.named_parameters()))
print(list(module.named_buffers()))

prune.l1_unstructured(module,name="weight",amount=0.3)
prune.l1_unstructured(module,name="bias",amount=3)
print(list(module.named_parameters()))
print(list(module.named_buffers()))
print(module.bias)
print(module.weight)
print(module._forward_pre_hooks)
prune.remove(module,'weight')
prune.remove(module,'bias')
print(list(module.named_parameters()))
print(model.state_dict())
torch.save(model.state_dict(),save_path)

结果是：

[('weight',Parameter containing:
tensor([[[-0.0000,-0.3137,-0.3221,...,0.5055,0.3614,-0.0000]],[[ 0.8889,0.2697,-0.3400,0.8546,0.2311,[[-0.2649,-0.1566,-0.0000,0.0000,0.3855]],[[-0.2836,0.2155,-0.8894,-0.7676,-0.6271]],[[-0.7908,-0.6732,-0.5024,0.2011,0.4627,1.0227]],[[ 0.4433,0.5048,0.7685,-1.0530,-0.8908,-0.4799]]],device='cuda:0',requires_grad=True)),('bias',Parameter containing:
tensor([-0.7497,-1.3594,-1.7613,-2.0137,-1.1763,0.4150,-1.6996,-1.5354,0.4330,-0.9259,0.4156,-2.3099,-0.4282,-0.5199,0.1188,-1.1725,-0.9064,-1.6639,-1.5834,-0.3655,-2.0727,-2.1078,-1.6431,-0.0694,-0.5435,-1.9623,0.5481,-0.8255,-1.5108,-0.4029,-1.9759,0.0522,0.0599,-2.2469,-0.5599,0.1039,-0.4472,-1.1706,-0.0398,-1.9441,-1.5310,-0.0837,-1.3250,-0.2098,-0.1919,0.4600,-0.8268,-1.0041,-0.8168,-0.8701,0.3869,0.1706,-0.0226,-1.2711,-0.9302,-2.0696,-1.1838,0.4497,-1.1426,0.0772,-2.4356,-0.3138,0.6297,0.2022,-0.4024,-1.2337,0.2840,0.4515,0.2999,0.0273,0.0374,0.1325,-0.4890,-2.3845,-1.9663,0.2108,-0.1144,0.0544,-0.2629,0.0393,-0.6728,-0.9645,0.3118,-0.5142,-0.4097,-1.5142,-1.2798,0.2871,-2.0122,-0.9346,-0.4931,-1.4895,-1.1401,-0.8823,0.2210,0.4282,0.1685,-1.8876,-0.7459,0.2505,-0.6315,0.3827,-0.3348,0.1862,0.0806,-2.0277,0.2068,0.3281,-1.8045,-2.2377,-1.9742,-0.5164,-0.0660,0.8392,0.5863,-0.7301,0.0778,0.1611,0.0260,0.3183,-0.9097,-1.6152,0.4712,-0.2378,-0.4972],requires_grad=True))]

存在许多零权重。如何计算FLOP和参数而不计算与这些零值相关的计算？

我使用以下代码来计算FLOP和参数。

import torch
from cnn import net
from ptflops import get_model_complexity_info

ori_model = '/content/drive/My Drive/ECG_weight_prune/checkpoint_dir/model.pth'
pthfile = '/content/drive/My Drive/ECG_weight_prune/checkpoint_dir/pruned_model.pth'

model = net()
# model.load_state_dict(torch.load(ori_model))  
model.load_state_dict(torch.load(pthfile))  
# print(model.state_dict())

macs,params = get_model_complexity_info(model,(1,260),as_strings=False,print_per_layer_stat=True,verbose=True)
print('{:<30}  {:<8}'.format('Computational complexity: ',macs))
print('{:<30}  {:<8}'.format('Number of parameters: ',params))

两个ori_model nad pthfile的输出相同，如下所示。

Warning: module Dropout2d is treated as a zero-op.
Warning: module Flatten is treated as a zero-op.
Warning: module net is treated as a zero-op.
net(
  0.05 M,100.000% Params,0.001 GMac,100.000% MACs,(conv1): Conv1d(0.007 M,13.143% Params,0.0 GMac,45.733% MACs,1,128,kernel_size=(50,),stride=(3,))
  (conv2): Conv1d(0.029 M,57.791% Params,50.980% MACs,32,kernel_size=(7,stride=(1,))
  (conv3): Conv1d(0.009 M,18.619% Params,0.913% MACs,kernel_size=(9,))
  (fc1): Linear(0.004 M,8.504% Params,0.404% MACs,in_features=32,out_features=128,bias=True)
  (fc2): Linear(0.001 M,1.299% Params,0.063% MACs,in_features=128,out_features=5,bias=True)
  (bn1): BatchNorm1d(0.0 M,0.515% Params,1.793% MACs,eps=1e-05,momentum=0.1,affine=True,track_running_stats=True)
  (bn2): BatchNorm1d(0.0 M,0.129% Params,0.114% MACs,track_running_stats=True)
  (dropout): Dropout2d(0.0 M,0.000% Params,0.000% MACs,p=0.5,inplace=False)
  (faltten): Flatten(0.0 M,)
)
Computational complexity:       1013472.0
Number of parameters:           49669 

解决方法

您可以做的一件事是从FLOPs计算中排除低于特定阈值的权重。为此，您必须修改触发器计数器功能。

我将在下面提供对fc和conv层进行修改的示例。

def linear_flops_counter_hook(module,input,output):
    input = input[0]
    output_last_dim = output.shape[-1]  # pytorch checks dimensions,so here we don't care much
    # MODIFICATION HAPPENS HERE
    num_zero_weights = (module.weight.data.abs() < 1e-9).sum()
    zero_weights_factor = 1 - torch.true_divide(num_zero_weights,module.weight.data.numel())
    module.__flops__ += int(np.prod(input.shape) * output_last_dim) * zero_weights_factor.numpy()
    # MODIFICATION HAPPENS HERE

def conv_flops_counter_hook(conv_module,output):
    # Can have multiple inputs,getting the first one
    input = input[0]

    batch_size = input.shape[0]
    output_dims = list(output.shape[2:])

    kernel_dims = list(conv_module.kernel_size)
    in_channels = conv_module.in_channels
    out_channels = conv_module.out_channels
    groups = conv_module.groups

    filters_per_channel = out_channels // groups
    conv_per_position_flops = int(np.prod(kernel_dims)) * in_channels * filters_per_channel

    active_elements_count = batch_size * int(np.prod(output_dims))

    # MODIFICATION HAPPENS HERE
    num_zero_weights = (conv_module.weight.data.abs() < 1e-9).sum()
    zero_weights_factor = 1 - torch.true_divide(num_zero_weights,conv_module.weight.data.numel())
    overall_conv_flops = conv_per_position_flops * active_elements_count * zero_weights_factor.numpy()
    # MODIFICATION HAPPENS HERE
    
    bias_flops = 0

    if conv_module.bias is not None:

        bias_flops = out_channels * active_elements_count

    overall_flops = overall_conv_flops + bias_flops

    conv_module.__flops__ += int(overall_flops)

请注意，我使用1e-9作为权重计数为零的阈值。