-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathconfigs.py
111 lines (76 loc) · 2.32 KB
/
configs.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
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
from dataclasses import dataclass
from typing import Optional
""" MobileViT Model Configurations & Helper Function"""
@dataclass
class Config_MobileViT_v1_S:
block_1_1_dims: int = 16
block_1_2_dims: int = 32
block_2_1_dims: int = 64
block_2_2_dims: int = 64
block_2_3_dims: int = 64
block_3_1_dims: int = 96
block_3_2_dims: int = 96
block_4_1_dims: int = 128
block_4_2_dims: int = 128
block_5_1_dims: int = 160
block_5_2_dims: int = 160
final_conv_dims: int = 640
tf_block_3_dims: int = 144
tf_block_4_dims: int = 192
tf_block_5_dims: int = 240
tf_block_3_repeats: int = 2
tf_block_4_repeats: int = 4
tf_block_5_repeats: int = 3
depthwise_expansion_factor: int = 4
@dataclass
class Config_MobileViT_v1_XS:
block_1_1_dims: int = 16
block_1_2_dims: int = 32
block_2_1_dims: int = 48
block_2_2_dims: int = 48
block_2_3_dims: int = 48
block_3_1_dims: int = 64
block_3_2_dims: int = 64
block_4_1_dims: int = 80
block_4_2_dims: int = 80
block_5_1_dims: int = 96
block_5_2_dims: int = 96
final_conv_dims: int = 384
tf_block_3_dims: int = 96
tf_block_4_dims: int = 120
tf_block_5_dims: int = 144
tf_block_3_repeats: int = 2
tf_block_4_repeats: int = 4
tf_block_5_repeats: int = 3
depthwise_expansion_factor: int = 4
@dataclass
class Config_MobileViT_v1_XXS:
block_1_1_dims: int = 16
block_1_2_dims: int = 16
block_2_1_dims: int = 24
block_2_2_dims: int = 24
block_2_3_dims: int = 24
block_3_1_dims: int = 48
block_3_2_dims: int = 48
block_4_1_dims: int = 64
block_4_2_dims: int = 64
block_5_1_dims: int = 80
block_5_2_dims: int = 80
final_conv_dims: int = 320
tf_block_3_dims: int = 64
tf_block_4_dims: int = 80
tf_block_5_dims: int = 96
tf_block_3_repeats: int = 2
tf_block_4_repeats: int = 4
tf_block_5_repeats: int = 3
depthwise_expansion_factor: int = 2
def get_mobile_vit_v1_configs(model_type: str = "S", updates: Optional[dict] = None):
if model_type == "S":
base_config = Config_MobileViT_v1_S
elif model_type == "XS":
base_config = Config_MobileViT_v1_XS
else:
base_config = Config_MobileViT_v1_XXS
if updates:
return base_config(**updates)
return base_config