• 网络模型的参数量和FLOPs的计算 Pytorch

    目录

    1、torchstat 

    2、thop

    3、fvcore 

    4、flops_counter

    5、自定义统计函数

    FLOPS和FLOPs的区别：

    • FLOPS：注意全大写，是floating point operations per second的缩写，意指每秒浮点运算次数，理解为计算速度。是一个衡量硬件性能的指标。
    • FLOPs：注意s小写，是floating point operations的缩写（s表复数），意指浮点运算数，理解为计算量。可以用来衡量算法/模型的复杂度。

    在介绍torchstat包和thop包之前，先总结一下：

    • torchstat包可以统计卷积神经网络和全连接神经网络的参数和计算量。
    • thop包可以统计统计卷积神经网络、全连接神经网络以及循环神经网络的参数和计算量，程序示例等详见下文。

    1、torchstat 

    pip install torchstat -i https://pypi.tuna.tsinghua.edu.cn/simple

    在实际操作中，我们可以调用torchstat包，帮助我们统计模型的parameters和FLOPs。如果不修改这个包里面的一些代码，那么这个包只适用于输入为3通道的图像的模型。

    1. import torch
    2. import torch.nn as nn
    3. from torchstat import stat
    4.  
    5.  
    6. class Simple(nn.Module):
    7.     def __init__(self):
    8.         super().__init__()
    9.         self.conv1 = nn.Conv2d(3, 16, 3, 1, padding=1, bias=False)
    10.         self.conv2 = nn.Conv2d(16, 32, 3, 1, padding=1, bias=False)
    11.  
    12.     def forward(self, x):
    13.         x = self.conv1(x)
    14.         x = self.conv2(x)
    15.         return x
    16.  
    17.  
    18. model = Simple()
    19. stat(model, (3, 244, 244))   # 统计模型的参数量和FLOPs，（3,244,244）是输入图像的size

     如果把torchstat包中的一行程序进行一点点改动，那么这个包可以用来统计全连接神经网络的参数量和计算量。当然手动计算全连接神经网络的参数量和计算量也很快 =_= 。进入torchstat源代码之后，如下图所示，注释掉圈红的地方，就可以用torchstat包统计全连接神经网络的参数量和计算量了。

    2、thop

    pip install thop -i https://pypi.tuna.tsinghua.edu.cn/simple
    1. import torch
    2. import torch.nn as nn
    3. from thop import profile
    4.  
    5. class Simple(nn.Module):
    6.     def __init__(self):
    7.         super().__init__()
    8.         self.fc1 = nn.Linear(10, 10)
    9.  
    10.     def forward(self, x):
    11.         x = self.fc1(x)
    12.         return x
    13.  
    14. net = Simple()
    15. input = torch.randn(1, 10)  # batchsize=1, 输入向量长度为10
    16. macs, params = profile(net, inputs=(input, ))
    17. print(' FLOPs: ', macs*2)   # 一般来讲，FLOPs是macs的两倍
    18. print('params: ', params)

    3、fvcore 

    pip install fvcore -i https://pypi.tuna.tsinghua.edu.cn/simple

    用它比较好

    1. import torch
    2. from torchvision.models import resnet50
    3. from fvcore.nn import FlopCountAnalysis, parameter_count_table
    4.  
    5. # 创建resnet50网络
    6. model = resnet50(num_classes=1000)
    7.  
    8. # 创建输入网络的tensor
    9. tensor = (torch.rand(1, 3, 224, 224),)
    10.  
    11. # 分析FLOPs
    12. flops = FlopCountAnalysis(model, tensor)
    13. print("FLOPs: ", flops.total())
    14.  
    15. # 分析parameters
    16. print(parameter_count_table(model))

     终端输出结果如下，FLOPs为4089184256，模型参数数量约为25.6M（这里的参数数量和我自己计算的有些出入，主要是在BN模块中，这里只计算了beta和gamma两个训练参数，没有统计moving_mean和moving_var两个参数），具体可以看下我在官方提的issue。
    通过终端打印的信息我们可以发现在计算FLOPs时并没有包含BN层，池化层还有普通的add操作（我发现计算FLOPs时并没有统一的规定，在github上看的计算FLOPs项目基本每个都不同，但计算出来的结果大同小异）。

    注意：在使用fvcore模块计算模型的flops时，遇到了问题，记录一下解决方案。首先是在jit_analysis.py的589行出错。经过调试发现，op_counts.values()的类型是int32，但是计算要求的类型只能是int、float、np.float64和np.int64，因此需要手动进行强制转换。修改如下：

    4、flops_counter

    pip install ptflops -i https://pypi.tuna.tsinghua.edu.cn/simple

    用它也很好，结果和fvcore一样

    1. from ptflops import get_model_complexity_info
    2.  
    3. macs, params = get_model_complexity_info(model, (112, 9, 9), as_strings=True,
    4.                                          print_per_layer_stat=True, verbose=True)
    5. print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
    6. print('{:<30}  {:<8}'.format('Number of parameters: ', params))

    5、自定义统计函数

    1. import torch
    2. import numpy as np
    3.  
    4. def calc_flops(model, input):
    5.     def conv_hook(self, input, output):
    6.         batch_size, input_channels, input_height, input_width = input[0].size()
    7.         output_channels, output_height, output_width = output[0].size()
    8.  
    9.         kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups) * (
    10.             2 if multiply_adds else 1)
    11.         bias_ops = 1 if self.bias is not None else 0
    12.  
    13.         params = output_channels * (kernel_ops + bias_ops)
    14.         flops = batch_size * params * output_height * output_width
    15.  
    16.         list_conv.append(flops)
    17.  
    18.     def linear_hook(self, input, output):
    19.         batch_size = input[0].size(0) if input[0].dim() == 2 else 1
    20.         num_steps = input[0].size(0)
    21.         weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
    22.         bias_ops = self.bias.nelement() if self.bias is not None else 0
    23.  
    24.         flops = batch_size * (weight_ops + bias_ops)
    25.         flops *= num_steps
    26.         list_linear.append(flops)
    27.  
    28.     def fsmn_hook(self, input, output):
    29.         batch_size = input[0].size(0) if input[0].dim() == 2 else 1
    30.  
    31.         weight_ops = self.filter.nelement() * (2 if multiply_adds else 1)
    32.         num_steps = input[0].size(0)
    33.         flops = num_steps * weight_ops
    34.         flops *= batch_size
    35.         list_fsmn.append(flops)
    36.  
    37.     def gru_cell(input_size, hidden_size, bias=True):
    38.         total_ops = 0
    39.         # r = \sigma(W_{ir} x + b_{ir} + W_{hr} h + b_{hr}) \\
    40.         # z = \sigma(W_{iz} x + b_{iz} + W_{hz} h + b_{hz}) \\
    41.         state_ops = (hidden_size + input_size) * hidden_size + hidden_size
    42.         if bias:
    43.             state_ops += hidden_size * 2
    44.         total_ops += state_ops * 2
    45.  
    46.         # n = \tanh(W_{in} x + b_{in} + r * (W_{hn} h + b_{hn})) \\
    47.         total_ops += (hidden_size + input_size) * hidden_size + hidden_size
    48.         if bias:
    49.             total_ops += hidden_size * 2
    50.         # r hadamard : r * (~)
    51.         total_ops += hidden_size
    52.  
    53.         # h' = (1 - z) * n + z * h
    54.         # hadamard hadamard add
    55.         total_ops += hidden_size * 3
    56.  
    57.         return total_ops
    58.  
    59.     def gru_hook(self, input, output):
    60.  
    61.         batch_size = input[0].size(0) if input[0].dim() == 2 else 1
    62.         if self.batch_first:
    63.             batch_size = input[0].size(0)
    64.             num_steps = input[0].size(1)
    65.         else:
    66.             batch_size = input[0].size(1)
    67.             num_steps = input[0].size(0)
    68.         total_ops = 0
    69.         bias = self.bias
    70.         input_size = self.input_size
    71.         hidden_size = self.hidden_size
    72.         num_layers = self.num_layers
    73.         total_ops = 0
    74.         total_ops += gru_cell(input_size, hidden_size, bias)
    75.         for i in range(num_layers - 1):
    76.             total_ops += gru_cell(hidden_size, hidden_size, bias)
    77.         total_ops *= batch_size
    78.         total_ops *= num_steps
    79.  
    80.         list_lstm.append(total_ops)
    81.  
    82.     def lstm_cell(input_size, hidden_size, bias):
    83.         total_ops = 0
    84.         state_ops = (input_size + hidden_size) * hidden_size + hidden_size
    85.         if bias:
    86.             state_ops += hidden_size * 2
    87.         total_ops += state_ops * 4
    88.         total_ops += hidden_size * 3
    89.         total_ops += hidden_size
    90.         return total_ops
    91.  
    92.     def lstm_hook(self, input, output):
    93.  
    94.         batch_size = input[0].size(0) if input[0].dim() == 2 else 1
    95.         if self.batch_first:
    96.             batch_size = input[0].size(0)
    97.             num_steps = input[0].size(1)
    98.         else:
    99.             batch_size = input[0].size(1)
    100.             num_steps = input[0].size(0)
    101.         total_ops = 0
    102.         bias = self.bias
    103.         input_size = self.input_size
    104.         hidden_size = self.hidden_size
    105.         num_layers = self.num_layers
    106.         total_ops = 0
    107.         total_ops += lstm_cell(input_size, hidden_size, bias)
    108.         for i in range(num_layers - 1):
    109.             total_ops += lstm_cell(hidden_size, hidden_size, bias)
    110.         total_ops *= batch_size
    111.         total_ops *= num_steps
    112.  
    113.         list_lstm.append(total_ops)
    114.  
    115.     def bn_hook(self, input, output):
    116.         list_bn.append(input[0].nelement())
    117.  
    118.     def relu_hook(self, input, output):
    119.         list_relu.append(input[0].nelement())
    120.  
    121.     def pooling_hook(self, input, output):
    122.         batch_size, input_channels, input_height, input_width = input[0].size()
    123.         output_channels, output_height, output_width = output[0].size()
    124.  
    125.         kernel_ops = self.kernel_size * self.kernel_size
    126.         bias_ops = 0
    127.         params = output_channels * (kernel_ops + bias_ops)
    128.         flops = batch_size * params * output_height * output_width
    129.  
    130.         list_pooling.append(flops)
    131.  
    132.     def foo(net):
    133.         childrens = list(net.children())
    134.         if not childrens:
    135.             print(net)
    136.             if isinstance(net, torch.nn.Conv2d) or isinstance(net, torch.nn.ConvTranspose2d):
    137.                 net.register_forward_hook(conv_hook)
    138.                 # print('conv_hook_ready')
    139.             if isinstance(net, torch.nn.Linear):
    140.                 net.register_forward_hook(linear_hook)
    141.                 # print('linear_hook_ready')
    142.             if isinstance(net, torch.nn.BatchNorm2d):
    143.                 net.register_forward_hook(bn_hook)
    144.                 # print('batch_norm_hook_ready')
    145.             if isinstance(net, torch.nn.ReLU) or isinstance(net, torch.nn.PReLU):
    146.                 net.register_forward_hook(relu_hook)
    147.                 # print('relu_hook_ready')
    148.             if isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d):
    149.                 net.register_forward_hook(pooling_hook)
    150.                 # print('pooling_hook_ready')
    151.             if isinstance(net, torch.nn.LSTM):
    152.                 net.register_forward_hook(lstm_hook)
    153.                 # print('lstm_hook_ready')
    154.             if isinstance(net, torch.nn.GRU):
    155.                 net.register_forward_hook(gru_hook)
    156.  
    157.             # if isinstance(net, FSMNZQ):
    158.             #     net.register_forward_hook(fsmn_hook)
    159.                 # print('fsmn_hook_ready')
    160.             return
    161.         for c in childrens:
    162.             foo(c)
    163.  
    164.     multiply_adds = False
    165.     list_conv, list_bn, list_relu, list_linear, list_pooling, list_lstm, list_fsmn = [], [], [], [], [], [], []
    166.     foo(model)
    167.  
    168.     _ = model(input)
    169.  
    170.     total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling) + sum(
    171.         list_lstm) + sum(list_fsmn))
    172.     fsmn_flops = (sum(list_fsmn) + sum(list_linear))
    173.     lstm_flops = sum(list_lstm)
    174.  
    175.     model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    176.     params = sum([np.prod(p.size()) for p in model_parameters])
    177.     print('The network has {} params.'.format(params))
    178.  
    179.     print(total_flops, fsmn_flops, lstm_flops)
    180.     print('  + Number of FLOPs: %.2f M' % (total_flops / 1000 ** 2))
    181.     return total_flops
    182.  
    183. if __name__ == '__main__':
    184.     from torchvision.models import resnet18
    185.  
    186.     model = resnet18(num_classes=1000)
    187.     imput_size = torch.rand((1,3,224,224))
    188.     calc_flops(model, imput_size)

  • 相关阅读:
    SSMP整合案例交互之在idea中利用vue和axios发送异步请求进行前后端调用
    集合贴——智能客服是什么
    python毕业设计项目源码选题（7）校园排课选课系统毕业设计毕设作品开题报告开题答辩PPT
    百度ERNIE 3.0——中文情感分析实战
    攻防世界心仪的公司
    论文学习记录随笔 多标签之GLOCAL
    vue中预览epub文件
    修改iframe生成的pdf的比例
    java培训技术自定义类型转换器示例
    qemu侧 网络包发送调试记录（二）
  • 原文地址：https://blog.csdn.net/qq_45100200/article/details/127728053