Mobile_net_v2

import tensorflow as tf
import pandas as pd
import matplotlib.pyplot as plt


class BottleNeckDown(tf.keras.layers.Layer):
    def __init__(self, input_channel, t, c, strides, **kwargs):
        super().__init__(**kwargs)
        self.branch = tf.keras.models.Sequential([
            tf.keras.layers.Conv2D(filters=t * input_channel, kernel_size=1, strides=1, use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(max_value=6.0),
            tf.keras.layers.DepthwiseConv2D(kernel_size=3, padding='same', strides=strides, use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(max_value=6.0),
            tf.keras.layers.Conv2D(filters=c, kernel_size=1, strides=1, use_bias=False),
            tf.keras.layers.BatchNormalization()
        ])

    def call(self, inputs, **kwargs):
        x = self.branch(inputs, training=False)
        return x


class BottleNeck(tf.keras.layers.Layer):
    def __init__(self, input_channel, t, c, **kwargs):
        super().__init__(**kwargs)
        self.branch1 = tf.keras.models.Sequential([
            tf.keras.layers.Conv2D(filters=t * input_channel, kernel_size=1, strides=1, use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(max_value=6.0),
            tf.keras.layers.DepthwiseConv2D(kernel_size=3, padding='same', strides=1, use_bias=False),
            tf.keras.layers.BatchNormalization(),
            tf.keras.layers.ReLU(max_value=6.0),
            tf.keras.layers.Conv2D(filters=c, kernel_size=1, strides=1, use_bias=False)
        ])
        self.bn = tf.keras.layers.BatchNormalization()
        self.add = tf.keras.layers.Add()

    def call(self, inputs, **kwargs):
        branch1 = self.branch1(inputs, training=False)
        branch2 = inputs
        branch = self.add([branch1, branch2])
        x = self.bn(branch)
        return x


def build_mobile_net(num_classes, n_height=224, n_width=224):
    net_shape = [(1, 16, 1, 1),
                 (6, 24, 2, 2),
                 (6, 32, 3, 2),
                 (6, 64, 4, 2),
                 (6, 96, 3, 1),
                 (6, 160, 3, 2),
                 (6, 320, 1, 1)]
    input_layer = tf.keras.layers.Input(shape=(n_height, n_width, 3), dtype=tf.float32)
    x = tf.keras.layers.Conv2D(32, 3, 2, padding='same', use_bias=False)(input_layer)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.ReLU(max_value=6.0)(x)
    for t, c, n, s in net_shape:
        x = BottleNeckDown(x.shape[-1], t, c, s)(x)
        for i in range(1, n):
            x = BottleNeck(x.shape[-1], t, c)(x)
    x = tf.keras.layers.Conv2D(1280, 1, 1, use_bias=False)(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.ReLU(max_value=6.0)(x)
    x = tf.keras.layers.AvgPool2D(pool_size=7)(x)
    x = tf.keras.layers.Conv2D(num_classes, 1, 1, use_bias=False)(x)
    x = tf.squeeze(x, axis=[1, 2])
    prediction = tf.keras.layers.Softmax()(x)
    model = tf.keras.models.Model(inputs=input_layer, outputs=prediction)
    return model


def draw_curve(history):
    pd.DataFrame(history.history).plot()
    plt.gca().set_ylim(0, 1)
    plt.grid(True)
    plt.show()


mobile_net = build_mobile_net(10)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
mobile_net.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=['acc'])
mobile_net.summary()
train_dir = './training/training'
valid_dir = './validation/validation'
train_dataset = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255.,
    horizontal_flip=True,
    vertical_flip=True,
    width_shift_range=0.2,
    height_shift_range=0.2,
    rotation_range=90,
    shear_range=0.2,
    zoom_range=0.2,
    fill_mode='nearest'
)
valid_dataset = tf.keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255.
)
train_dataset = train_dataset.flow_from_directory(
    directory=train_dir,
    batch_size=32,
    target_size=(224, 224),
    class_mode='categorical',
    shuffle=True
)
valid_dataset = valid_dataset.flow_from_directory(
    directory=valid_dir,
    batch_size=32,
    target_size=(224, 224),
    class_mode='categorical',
    shuffle=False
)
history = mobile_net.fit(train_dataset,
                         steps_per_epoch=train_dataset.samples // 32,
                         validation_data=valid_dataset,
                         validation_steps=valid_dataset.samples // 32,
                         epochs=5)
# draw_curve(history)
mobile_net.evaluate(valid_dataset, steps=valid_dataset.samples // 32)