Pytorch统计模型参数量和计算量

需要使用官方库THOP: PyTorch-OpCounterGitHub 地址：https://github.com/Lyken17/pytorch-OpCounter1.安装THOP：pip install thop# orpip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git2.使用方...

文章共258字 · 阅读需要大约1分钟

一键AI生成摘要，助你高效阅读

问答

chenzy_hust

8793人浏览 · 2020-01-13 15:08:31

chenzy_hust · 2020-01-13 15:08:31 发布

需要使用官方库THOP: PyTorch-OpCounter
GitHub 地址：https://github.com/Lyken17/pytorch-OpCounter

1.安装THOP：

pip install thop 
# or
pip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git

2.使用方法：

from torchvision.models import resnet50
from thop import profile
# 增加可读性
from thop import clever_format

# 可替换为自己的模型及输入
model = resnet50()
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ))
flops, params = clever_format([flops, params], "%.3f")