计算深度学习模型的参数量、复杂度(MACS、params)

目录

背景:

依赖包介绍:

示例:

参考:


背景:

计算深度学习模型参数量有现成的脚本可供使用,如下。简单的使用方法参考了代码【2】

依赖包介绍:

pytorch框架中卷积网络的flops计数器:ptflops

安装:

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

##or
#pip install ptflops

说明【1】:

此脚本用于计算乘法加法运算的理论量 在卷积神经网络中。它还可以计算参数的数量和 给定网络的每层打印计算成本。

支持的层:

  • conv1d/2d/3d(包括分组)
  • convTranspasse2d(包括分组)
  • 批次标准1d/2d/3d
  • 激活(relu、prelu、elu、relu6、leakyrelu)
  • 线性
  • 向上采样
  • 池(avgpool1d/2d/3d、maxpool1d/2d/3d和自适应池)

要求:Pythorch>;=0.4.1,TorchVision>;=0.2.1

示例:

1【3】.

import torchvision.models as models
import torch
from ptflops import get_model_complexity_info

with torch.cuda.device(0):
  net = models.densenet161()
  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True,
                                           print_per_layer_stat=True, verbose=True)
  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
  print('{:<30}  {:<8}'.format('Number of parameters: ', params))

2.【2】

    net = model()

    inp_shape = (6, 64, 64) #输入的分辨率

    from ptflops import get_model_complexity_info
    FLOPS = 0
    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=True)

    # params = float(params[:-4])
    print(params)
    macs = float(macs[:-4]) + FLOPS / 10 ** 9

    print('mac', macs, params)

参考:

【1】Python ptflops包_程序模块 - PyPI - Python中文网

【2】NAFNet/NAFSSR_arch.py at main · megvii-research/NAFNet · GitHub

【3】GitHub - sovrasov/flops-counter.pytorch: Flops counter for convolutional networks in pytorch framework

你可能感兴趣的:(python,深度学习,python,torch,参数量)