forked from BR-IDL/PaddleViT
-
Notifications
You must be signed in to change notification settings - Fork 0
/
stat_define.py
60 lines (50 loc) · 1.69 KB
/
stat_define.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
import os
import glob
import paddle
from config import get_config
from swin_transformer import build_swin as build_model
def count_gelu(layer, input, output):
activation_flops = 8
x = input[0]
num = x.numel()
layer.total_ops += num * activation_flops
def count_softmax(layer, input, output):
softmax_flops = 5 # max/substract, exp, sum, divide
x = input[0]
num = x.numel()
layer.total_ops += num * softmax_flops
def count_layernorm(layer, input, output):
layer_norm_flops = 5 # get mean (sum), get variance (square and sum), scale(multiply)
x = input[0]
num = x.numel()
layer.total_ops += num * layer_norm_flops
cfg = './configs/swin_tiny_patch4_window7_224.yaml'
input_size = (1, 3, 224, 224)
config = get_config(cfg)
model = build_model(config)
custom_ops = {paddle.nn.GELU: count_gelu,
paddle.nn.LayerNorm: count_layernorm,
paddle.nn.Softmax: count_softmax,
}
print(os.path.basename(cfg))
paddle.flops(model,
input_size=input_size,
custom_ops=custom_ops,
print_detail=False)
#for cfg in glob.glob('./configs/*.yaml'):
# #cfg = './configs/swin_base_patch4_window7_224.yaml'
# input_size = (1, 3, int(cfg[-8:-5]), int(cfg[-8:-5]))
# config = get_config(cfg)
# model = build_model(config)
#
#
# custom_ops = {paddle.nn.GELU: count_gelu,
# paddle.nn.LayerNorm: count_layernorm,
# paddle.nn.Softmax: count_softmax,
# }
# print(os.path.basename(cfg))
# paddle.flops(model,
# input_size=input_size,
# custom_ops=custom_ops,
# print_detail=False)
# print('-----------')