公式来源:卷积神经网络模型参数量和运算量计算方法 - 简书 (jianshu.com)
通过遍历模型的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))
效果如下:
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))
还是一样,没有的先pip install torchstat即可,还是以restnet18为例:
from torchstat import stat
from torchvision.models import resnet18
model = resnet18()
stat(model, (3, 224, 224))
还是一样,用pip install torchsummay进行安装,通过下方例子进行调用。
from torchsummary import summary
summary(model.to(device), input_size=(1, 64, 64), batch_size=-1)