东方耀AI技术分享

标题: htop和thop傻傻分不清楚 [打印本页]

作者: 东方耀 时间: 2020-10-31 09:42
标题: htop和thop傻傻分不清楚

htop和thop傻傻分不清楚

Pytorch: 用thop计算pytorch模型的FLOPs
pip3 install thop

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

作者: bixintao 时间: 2020-11-10 19:50
积分积分积分

欢迎光临东方耀AI技术分享 (http://www.ai111.vip/)