MobileNet V2 代码

class Bottleneck(tf.keras.layers.Layer):
  def __init__(self, input_channels, output_channels, expansion_factor, stride):
    self.stride = stride
    self.input_channels = input_channels
    self.output_channels = output_channels
    super(Bottleneck, self).__init__()

    self.conv1 = tf.keras.layers.Conv2D(filters = output_channels * expansion_factor,
                      strides = 1, kernel_size = (1,1), padding = "same")
    self.dwconv = tf.keras.layers.DepthwiseConv2D(kernel_size=(3,3),strides = stride, padding = "same")

    self.conv2 = tf.keras.layers.Conv2D(filters = output_channels, kernel_size = (1,1),strides = 1,
                       activation = 'linear')
    
  def call(self, inputs, training = None, **kwargs):
    x = self.conv1(inputs)
    x = self.dwconv(x)
    x = self.conv2(x)
    if self.stride == 1 and self.input_channels == self.output_channels:
      x = tf.keras.layers.add([x, inputs])

    return x

def build_bottleneck(t, input_channels, output_channels, n, s):
  bottleneck = tf.keras.Sequential()
  for i in range(n):
    if i == 0:
      bottleneck.add(Bottleneck(input_channels, output_channels, t, s))
    else:
      bottleneck.add(Bottleneck(output_channels, output_channels, t, stride = 1))

  return bottleneck

class MobileNetV2(tf.keras.Model):
  def __init__(self):
    super(MobileNetV2, self).__init__()
    self.conv1 = tf.keras.layers.Conv2D(filters = 32, kernel_size = (3,3),
                      strides = 2, padding = "same")
    self.bottleneck1 = build_bottleneck(t = 1, input_channels = 32, output_channels=16,
                      n = 1, s = 1)
    self.bottleneck2 = build_bottleneck(t = 6, input_channels=16, output_channels=24,
                      n = 2, s = 2)
    self.bottleneck3 = build_bottleneck(t = 6, input_channels=24, output_channels=32,
                      n = 3, s = 2)
    self.bottleneck4 = build_bottleneck(t = 6, input_channels=32, output_channels=64,
                      n = 4, s = 2)
    self.bottleneck5 = build_bottleneck(t = 6, input_channels=64, output_channels=96,
                      n = 3, s = 1)
    self.bottleneck6 = build_bottleneck(t = 6, input_channels=96, output_channels=160,
                      n = 3, s = 2)
    self.bottleneck7 = build_bottleneck(t = 6, input_channels=160, output_channels=320,
                      n = 1, s = 1)
    self.conv2 = tf.keras.layers.Conv2D(filters = 1280, kernel_size =(1,1), 
                      strides = 1, padding = "same")
    self.avgpool = tf.keras.layers.AveragePooling2D(pool_size = (7,7))
    self.conv3 = tf.keras.layers.Conv2D(filters = 10, kernel_size = (1,1),
                      strides = 1, padding = "same", activation = "softmax")
  
  def call(self, inputs):
    x = inputs
    x = self.conv1(x)
    x = self.bottleneck1(x)
    x = self.bottleneck2(x)
    x = self.bottleneck3(x)
    x = self.bottleneck4(x)
    x = self.bottleneck5(x)
    x = self.bottleneck6(x)
    x = self.bottleneck7(x)

    x = self.conv2(x)
    x = self.avgpool(x)
    x = self.conv3(x)

    return x 

model = MobileNetV2()
model.build(input_shape = (None, 224, 224, 3))
model.summary()

参考资料：
https://github.com/calmisential/Basic_CNNs_TensorFlow2/blob/master/models/mobilenet_v1.py
https://stackoverflow.com/questions/55235212/model-summary-cant-print-output-shape-while-using-subclass-model