用fvcore计算PyTorch网络的参数量和FLOPs

依赖库安装命令

pip install fvcore

FLOPs： 注意s小写，是floating point operations的缩写（s表复数），意指浮点运算数，理解为计算量，可以用来衡量算法/模型的复杂度。

示例

import torch
from torchvision.models import resnet50
from fvcore.nn import FlopCountAnalysis, parameter_count_table

# 创建resnet50网络
model = resnet50(num_classes=1000)

# 创建输入网络的tensor
tensor = (torch.rand(1, 3, 224, 224),)

# 分析FLOPs，格式为10次幂
flops = FlopCountAnalysis(model, tensor)
print("FLOPs: {:.2e}".format(flops.total()))

# 分析parameters
print(parameter_count_table(model))

Reference

使用fvcore计算Pytorch中模型的参数数量以及FLOPs_太阳花的小绿豆的博客-CSDN博客