python-基础-计算模型参数总量

计算网络模型参数总量

方法1：pytorch自带方法，计算模型参数总量

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

例子：

红色断点放到下一行，并用debug运行。结果为：

方法2：用工具包thop

要安装：pip install thop
在程序里输入：
from thop import profile
from thop import clever_format

input1 = torch.randn(1,1, 16,16)
input2 = torch.randn(1,1, 16,16)
input3 = torch.randn(1,1, 16,16)

flops, params = profile(model, inputs=(input1,input2,input2 ))
print(flops, params) # 1819066368.0 11689512.0
flops, params = clever_format([flops, params], “%.3f”)
print(flops, params) # 1.819G 11.690M

例如：

红色断点放到下一行，并用debug运行。结果为：

3.642M和 7.265K 分别代表FLOPs和参数量。

方法3：用torchstat

和thop一样，要安装包：pip install torchstat
但torchstat貌似只能测试一个输入的，像thop一样多输入的不行。torchstat 详细展示网络各层的参数，计算量，内存读写等等。
在程序里输入：
from torchstat import stat
stat(model, (1, 16, 16))

例如：

红色断点放到下一行，并用debug运行。结果为：

备注：例子里面(16,16)指的是图片的宽和长。

小结

以上三种方式，计算的结果会有些许差异，但是不会差太多。
方法1不需要另外装包。但只能计算模型的参数量。方法2和方法3需要额外安装包，不过这个包不难装，我用pip 清华源一次装好。
方法2可以适合单输入和多输入的模型。
方法3貌似只适合但输入的模型。方法3有层间的分析。