经验首页 前端设计 程序设计 Java相关 移动开发 数据库/运维 软件/图像 大数据/云计算 其他经验
当前位置:技术经验 » 程序设计 » Python » 查看文章
pytorch如何获得模型的计算量和参数量
来源:jb51  时间:2021/5/31 13:04:46  对本文有异议

方法1 自带

pytorch自带方法,计算模型参数总量

  1. total = sum([param.nelement() for param in model.parameters()])
  2. print("Number of parameter: %.2fM" % (total/1e6))

或者

  1. total = sum(p.numel() for p in model.parameters())
  2. print("Total params: %.2fM" % (total/1e6))

方法2 编写代码

计算模型参数总量和模型计算量

  1. def count_params(model, input_size=224):
  2. # param_sum = 0
  3. with open('models.txt', 'w') as fm:
  4. fm.write(str(model))
  5. # 计算模型的计算量
  6. calc_flops(model, input_size)
  7. # 计算模型的参数总量
  8. model_parameters = filter(lambda p: p.requires_grad, model.parameters())
  9. params = sum([np.prod(p.size()) for p in model_parameters])
  10. print('The network has {} params.'.format(params))
  11. # 计算模型的计算量
  12. def calc_flops(model, input_size):
  13. def conv_hook(self, input, output):
  14. batch_size, input_channels, input_height, input_width = input[0].size()
  15. output_channels, output_height, output_width = output[0].size()
  16. kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups) * (
  17. 2 if multiply_adds else 1)
  18. bias_ops = 1 if self.bias is not None else 0
  19. params = output_channels * (kernel_ops + bias_ops)
  20. flops = batch_size * params * output_height * output_width
  21. list_conv.append(flops)
  22. def linear_hook(self, input, output):
  23. batch_size = input[0].size(0) if input[0].dim() == 2 else 1
  24. weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
  25. bias_ops = self.bias.nelement()
  26. flops = batch_size * (weight_ops + bias_ops)
  27. list_linear.append(flops)
  28. def bn_hook(self, input, output):
  29. list_bn.append(input[0].nelement())
  30. def relu_hook(self, input, output):
  31. list_relu.append(input[0].nelement())
  32. def pooling_hook(self, input, output):
  33. batch_size, input_channels, input_height, input_width = input[0].size()
  34. output_channels, output_height, output_width = output[0].size()
  35. kernel_ops = self.kernel_size * self.kernel_size
  36. bias_ops = 0
  37. params = output_channels * (kernel_ops + bias_ops)
  38. flops = batch_size * params * output_height * output_width
  39. list_pooling.append(flops)
  40. def foo(net):
  41. childrens = list(net.children())
  42. if not childrens:
  43. if isinstance(net, torch.nn.Conv2d):
  44. net.register_forward_hook(conv_hook)
  45. if isinstance(net, torch.nn.Linear):
  46. net.register_forward_hook(linear_hook)
  47. if isinstance(net, torch.nn.BatchNorm2d):
  48. net.register_forward_hook(bn_hook)
  49. if isinstance(net, torch.nn.ReLU):
  50. net.register_forward_hook(relu_hook)
  51. if isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d):
  52. net.register_forward_hook(pooling_hook)
  53. return
  54. for c in childrens:
  55. foo(c)
  56. multiply_adds = False
  57. list_conv, list_bn, list_relu, list_linear, list_pooling = [], [], [], [], []
  58. foo(model)
  59. if '0.4.' in torch.__version__:
  60. if assets.USE_GPU:
  61. input = torch.cuda.FloatTensor(torch.rand(2, 3, input_size, input_size).cuda())
  62. else:
  63. input = torch.FloatTensor(torch.rand(2, 3, input_size, input_size))
  64. else:
  65. input = Variable(torch.rand(2, 3, input_size, input_size), requires_grad=True)
  66. _ = model(input)
  67. total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling))
  68. print(' + Number of FLOPs: %.2fM' % (total_flops / 1e6 / 2))

方法3 thop

需要安装thop

  1. pip install thop

调用方法:计算模型参数总量和模型计算量,而且会打印每一层网络的具体信息

  1. from thop import profile
  2. input = torch.randn(1, 3, 224, 224)
  3. flops, params = profile(model, inputs=(input,))
  4. print(flops)
  5. print(params)

或者

  1. from torchvision.models import resnet50
  2. from thop import profile
  3. # model = resnet50()
  4. checkpoints = '模型path'
  5. model = torch.load(checkpoints)
  6. model_name = 'yolov3 cut asff'
  7. input = torch.randn(1, 3, 224, 224)
  8. flops, params = profile(model, inputs=(input, ),verbose=True)
  9. print("%s | %.2f | %.2f" % (model_name, params / (1000 ** 2), flops / (1000 ** 3)))#这里除以1000的平方,是为了化成M的单位,

注意:输入必须是四维的

提高输出可读性, 加入一下代码。

  1. from thop import clever_format
  2. macs, params = clever_format([flops, params], "%.3f")

方法4 torchstat

  1. from torchstat import stat
  2. from torchvision.models import resnet50, resnet101, resnet152, resnext101_32x8d
  3. model = resnet50()
  4. stat(model, (3, 224, 224)) # (3,224,224)表示输入图片的尺寸

使用torchstat这个库来查看网络模型的一些信息,包括总的参数量params、MAdd、显卡内存占用量和FLOPs等。需要安装torchstat:

  1. pip install torchstat

方法5 ptflops

作用:计算模型参数总量和模型计算量

安装方法:pip install ptflops

或者

  1. pip install --upgrade git+https://github.com/sovrasov/flops-counter.pytorch.git

使用方法

  1. import torchvision.models as models
  2. import torch
  3. from ptflops import get_model_complexity_info
  4. with torch.cuda.device(0):
  5. net = models.resnet18()
  6. flops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, print_per_layer_stat=True) #不用写batch_size大小,默认batch_size=1
  7. print('Flops: ' + flops)
  8. print('Params: ' + params)

或者

  1. from torchvision.models import resnet50
  2. import torch
  3. import torchvision.models as models
  4. # import torch
  5. from ptflops import get_model_complexity_info
  6. # model = models.resnet50() #调用官方的模型,
  7. checkpoints = '自己模型的path'
  8. model = torch.load(checkpoints)
  9. model_name = 'yolov3 cut'
  10. flops, params = get_model_complexity_info(model, (3,320,320),as_strings=True,print_per_layer_stat=True)
  11. print("%s |%s |%s" % (model_name,flops,params))

注意,这里输入一定是要tuple类型,且不需要输入batch,直接输入输入通道数量与尺寸,如(3,320,320) 320为网络输入尺寸。

输出为网络模型的总参数量(单位M,即百万)与计算量(单位G,即十亿)

方法6 torchsummary

安装:pip install torchsummary

使用方法:

  1. from torchsummary import summary
  2. ...
  3. summary(your_model, input_size=(channels, H, W))

作用:

1、每一层的类型、shape 和 参数量

2、模型整体的参数量

3、模型大小,和 fp/bp 一次需要的内存大小,可以用来估计最佳 batch_size

补充:pytorch计算模型算力与参数大小

ptflops介绍

官方链接

这个脚本设计用于计算卷积神经网络中乘法-加法操作的理论数量。它还可以计算参数的数量和打印给定网络的每层计算成本。

支持layer:Conv1d/2d/3d,ConvTranspose2d,BatchNorm1d/2d/3d,激活(ReLU, PReLU, ELU, ReLU6, LeakyReLU),Linear,Upsample,Poolings (AvgPool1d/2d/3d、MaxPool1d/2d/3d、adaptive ones)

安装要求:Pytorch >= 0.4.1, torchvision >= 0.2.1

get_model_complexity_info()

get_model_complexity_info是ptflops下的一个方法,可以计算出网络的算力与模型参数大小,并且可以输出每层的算力消耗。

栗子

以输出Mobilenet_v2算力信息为例:

  1. from ptflops import get_model_complexity_info
  2. from torchvision import models
  3. net = models.mobilenet_v2()
  4. ops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, print_per_layer_stat=True, verbose=True)

在这里插入图片描述

从图中可以看到,MobileNetV2在输入图像尺寸为(3, 224, 224)的情况下将会产生3.505MB的参数,算力消耗为0.32G,同时还打印出了每个层所占用的算力,权重参数数量。当然,整个模型的算力大小与模型大小也被存到了变量ops与params中。

以上为个人经验,希望能给大家一个参考,也希望大家多多支持w3xue。

 友情链接:直通硅谷  点职佳  北美留学生论坛

本站QQ群:前端 618073944 | Java 606181507 | Python 626812652 | C/C++ 612253063 | 微信 634508462 | 苹果 692586424 | C#/.net 182808419 | PHP 305140648 | 运维 608723728

W3xue 的所有内容仅供测试,对任何法律问题及风险不承担任何责任。通过使用本站内容随之而来的风险与本站无关。
关于我们  |  意见建议  |  捐助我们  |  报错有奖  |  广告合作、友情链接(目前9元/月)请联系QQ:27243702 沸活量
皖ICP备17017327号-2 皖公网安备34020702000426号