MobileFaceNet 代码

class Bottleneck(tf.keras.layers.Layer):
  def __init__(self, input_dim, output_dim, expansion_factor, stride):
    self.input_dim = input_dim
    self.output_dim = output_dim
    self.stride = stride
    super(Bottleneck, self).__init__()
    self.conv1 = tf.keras.layers.Conv2D(filters=input_dim * expansion_factor, kernel_size = (1,1),
                      strides = (1,1),padding = "same")
    self.dwise = tf.keras.layers.DepthwiseConv2D(kernel_size = (3,3), strides = self.stride,
                      padding = "same")
    self.conv2 = tf.keras.layers.Conv2D(filters = output_dim, kernel_size = (1,1),
                      strides = (1,1),activation = 'linear')
  def call(self, inputs):
    x = inputs
    x = self.conv1(x)
    x = self.dwise(x)
    out = self.conv2(x)
    if self.stride == 1 and self.input_dim == self.output_dim:
      return tf.keras.layers.add([inputs, out])
    
    return out

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

class GDConv(tf.keras.layers.Layer):
  def __init__(self, input_dim, stride):
    self.input_dim = input_dim
    super(GDConv, self).__init__()
  def build(self, input_shape):
    self.W = self.add_weight(name = "W",shape = (input_shape[1], input_shape[2], input_shape[3]),
                initializer ='uniform', trainable = True)
  def call(self, inputs):
    x = inputs
    out = tf.multiply(x, self.W)
    out = tf.reduce_sum(out, [1,2])
    out = tf.expand_dims(out,[1])
    out = tf.expand_dims(out,[1])
    return out

def GDConv_build(input_channels, s):
  gdconv = tf.keras.Sequential()
  gdconv.add(GDConv(input_dim=input_channels, stride=s))
  return gdconv

class MobileFaceNet(tf.keras.Model):
  def __init__(self):
    super(MobileFaceNet, self).__init__()
    self.conv1 = tf.keras.layers.Conv2D(filters = 64, kernel_size=(3,3),
                      strides=(2,2),padding="same")
    self.dwise = tf.keras.layers.DepthwiseConv2D(kernel_size = (3,3),strides = (1,1),
                          padding="same")
    self.conv2 = tf.keras.layers.Conv2D(filters = 64, kernel_size=(1,1),strides =(1,1),
                      padding = "same")
    self.bottleneck1 = bottleneck_build(input_channels=64, output_channels=64, t = 2,
                      n=5, s=2)
    self.bottleneck2 = bottleneck_build(input_channels=64, output_channels=128, t = 4,
                      n=1, s=2)
    self.bottleneck3 = bottleneck_build(input_channels=128, output_channels=128, t = 2,
                      n=6, s=1)
    self.bottleneck4 = bottleneck_build(input_channels=128, output_channels=128, t = 4,
                      n=1, s=2)
    self.bottleneck5 = bottleneck_build(input_channels=128, output_channels=128, t = 2,
                      n=2, s=1)
    self.conv3 = tf.keras.layers.Conv2D(filters=512, kernel_size=(1,1), strides=(1,1))
    self.gdconv = GDConv_build(input_channels=512, s=1)
    self.conv4 = tf.keras.layers.Conv2D(filters = 128, kernel_size=(1,1),
                      strides=(1,1),activation="linear")
    
  def call(self,inputs):
    x = inputs
    x = self.conv1(x)
    x = self.dwise(x)
    x = self.conv2(x)
    x = self.bottleneck1(x)
    x = self.bottleneck2(x)
    x = self.bottleneck3(x)
    x = self.bottleneck4(x)
    x = self.bottleneck5(x)
    x = self.conv3(x)
    x = self.gdconv(x)
    x = self.conv4(x)
    return x
  def model(self):
    x = tf.keras.layers.Input(shape=(112,112, 3))
    return tf.keras.Model(inputs=[x], outputs=self.call(x))

model = MobileFaceNet()
model.model().summary()