如何用Pytorch计算模型参数量Params和计算量Flops查看框架结构等

公式

公式来源：卷积神经网络模型参数量和运算量计算方法 - 简书 (jianshu.com)

计算模型参数量

1.parameters法

通过遍历模型的parameters获取结构以及参数：

    model = 你自己的模型，eg:CNN()  ResNet()  SegNet()....
    params = list(model.parameters())
    k = 0
    for i in params:
        l = 1
        print("该层的结构：" + str(list(i.size())))
        for j in i.size():
            l *= j
        print("该层参数和：" + str(l))
        k = k + l
    print("总参数数量和：" + str(k))

效果如下： 

 

2.thop方法

thop是一个库，可以通过pip install thop进行安装，安装后通过里面的profile可以获取参数量params以及计算量flops

from thop import profile
model = 自己的模型实例化
flops, params = profile(model, inputs=传入模型的输入shape，这个必须要填)

 以resnet18进行测试：

import torch
from torchvision.models import resnet18
from thop import profile
model = resnet18()
input = torch.randn(1, 3, 128, 128)
flops, params = profile(model, inputs=(input, ))
print('flops:{}'.format(flops))
print('params:{}'.format(params))

 3.torchstat法

还是一样，没有的先pip install torchstat即可，还是以restnet18为例：

from torchstat import stat
from torchvision.models import resnet18

model = resnet18()
stat(model, (3, 224, 224))

 

4.torchsummary法

还是一样，用pip install torchsummay进行安装，通过下方例子进行调用。

from torchsummary import summary
summary(model.to(device), input_size=(1, 64, 64), batch_size=-1)