参考文章
神经网络模型可视化工具、参数量、Flops 统计工具
使用torchsummary打印torch每层形状
AttributeError: ‘list’ object has no attribute 'size’解决方案
方式一:使用torchsummary
import torch
from torchsummary import summary
from pytorch_model_attention import densenet121
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
densenet = densenet121().to(device)
summary(densenet, (3, 224, 224))
方式二:使用ptflops工具
import torch
from ptflops import get_model_complexity_info
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = densenet121().to(device)
###参数计算
flops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, print_per_layer_stat=True)
print('Flops: ' + flops)
print('Params: ' + params)