神经网络参数量和运算量的计算- 基于deepspeed库和thop库函数

引言

最近需要对神经网络的参数量和运算量进行统计。找到一个基于deepspeed库函数计算参数量和运算量的例子。而我之前一直用thop库函数来计算。

看到有一篇勘误博文写道使用thops库得到的运算量是MACs (Multiply ACcumulate operations，乘加累积操作次数），而很多其他文章提到的还是FLOPs（Floating Point Operations，浮点运算次数）。
Pytorch: 采用thop库正确计算模型计算量FLOPs和模型参数Params 【误区更正】

因此对这两种方法进行测试，来验证thop库函数得到的运算量到底是MACs还是Flops。

1 使用deepspeed库函数计算参数量和运算量

对于deepspeed库的安装就不多介绍了，对于window系统，deepspeed的最新版本可以直接通过pip下载。不需要像以前一样安装过程一把辛酸泪。(2025.2.3)

win10上安装看一下文档：

链接: windows系统安装deepspeed说明文档

以下以resnet18为例子

python 复制代码

import sys
import torch
from deepspeed.profiling.flops_profiler import get_model_profile
torch.backends.cudnn.deterministic = True
import torchvision.models as models

def main(argv):
    device = "cuda:0"
    net = models.resnet18()
    net.to(device).eval()
    width, height = 224, 224
    flops, macs, params = get_model_profile(net, (1,3,width,height))
    print("params: ", params)
    print("flops: ", flops)
    print("macs: ", macs)
if __name__ == "__main__":
    main(sys.argv)

结果如下:

打印了每一层的运算量和参数量：

最后打印的结果如下:

2 使用thop库函数计算运算量和参数量

复制代码

import torch
from thop import profile
from thop import clever_format
import torchvision.models as models

# 假设我们有一个预训练的模型
model = models.resnet18()
model.eval()

# 使用thop分析模型的运算量和参数量
input = torch.randn(1, 3, 224, 224)  # 随机生成一个输入张量，这个尺寸应该与模型输入的尺寸相匹配
MACs, params = profile(model, inputs=(input,))

# 将结果转换为更易于阅读的格式
MACs, params = clever_format([MACs, params], '%.3f')

print(f"运算量：{MACs}, 参数量：{params}")

3.结论

从以上两种方法对于ResNet-18的运算量的比较可以得知，

deepspeed库统计ResNet-18的运算量FLOPs为3.64G，MACs为1.81G。

thop库统计 ResNet-18的运算量为1.824G，这个数值上更接近deepspeed库的MACs或者是FLOPs/2。

所以 thop库获得的运算量更可能是MACs而不是Flops。