THOP: 估算PyTorch模型的MACs 模块

最新推荐文章于 2024-08-08 08:10:48 发布

hajungong007

最新推荐文章于 2024-08-08 08:10:48 发布

阅读量1.5w

点赞数 7

分类专栏： pytorch

本文链接：https://blog.csdn.net/hajungong007/article/details/90169808

版权

pytorch 专栏收录该内容

11 篇文章 1 订阅

订阅专栏

THOP: PyTorch-OpCounter

How to install

Through PyPi

pip install thop
Using GitHub (always latest)

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

How to use

Basic usage

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

Define the rule for 3rd party module.

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

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

Improve the output readability

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

## Results on Recent Models
Model | Params(M) |MACs(G)
---|---|---
alexnet | 61.10 | 0.71
vgg11 | 132.86 | 7.75
vgg11_bn | 132.87 | 7.76
vgg13 | 133.05 | 11.46
vgg13_bn | 133.05 | 11.48
vgg16 | 138.36 | 15.62
vgg16_bn | 138.37 | 15.65
vgg19 | 143.67 | 19.79
vgg19_bn | 143.68 | 19.82
densenet121 | 7.98 | 2.79
densenet161 | 28.68 | 7.69
densenet169 | 14.15 | 3.33
densenet201 | 20.01 | 4.28
resnet18 | 11.69 | 1.58
resnet34 | 21.80 | 3.44
resnet50 | 25.56 | 3.53
resnet101 | 44.55 | 7.26
resnet152 | 60.19 | 10.99
squeezenet1_0 | 1.25 | 0.70
squeezenet1_1 | 1.24 | 0.34