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

你可能感兴趣的:(深度学习,pytorch,python,人工智能)