pytorch统计模型参数量

时间:2023-03-09 17:16:16
pytorch统计模型参数量

用resnet50 来举例子

from torchvision.models import resnet18, resnet34, resnet50
model = resnet50()
print("resnet50 have {} paramerters in total".format(sum(x.numel() for x in model.parameters())))

其中numel表示含有多少element,通过此操作可以统计模型的参数量有多少

另外,两个是一样的,方便debug看其中到底有什么东西

    num_params = 0
for param in model.parameters():
num_params += param.numel()
print(num_params / 1e6)