google-research

1
# coding=utf-8
2
# Copyright 2024 The Google Research Authors.
3
#
4
# Licensed under the Apache License, Version 2.0 (the "License");
5
# you may not use this file except in compliance with the License.
6
# You may obtain a copy of the License at
7
#
8
#     http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""Sparse/dense MobileNetV1 model configurations."""
17

18

19
class Config(object):
20
  """Config base class."""
21

22
  def __init__(self, fuse_bnbr=False):
23
    self.num_blocks = 13
24
    self.width = None
25
    self.block_nonzeros = None
26
    self.block_config = ["dense"] * self.num_blocks
27
    self.fuse_bnbr = fuse_bnbr
28
    self.num_classes = 1000
29

30

31
class W14(Config):
32
  """Width 1.4."""
33

34
  def __init__(self, fuse_bnbr=False):
35
    super(W14, self).__init__(fuse_bnbr)
36
    self.width = 1.4
37

38

39
class W12(Config):
40
  """Width 1.2."""
41

42
  def __init__(self, fuse_bnbr=False):
43
    super(W12, self).__init__(fuse_bnbr)
44
    self.width = 1.2
45

46

47
class W10(Config):
48
  """Width 1.0."""
49

50
  def __init__(self, fuse_bnbr=False):
51
    super(W10, self).__init__(fuse_bnbr)
52
    self.width = 1.0
53

54

55
class W18S90D0(Config):
56
  """Width 1.8, sparsity 90%, dense first block."""
57

58
  def __init__(self, fuse_bnbr=False):
59
    super(W18S90D0, self).__init__(fuse_bnbr)
60
    self.block_nonzeros = [
61
        6440,
62
        2668,
63
        5382,
64
        10765,
65
        21530,
66
        42688,
67
        84640,
68
        84640,
69
        84640,
70
        84640,
71
        84640,
72
        169280,
73
        338560,
74
    ]
75
    self.block_config = ["dense"] + ["sparse"] * (self.num_blocks - 1)
76
    self.width = 1.8
77

78

79
class W17S90D0(W18S90D0):
80
  """Width 1.7, sparsity 90%, dense first block."""
81

82
  def __init__(self, fuse_bnbr=False):
83
    super(W17S90D0, self).__init__(fuse_bnbr)
84
    self.block_nonzeros = [
85
        6048,
86
        2333,
87
        4666,
88
        9331,
89
        18662,
90
        37670,
91
        76038,
92
        76038,
93
        76038,
94
        76038,
95
        76038,
96
        152077,
97
        304154,
98
    ]
99
    self.width = 1.7
100

101

102
class W16S90D0(W18S90D0):
103
  """Width 1.6, sparsity 90%, dense first block."""
104

105
  def __init__(self, fuse_bnbr=False):
106
    super(W16S90D0, self).__init__(fuse_bnbr)
107
    self.block_nonzeros = [
108
        4896,
109
        2122,
110
        4326,
111
        8486,
112
        16646,
113
        33293,
114
        66586,
115
        66586,
116
        66586,
117
        66586,
118
        66586,
119
        133824,
120
        268960,
121
    ]
122
    self.width = 1.6
123

124

125
class W15S90D0(W18S90D0):
126
  """Width 1.5, sparsity 90%, dense first block."""
127

128
  def __init__(self, fuse_bnbr=False):
129
    super(W15S90D0, self).__init__(fuse_bnbr)
130
    self.block_nonzeros = [
131
        4608,
132
        1843,
133
        3686,
134
        7373,
135
        14746,
136
        29491,
137
        58982,
138
        58982,
139
        58982,
140
        58982,
141
        58982,
142
        117965,
143
        235930,
144
    ]
145
    self.width = 1.5
146

147

148
class W14S90D0(W18S90D0):
149
  """Width 1.4, sparsity 90%, dense first block."""
150

151
  def __init__(self, fuse_bnbr=False):
152
    super(W14S90D0, self).__init__(fuse_bnbr)
153
    self.block_nonzeros = [
154
        4272,
155
        1566,
156
        3098,
157
        6336,
158
        12960,
159
        25920,
160
        51840,
161
        51840,
162
        51840,
163
        51840,
164
        51840,
165
        103104,
166
        205062,
167
    ]
168
    self.width = 1.4
169

170

171
class W13S90D0(W18S90D0):
172
  """Width 1.3, sparsity 90%, dense first block."""
173

174
  def __init__(self, fuse_bnbr=False):
175
    super(W13S90D0, self).__init__(fuse_bnbr)
176
    self.block_nonzeros = [
177
        3320,
178
        1394,
179
        2822,
180
        5645,
181
        11290,
182
        22310,
183
        44090,
184
        44090,
185
        44090,
186
        44090,
187
        44090,
188
        88179,
189
        176358,
190
    ]
191
    self.width = 1.3
192

193

194
def get_config(width, sparsity):
195
  """Getter for model configuration."""
196
  tag = "_".join(["mb"] + [str(int(x)) for x in [width * 10, sparsity * 100]])
197
  configs = {
198
      "mb_10_0": W10,
199
      "mb_12_0": W12,
200
      "mb_14_0": W14,
201
      "mb_13_90": W13S90D0,
202
      "mb_14_90": W14S90D0,
203
      "mb_15_90": W15S90D0,
204
      "mb_16_90": W16S90D0,
205
      "mb_17_90": W17S90D0,
206
      "mb_18_90": W18S90D0,
207
  }
208
  if tag not in configs:
209
    raise ValueError("Could not find config for width {}, sparsity {}.".format(
210
        width, sparsity))
211
  return configs[tag]
212