mobilenetv2 tensorflow2.x版本重构

tensorflow1和2版本的api变化很大，1.x版本又已经过期了，所以今天花了点时间用tensorflow2重构了一下mobilenet的代码，这里直接上代码，原本1.x代码在转载的文章里面， 有需要的同学自取即可，如有疑问，feel free to contact me：

Mobilenetv2.py

"""
tensorflow 2.0重构 mobilenetv2
"""
import tensorflow as tf
from keras import layers, initializers, activations, Model


class Conv2DBlock(Model):
    def __init__(self, kernel_size, out_channels, stride, is_training=False):
        super(Conv2DBlock, self).__init__()
        self.conv = layers.Conv2D(filters=out_channels, kernel_size=kernel_size,
                                  strides=[stride, stride], padding='SAME',
                                  kernel_initializer="truncated_normal",
                                  bias_initializer=initializers.Constant(value=0.01))
        self.bn = layers.BatchNormalization(trainable=is_training)

    def call(self, inputs, training=None, mask=None):
        output_conv = self.conv(inputs)
        output_bn = self.bn(output_conv)
        return output_bn


class DepthwiseConv2DBlock(Model):
    def __init__(self, stride, width_multiplier, is_training):
        super(DepthwiseConv2DBlock, self).__init__()
        self.conv = layers.DepthwiseConv2D(kernel_size=3, strides=stride, depth_multiplier=width_multiplier,
                                           padding="SAME", kernel_initializer="truncated_normal",
                                           bias_initializer=initializers.Constant(value=0.01))
        self.bn = layers.BatchNormalization(trainable=is_training)

    def call(self, inputs, training=None, mask=None):
        output_conv = self.conv(inputs)
        output_bn = self.bn(output_conv)
        return activations.relu6(output_bn)


class BottleNeck(Model):
    def __init__(self, in_channels, t, out_channels, stride, width_multiplier, r=False, is_training=False):
        """
        残差模块
        :param in_channels:
        :param t:
        :param out_channels:
        :param stride:
        :param r: 表示是否用残差(只有当步长为1并且输入维度等于输出维度时才用)
        """
        super(BottleNeck, self).__init__()
        expand_dim = t * in_channels
        self.conv1 = Conv2DBlock(kernel_size=1, out_channels=expand_dim, stride=1, is_training=is_training)
        self.dwconv = DepthwiseConv2DBlock(stride=stride, width_multiplier=width_multiplier, is_training=is_training)
        self.conv2 = Conv2DBlock(kernel_size=1, out_channels=out_channels, stride=1, is_training=is_training)
        self.r = r

    def call(self, inputs, training=None, mask=None):
        output_conv = self.conv1(inputs)
        output_relu6 = activations.relu6(output_conv)
        output_dwconv = self.dwconv(output_relu6)
        outputs = self.conv2(output_dwconv)
        if self.r:
            outputs = tf.add(outputs, inputs)
        return outputs


class MobilenetV2(Model):
    def __init__(self, num_classes=8, is_training=True, width_multiplier=1):
        super(MobilenetV2, self).__init__()
        self.num_classes = num_classes
        self.is_training = is_training
        self.width_multiplier = width_multiplier

        self.conv1 = Conv2DBlock(kernel_size=3, out_channels=32, stride=2)  # 输出大小112*112*32
        # 只有步长为1并且输入维度等于输出时，才使用残差
        self.b1_1 = BottleNeck(in_channels=32, t=1, out_channels=16, stride=1, r=False,
                               width_multiplier=self.width_multiplier)  # 输出大小112*112*16，一层不使用残差

        self.b2_1 = BottleNeck(in_channels=16, t=6, out_channels=24, stride=2, r=False,
                               width_multiplier=self.width_multiplier) # 输出大小56*56*24
        self.b2_2 = BottleNeck(in_channels=24, t=6, out_channels=24, stride=1, r=True,
                               width_multiplier=self.width_multiplier) # 输出大小56*56*24

        self.b3_1 = BottleNeck(in_channels=24, t=6, out_channels=32, stride=2, r=False,
                               width_multiplier=self.width_multiplier) # 输出大小28*28*32
        self.b3_2 = BottleNeck(in_channels=32, t=6, out_channels=32, stride=1, r=True,
                               width_multiplier=self.width_multiplier) # 输出大小输出大小28*28*32
        self.b3_3 = BottleNeck(in_channels=32, t=6, out_channels=32, stride=1, r=True,
                               width_multiplier=self.width_multiplier) # 输出大小输出大小28*28*32

        self.b4_1 = BottleNeck(in_channels=32, t=6, out_channels=64, stride=2, r=False,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*64
        self.b4_2 = BottleNeck(in_channels=64, t=6, out_channels=64, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*64
        self.b4_3 = BottleNeck(in_channels=64, t=6, out_channels=64, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*64
        self.b4_4 = BottleNeck(in_channels=64, t=6, out_channels=64, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*64

        self.b5_1 = BottleNeck(in_channels=64, t=6, out_channels=96, stride=1, r=False,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*96
        self.b5_2 = BottleNeck(in_channels=96, t=6, out_channels=96, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*96
        self.b5_3 = BottleNeck(in_channels=96, t=6, out_channels=96, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小14*14*96

        self.b6_1 = BottleNeck(in_channels=96, t=6, out_channels=160, stride=2, r=False,
                               width_multiplier=self.width_multiplier)  # 输出大小7*7*160
        self.b6_2 = BottleNeck(in_channels=160, t=6, out_channels=160, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小7*7*160
        self.b6_3 = BottleNeck(in_channels=160, t=6, out_channels=160, stride=1, r=True,
                               width_multiplier=self.width_multiplier)  # 输出大小7*7*160

        self.b7_1 = BottleNeck(in_channels=160, t=6, out_channels=320, stride=1, r=False,
                               width_multiplier=self.width_multiplier)  # 输出大小7*7*320

        self.conv8 = Conv2DBlock(kernel_size=1, out_channels=1280, stride=1) # 输出大小7*7*1280

        self.conv10_1 = Conv2DBlock(kernel_size=1, out_channels=self.num_classes, stride=1) # 输出大小1*1*8

    def call(self, inputs, training=None, mask=None):
        conv1_output = self.conv1(inputs)

        b1_1_output = self.b1_1(conv1_output)

        b2_1_output = self.b2_1(b1_1_output)
        b2_2_output = self.b2_2(b2_1_output)

        b3_1_output = self.b3_1(b2_2_output)
        b3_2_output = self.b3_2(b3_1_output)
        b3_3_output = self.b3_3(b3_2_output)

        b4_1_output = self.b4_1(b3_3_output)
        b4_2_output = self.b4_2(b4_1_output)
        b4_3_output = self.b4_3(b4_2_output)
        b4_4_output = self.b4_4(b4_3_output)

        b5_1_output = self.b5_1(b4_4_output)
        b5_2_output = self.b5_2(b5_1_output)
        b5_3_output = self.b5_3(b5_2_output)

        b6_1_output = self.b6_1(b5_3_output)
        b6_2_output = self.b6_2(b6_1_output)
        b6_3_output = self.b6_3(b6_2_output)

        b7_1_output = self.b7_1(b6_3_output)

        conv8_output = self.conv8(b7_1_output)

        pool9_1_output = tf.reduce_mean(conv8_output, [1, 2], keepdims=True)

        conv10_1_output = self.conv10_1(pool9_1_output)

        result = tf.squeeze(conv10_1_output)

        return result


if __name__ == '__main__':
    # model = DepthwiseConv2DBlock(2, 1, False)
    # model.build(input_shape=(None, 224, 224, 3))
    # model.call(layers.Input(shape=(224, 224, 3)))
    # model.summary()
    #
    # input_tensor = tf.random.normal(shape=(10, 224, 224, 3))
    # y = model(input_tensor)
    # print(y.shape)

    # model = BottleNeck(in_channels=3, t=1, out_channels=10, stride=1, width_multiplier=1)
    # model.build(input_shape=(None, 224, 224, 3))
    # model.call(layers.Input(shape=(224, 224, 3)))
    # model.summary()
    #
    # input_tensor = tf.random.normal(shape=(10, 224, 224, 3))
    # y = model(input_tensor)
    # print(y.shape)

    model = MobilenetV2()
    model.build(input_shape=(None, 224, 224, 3))
    model.call(layers.Input(shape=(224, 224, 3)))
    model.summary()

    input_tensor = tf.random.normal(shape=(10, 224, 224, 3))
    y = model(input_tensor)
    print(y.shape)