转自: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
来源:简书
著作权归作者所有。商业转载请联系作者获得授权,非商业转载请注明出处。