统计模型参数量、计算复杂度ptflops

一、安装ptflops

https://pypi.org/project/ptflops/
pip install ptflops

二、计算模型参数量和复杂度 get_model_complexity_info

from ptflops import get_model_complexity_info
model = MyModel()
macs, params = get_model_complexity_info(model, (2,3,64,64), print_per_layer_stat=True)
print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
print('{:<30}  {:<8}'.format('Number of parameters: ', params))

下面是我测试的模型输出

Computational complexity:       1471.32 GMac
Number of parameters:           111.61 M

关于GMac是什么，我觉得下面的链接讲得比较好：
FLOPS，FLOPs，GMac

三、其他计算方法profile，stat

使用了profile来计算过，from thop import profile
出现warning, cannot find relu...
而且计算出来模型的参数量也少了很多，所以后面用了ptflops

stat

from torchstat import stat
model = MyModel()
stat(model, (3,128,128))

Total params: 39,815,151
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
Total memory: 750.35MB
Total MAdd: 168.22GMAdd
Total Flops: 91.54GFlops
Total MemR+W: 2.01GB