Pytorch查看模型参数量和计算量

1、首先声明一个网络

使用torchsummary可以查看模型的参数，和输入输出尺寸，但不能看FLOPs。

import torchvision.models
import torch
import torchsummary

model = torchvision.models.vgg16(pretrained=False)
device = torch.device('cpu')
model.to(device)

torchsummary.summary(model.cuda(), (3, 224, 224))

 2、pip install thop

可以显示总的FLOPs和参数量params，但不能显示每层的结构，有一些警告，但是没关系，因为relu是不参与计算参数量和FLOPs的。

import torchvision.models
import torch
from thop import profile
from thop import clever_format
import torchsummary
model = torchvision.models.vgg16(pretrained=False)
device = torch.device('cpu')
model.to(device)
myinput = torch.zeros((1, 3, 224, 224)).to(device)
flops, params = profile(model.to(device), inputs=(myinput,))
flops, params = clever_format([flops, params], "%.3f")
print(flops, params)

3、pip install torchstat

可以显示总的FLOPs和参数量params，也可能显示每层的结构。

import torchvision.models
import torch
from torchstat import stat
model = torchvision.models.vgg16(pretrained=False)
device = torch.device('cpu')
model.to(device)
stat(model.to(device), (3, 224, 224))