Pytorch计算模型的参数量和GFLOPs

最新推荐文章于 2024-08-10 18:01:08 发布

王小凡wxf

最新推荐文章于 2024-08-10 18:01:08 发布

阅读量1.3k

点赞数 1

文章标签： pytorch 人工智能 python

本文链接：https://blog.youkuaiyun.com/qq_42242829/article/details/131601358

版权

文章介绍了如何在Pytorch中利用thop包来计算模型的参数数量和GFLOPs，强调了thop相对于torchstat的优势，特别是在处理复杂模型和自定义模型时的适用性。通过示例展示了如何导入和使用thop对ResNet18模型进行分析。

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

Pytorch计算模型的参数量和GFLOPs

Pytorch计算模型的参数量和GFLOPs

Pytorch计算模型的参数量和GFLOPs

最近在研究模型的压缩技术，使用Pytorch框架，为了计算模型的参数量和GFLOPs。

thop包

相较于torchstat，torchstat无法适用的模型某一个layer的输入为多个变量和Pytorch-0.4.1版本等情况，可以尝试使用thop工具包进行模型分析。此外，thop更加适用于自己定义的模型，在作者定义的各种模型中较为普适。

安装

pip install thop

thop使用

from thop import profile                            
from thop import clever_format
import torchvision.models as models

resnet18 = models.resnet18()

input = torch.randn(1, 3, 224, 224)                        
flops, params = profile(resnet18, inputs=(input, ))                   
print(flops,params) 
flops,params = clever_format([flops, params],"%.3f")
print(flops,params)