[人工智能]Pytorch 如何计算模型的【参数量 & 计算量】

1、安装 thop

注意:输入必须是四维的

# 安装 thop
pip install thop


------------------------------------------------------------------------------
# 基本用法
from nets.yolo4 import YoloBody  	# 导入模型
from thop import profile		 	# 导入thop模块

model = YoloBody(3,20)
input = torch.randn(1, 3, 420, 420)
flops, params = profile(model, inputs=(input,))

print('flops' ,flops)				# 打印计算量
print('params',params)				# 打印参数量

或者

from torchvision.models import resnet50
from thop import profile
 
# model = resnet50()
checkpoints = '模型path'
model = torch.load(checkpoints)
model_name = 'yolov3 cut asff'
input = torch.randn(1, 3, 224, 224)


flops, params = profile(model, inputs=(input, ),verbose=True)
print("%s | %.2f | %.2f" % (model, params / (1000 ** 2), flops / (1000 ** 3)))			# 转化为MB

2、torchstat

使用torchstat这个库来查看网络模型的一些信息，包括总的参数量params、MAdd、显卡内存占用量和FLOPs等。

from torchstat import stat
from torchvision.models import resnet50, resnet101, resnet152, resnext101_32x8d
 
model = resnet50()
input = torch.randn(1, 3, 224, 224)

stat(model, (3, 224, 224))  #  (3,224,224)表示输入图片的尺寸

参考

• https://www.jb51.net/article/213731.htm