Pytorch统计网络参数计算工具、模型 FLOPs, MACs, MAdds 关系

最新推荐文章于 2025-03-26 10:03:59 发布

李代数

最新推荐文章于 2025-03-26 10:03:59 发布

阅读量6.7k

点赞数 5

文章标签： pytorch tensorflow 深度学习

本文链接：https://blog.youkuaiyun.com/user_lib/article/details/123452572

版权

本文介绍了如何在PyTorch中统计神经网络的参数数量，包括总参数数和训练参数数，并提供了相关工具来计算FLOPs、MACs和MAdds。通过`get_parameter_number`函数可以获取模型参数详情，使用torchstat或torchsummary库可以进一步查看模型的FLOPs、显存占用等信息，帮助理解模型复杂度。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Pytorch统计网络参数

#网络参数数量
def get_parameter_number(net):
    total_num = sum(p.numel() for p in net.parameters())
    trainable_num = sum(p.numel() for p in net.parameters() if p.requires_grad)
    return {'Total': total_num, 'Trainable': trainable_num}

#查看网络参数
print(model.state_dict())

FLOPs, MACs, MAdds 关系

在这里插入图片描述
见文章：CNN模型复杂度（FLOPs、MAC）、参数量与运行速度

计算工具：

地址	备注
https://github.com/Lyken17/pytorch-OpCounter	Pytorch
https://github.com/sovrasov/flops-counter.pytorch	Pytorch
https://stackoverflow.com/questions/45085938/tensorflow-is-there-a-way-to-measure-flops-for-a-model	TensorFlow: 自带tf.RunMetadata()

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

!pip install torchstat
from torchstat import stat
from torchvision.models import resnet50, resnet101, resnet152, resnext101_32x8d
 
model = resnet50()
# stat打印完整信息
stat(model, (3, 224, 224))
# 模型的总参数量
total = sum([param.nelement() for param in model.parameters()])
print("Number of parameters: %.2fM" % (total/1e6))

也可以使用torchsummary

!pip install torchsummary
from torchsummary import summary
summary(model, input_size=(ch, h, w), batch_size=-1)
#ch是指输入张量的channel数量，h表示输入张量的高，w表示输入张量的宽。