Insert dropout to MobileNet

mobilenet.py

import tensorflow as tf
from tensorflow.keras.applications import MobileNet
from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout
from tensorflow.keras.models import Model
from tensorflow.keras.callbacks import History
from tensorflow.contrib.tpu.python.tpu import keras_support
import tensorflow.keras.backend as K

from keras.datasets import cifar10
from keras.utils import to_categorical
import numpy as np
import pickle, os

def create_normal_model():
    model = MobileNet(include_top=False, weights="imagenet", input_shape=(128,128,3))
    x = GlobalAveragePooling2D()(model.layers[-1].output)
    x = Dense(10, activation="softmax")(x)

    return Model(model.inputs, x)

def create_dropout_model():
    model = create_normal_model()
    # ActivationのあとにDropoutを入れる
    for i, layer in enumerate(model.layers):
        if i==0:
            input = layer.input
            x = input
        else:
            if "relu" in layer.name:
                x = layer(x)
                x = Dropout(0.01)(x)
            else:
                x = layer(x)
    drop_model = Model(input, x)
    return drop_model

def generator(X, y, batch_size):
    # 32->128に引き伸ばす操作でメモリ食い過ぎるので自分でジェネレーター書く
    while True:
        indices = np.arange(X.shape[0])
        np.random.shuffle(indices)
        for i in range(X.shape[0]//batch_size):
            current_indices = indices[i*batch_size:(i+1)*batch_size]
            X_batch = X[current_indices]
            # 4倍に引き伸ばす
            X_batch = X_batch.repeat(4, axis=1).repeat(4, axis=2)
            X_batch = X_batch / 255.0
            y_batch = to_categorical(y[current_indices], 10)
            yield X_batch, y_batch

def train(use_dropout):
    (X_train, y_train), (X_test, y_test) = cifar10.load_data()

    if use_dropout:
        model = create_dropout_model()
    else:
        model = create_normal_model()

    # 元の係数を壊さないように低めの学習率を使う
    model.compile(tf.train.RMSPropOptimizer(1e-5), "categorical_crossentropy", ["acc"])
    model.summary()

    tpu_grpc_url = "grpc://"+os.environ["COLAB_TPU_ADDR"]
    tpu_cluster_resolver = tf.contrib.cluster_resolver.TPUClusterResolver(tpu_grpc_url)
    strategy = keras_support.TPUDistributionStrategy(tpu_cluster_resolver)
    model = tf.contrib.tpu.keras_to_tpu_model(model, strategy=strategy)

    hist = History()
    batch_size = 1024

    model.fit_generator(generator(X_train, y_train, batch_size), steps_per_epoch=X_train.shape[0]//batch_size,
                        validation_data=generator(X_test, y_test, batch_size), 
                        validation_steps=X_test.shape[0]//batch_size, callbacks=[hist],
                        epochs=100)
    with open(f"history_drop_{use_dropout}.dat", "wb") as fp:
        pickle.dump(hist.history, fp)

if __name__ == "__main__":
    K.clear_session()
    train(True)