flops-counter.pytorch

项目介绍

flops-counter.pytorch 是一个用于计算神经网络在 Pytorch 框架中的理论乘加运算量的工具。这款工具可以帮助用户计算网络的参数数量，并打印每一层的计算成本细节。对于神经网络模型的优化和评估，该工具提供了重要的参考指标。

功能与特点

后端支持

ptflops 提供了两个后端：pytorch 和 aten。两个后端之间的区别主要在于支持的模型架构和输出的详细程度。

pytorch 后端：这是一个遗留后端，主要支持 nn.Modules，在卷积神经网络（CNN）分析上表现较佳。然而，对于类似 transformer 等非 CNN 架构，建议使用 aten 后端。
aten 后端：这是默认后端，考虑到了更多操作类型，覆盖了更多的模型架构，包括 transformer。

`aten` 后端功能

在 aten 后端中，重点考虑了以下操作：

矩阵乘法及相关变体（如 aten.mm, aten.matmul, aten.addmm, aten.bmm）
卷积操作（如 aten.convolution）

使用建议：

若使用 verbose=True，工具可以显示在复杂度计算过程中被忽略的操作。
aten 后端仅打印直接嵌套在根 nn.Module 中的模块的统计信息，第二层嵌套模块不会显示在输入层统计中。
使用 ignore_modules 参数，可以让 ptflops 忽略指定的模块，便于进行特定的研究需求。例如，要在计算中忽略所有卷积，可以指定 ignore_modules=[torch.ops.aten.convolution, torch.ops.aten._convolution]。

`pytorch` 后端功能

这个后端支持多种层类型，包括但不限于：

卷积层 (Conv1d/2d/3d)
批归一化，分组归一化，实例归一化，层归一化等操作
激活函数 (如 ReLU, ELU, GELU)
线性层及池化层

使用建议：

部分 torch.nn.functional.* 和 tensor.* 操作不在支持列表中，因此它们不会对最终复杂度估算有贡献。
在复杂模型中，可以使用 input_constructor 参数来构建非平凡输入，通过传入函数来处理输入的空间分辨率。
verbose 参数可以提供关于不影响最终数据的模块的信息。
使用 ignore_modules 参数可以忽略特定模块，以适应研究需求。

使用方式

用户可以通过 PyPI 和项目存储库安装 ptflops。对于 PyTorch 2.0 及以上版本，可以使用以下命令进行安装：

pip install ptflops

示例

以下是一个简单的使用示例：

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, backend='pytorch',
                                           print_per_layer_stat=True, verbose=True)
  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
  print('{:<30}  {:<8}'.format('Number of parameters: ', params))

通过这个示例，用户可以看到每层计算的复杂度以及参数数量，为进一步的模型优化提供数据参考。

评价与引用

如果 ptflops 在您的论文或技术报告中发挥了作用，请引用以下信息：

@online{ptflops,
  author = {Vladislav Sovrasov},
  title = {ptflops: a flops counting tool for neural networks in pytorch framework},
  year = 2018-2024,
}

该工具得到了社区的支持与贡献，感谢 @warmspringwinds 和 Horace He 提供的最初版本脚本。