pytorch实现FLOPs和Params的计算

概念

• 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))