simple_softmax

simple_softmax.py

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Model
from tensorflow.keras.callbacks import Callback, History
import tensorflow.keras.backend as K
from tensorflow.contrib.tpu.python.tpu import keras_support

from train1000 import cifar10
import numpy as np
import os, json, tarfile

# VGG-like model
def create_model():
    input = layers.Input((32, 32, 3))
    x = input
    for i in range(4):
        for j in range(3):
            x = layers.Conv2D(64*(2**i), 3, padding="same")(x)
            x = layers.BatchNormalization()(x)
            x = layers.Activation("relu")(x)
        if i != 3:
            x = layers.AveragePooling2D(2)(x)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(10, activation="softmax")(x)

    return Model(input, x)

def data_augmentation(image):
    outputs = np.zeros(image.shape, dtype=np.float32)
    # crop
    crop_x = np.random.randint(0, 4)
    crop_y = np.random.randint(0, 4)
    outputs[crop_x:crop_x+28, crop_y:crop_y+28, :] = image[crop_x:crop_x+28, crop_y:crop_y+28, :]
    # flip
    if np.random.rand() >= 0.5:
        outputs = outputs[:, ::-1, :]
    return outputs

def generator(X, y, batch_size, use_augmentation):
    while True:
        X_cache, y_cache = [], []
        indices = np.random.permutation(X.shape[0])
        for i in indices:
            if use_augmentation:
                X_cache.append(data_augmentation(X[i]))
            else:
                X_cache.append(X[i])
            y_cache.append(y[i])
            if(len(y_cache)==batch_size):
                X_batch = np.asarray(X_cache, np.float32)
                y_batch = np.asarray(y_cache, np.float32)
                X_cache, y_cache = [], []
                yield X_batch, y_batch

def train(use_aug):
    (X_train, y_train), (X_test, y_test) = cifar10()

    model = create_model()
    model.compile(tf.train.AdamOptimizer(1e-3), "categorical_crossentropy", ["acc"])

    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)

    batch_size = 200
    hist = History()
    model.fit_generator(generator(X_train, y_train, batch_size, use_aug), steps_per_epoch=X_train.shape[0]//batch_size,
                        validation_data=generator(X_test, y_test, batch_size, False), validation_steps=X_test.shape[0]//batch_size,
                        max_queue_size=1, callbacks=[hist], epochs=200)
    history = hist.history
    print(max(history["val_acc"]))

    with open(f"softmax_aug_{use_aug}.json", "w") as fp:
        json.dump(history, fp)

if __name__ == "__main__":
    for aug in [False, True]:
        K.clear_session()
        print(aug, "starts")
        train(aug)