/
MobileNetV2.py
93 lines (77 loc) · 3 KB
/
MobileNetV2.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
import tensorflow as tf
from tensorflow.keras import layers, Sequential, Model
def ReLU6():
return layers.Lambda(lambda x: tf.nn.relu6(x))
class LinearBottleNeck(layers.Layer):
def __init__(self, in_channels, out_channels, strides=1, t=6):
super(LinearBottleNeck, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.strides = strides
self.residual = Sequential([
layers.Conv2D(in_channels * t,
(1, 1),
strides=1,
padding='same'),
layers.BatchNormalization(),
ReLU6(),
layers.DepthwiseConv2D((3, 3),
strides=strides,
padding='same'),
layers.BatchNormalization(),
ReLU6(),
layers.Conv2D(out_channels,
(1, 1),
strides=1,
padding='same'),
layers.BatchNormalization(),
])
def call(self, x, training=False):
residual = self.residual(x, training=training)
if self.strides == 1 and self.in_channels == self.out_channels:
residual += x
return residual
class MobileNetV2(Model):
def __init__(self, num_classes, input_shape=(32, 32, 3)):
super(MobileNetV2, self).__init__()
self.front = Sequential([
layers.Input(input_shape),
layers.BatchNormalization(),
ReLU6()
])
self.stage1 = LinearBottleNeck(32, 16, 1, 1)
self.stage2 = self._make_stage(2, 16, 24, 2, 6)
self.stage3 = self._make_stage(3, 24, 32, 2, 6)
self.stage4 = self._make_stage(4, 32, 64, 2, 6)
self.stage5 = self._make_stage(3, 64, 96, 1, 6)
self.stage6 = self._make_stage(3, 96, 160, 1, 6)
self.stage7 = LinearBottleNeck(160, 320, 1, 6)
self.conv1 = layers.Conv2D(filters=1280,
kernel_size=(1, 1),
strides=1,
padding="same")
self.ap = layers.AveragePooling2D((7, 7))
self.fc = layers.Dense(num_classes, activation='softmax')
def _make_stage(self, repeat, in_channels, out_channels, strides, t):
nets = Sequential()
nets.add(LinearBottleNeck(in_channels, out_channels, strides, t))
while repeat - 1:
nets.add(LinearBottleNeck(out_channels, out_channels, 1, t))
repeat -= 1
return nets
def call(self, inputs, training=False):
x = self.front(inputs)
x = self.stage1(x)
x = self.stage2(x)
x = self.stage3(x)
x = self.stage4(x)
x = self.stage5(x)
x = self.stage6(x)
x = self.stage7(x)
x = self.conv1(x)
x = self.ap(x)
x = tf.reshape(x, (x.shape[0], -1))
x = self.fc(x)
return x
def mobilenetv2(num_classes):
return MobileNetV2(num_classes)