统计模型参数量、计算复杂度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