-
Notifications
You must be signed in to change notification settings - Fork 28
/
get_flops.py
89 lines (73 loc) · 2.95 KB
/
get_flops.py
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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
# ------------------------------------------------------------
# Copyright (c) VCU, Nanjing University.
# Licensed under the Apache License 2.0 [see LICENSE for details]
# Written by Qing-Long Zhang
# ------------------------------------------------------------
import argparse
import torch
from timm.models import create_model
import models
try:
from mmcv.cnn import get_model_complexity_info
from mmcv.cnn.utils.flops_counter import get_model_complexity_info, flops_to_string, params_to_string
except ImportError:
raise ImportError('Please upgrade mmcv to >0.6.2')
def parse_args():
parser = argparse.ArgumentParser(description='Get FLOPS of a classification model')
parser.add_argument('--model', default='restv2_tiny', type=str, metavar='MODEL',
help='train config file path')
parser.add_argument(
'--shape',
type=int,
nargs='+',
default=[224, 224],
help='input image size')
args = parser.parse_args()
return args
def attn_flops(h, w, r, dim):
return 2 * h * w * (h // r) * (w // r) * dim
def get_flops(model, input_shape):
flops, params = get_model_complexity_info(model, input_shape, as_strings=False)
if 'rest' in model.name:
_, H, W = input_shape
# calculate flops of ResTv2
stage1 = attn_flops(H // 4, W // 4,
model.stage1[0].attn.sr_ratio,
model.stage1[0].attn.dim) * model.depths[0]
stage2 = attn_flops(H // 8, W // 8,
model.stage2[0].attn.sr_ratio,
model.stage2[0].attn.dim) * model.depths[1]
stage3 = attn_flops(H // 16, W // 16,
model.stage3[0].attn.sr_ratio,
model.stage3[0].attn.dim) * model.depths[2]
stage4 = attn_flops(H // 32, W // 32,
model.stage4[0].attn.sr_ratio,
model.stage4[0].attn.dim) * model.depths[3]
flops += stage1 + stage2 + stage3 + stage4
return flops_to_string(flops), params_to_string(params)
def main():
args = parse_args()
if len(args.shape) == 1:
input_shape = (3, args.shape[0], args.shape[0])
elif len(args.shape) == 2:
input_shape = (3,) + tuple(args.shape)
else:
raise ValueError('invalid input shape')
model = create_model(
args.model,
pretrained=False,
num_classes=1000
)
model.name = args.model
if torch.cuda.is_available():
model.cuda()
model.eval()
flops, params = get_flops(model, input_shape)
split_line = '=' * 30
print(f'{split_line}\nInput shape: {input_shape}\n'
f'Flops: {flops}\nParams: {params}\n{split_line}')
print('!!!Please be cautious if you use the results in papers. '
'You may need to check if all ops are supported and verify that the '
'flops computation is correct.')
if __name__ == '__main__':
main()