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))

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐