pytorch网络m参数量、flops计算方法

最新推荐文章于 2024-10-16 12:14:18 发布

cv_lz

最新推荐文章于 2024-10-16 12:14:18 发布

阅读量1.3k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： code 文章标签： pytorch python 深度学习

本文链接：https://blog.youkuaiyun.com/qq_42250789/article/details/120070976

code 专栏收录该内容

6 篇文章

订阅专栏

该篇博客主要介绍了如何使用thop库和ptflops库来计算PyTorch模型的FLOPs（浮点运算次数）和参数量。通过示例代码展示了针对模型G的计算过程，打印了模型的计算量和参数量，为模型复杂性的评估提供了参考。

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

# 1
from thop import profile

x = torch.randn(1, 3, 256, 256)
flops, params = profile(self.modelG, inputs=(x,))
print('flops is %.2fM' % (flops/1e6))  ## 打印计算量
print('params is %.2fM' % (params/1e6))  ## 打印参数量

# 2
num = 0
for param in self.modelG.parameters():
    if param.requires_grad:
        num += param.numel()
print('param is %.2fM' % (num / 1e6))

# 3
total = sum([param.nelement() for param in self.modelG.parameters()])
print("Number of parameter: %.2fM" % (total / 1e6))

# 4
from ptflops import get_model_complexity_info

flops, params = get_model_complexity_info(self.modelG, (3, 256, 256), as_strings=True,
                                          print_per_layer_stat=False)  # 不用写batch_size大小，默认batch_size=1
print('Flops:  ' + flops)
print('Params: ' + params)
exit(-1)