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)