统计官方模型的参数量和计算量

import torch
from torchvision import models

model_names = sorted(
    name
    for name in models.__dict__
    if name.islower()
    and not name.startswith("__")  # and "inception" in name
    and callable(models.__dict__[name])
)

device = "cpu"


for name in model_names:
    print("======================="+name+"=========================")
    model = models.__dict__[name]().to(device)
    dsize = (1, 3, 224, 224)
    size_ = (3, 224, 224)
    if "inception" in name:
        dsize = (1, 3, 299, 299)
        size_ = (3, 224, 224)

    '''
    # thop，功能单一，不推荐
    from thop.profile import profile
    inputs = torch.randn(dsize).to(device)
    total_ops, total_params = profile(model, (inputs,), verbose=False)
    print(
        "%s | %.2f | %.2f" % (name, total_params / (1000 ** 2), total_ops / (1000 ** 3))
    )
    '''

    '''
    # 直接测量参数量，不能统计计算量。
    total_params = sum(param.numel() for param in model.parameters())
    print(total_params / (1000 ** 2))
    '''

    # torchstat，打印内容最详细，但没打印层名，且支持某些层的种类有限（https://github.com/Swall0w/torchstat/blob/master/detail.md）。对于包含这些层的网络，计算结果比实际偏低。测量的是GPU模型。
    from torchstat import stat
    stat(model.to(device), size_)
    
    # torchinfo。框架比较新，还在更新中，各层的名称比较明确，但打印信息不如torchstat全面。测量的是CPU模型。
    from torchinfo import summary
    summary(model, input_size=dsize)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46