计算深度学习模型的参数量、复杂度

目录

背景：

依赖包介绍：

示例：

参考：

---

背景：

计算深度学习模型参数量有现成的脚本可供使用，如下。简单的使用方法参考了代码【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