Pytorch中计算自己模型的FLOPs | thop.profile() 方法 | yolov5s 网络模型参数量、计算量统计

转自：Pytorch中计算自己模型的FLOPs | thop.profile() 方法 | yolov5s 网络模型参数量、计算量统计_墨理学AI-优快云博客

Pytorch: 用thop计算pytorch模型的FLOPs - 简书

安装thop

pip install thop

基础用法

• 以查看resnet50的FLOPs为例

from torchvision.models import resnet50
from thop import profile
model = resnet50()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))

• 查看自己模型的FLOPs

class YourModule(nn.Module):
    # your definition
def count_your_model(model, x, y):
    # your rule here

input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ), 
                        custom_ops={YourModule: count_your_model})

• 提升输出结果的可读性
  调用thop.clever_format

from thop import clever_format
flops, params = clever_format([flops, params], "%.3f")

参考：https://github.com/Lyken17/pytorch-OpCounter

作者：wzNote
链接：https://www.jianshu.com/p/6514b8fb1ada
来源：简书
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。