概念
- FLOPS:指每秒浮点运算次数,理解为计算速度,是一个衡量硬件性能的指标。
计算公式:
对卷积层:(K_h * K_w * C_in * C_out) * (H_out * W_out)
对全连接层:C_in * C_out
- FLOPs: 注意s小写,指浮点运算数,理解为计算量。可以用来衡量算法/模型的复杂度。
- Params:是指模型训练中需要训练的参数总数。
模型参数量计算公式为:
对卷积层:(K_h * K_w * C_in)* C_out
对全连接层:C_in * C_out
代码实现
方法1
使用thop库
pip install thop
import torch
from torchvision import models
from thop import profile
if __name__ == '__main__':
net = models.mobilenet_v2()
inputs = torch.randn(1, 3, 224, 224)
flops, params = profile(net, inputs=(inputs, ))
print("FLOPs=", str(flops/1e9) +'{}'.format("G"))
print("params=", str(params/1e6)+'{}'.format("M"))
方法2
使用torchstat库
pip install torchstat
import torch
from torchvision import models
from torchstat import stat
if __name__ == '__main__':
net = models.mobilenet_v2()
stat(net, (3, 224, 224))