RICAP on CIFAR

static_ricap.py

import tensorflow as tf
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation, Input, AveragePooling2D, GlobalAveragePooling2D, Dense
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.image import ImageDataGenerator
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
from ricap import ricap
import os, pickle, glob, zipfile

def create_block(input, ch, reps):
    x = input
    for i in range(reps):
        x = Conv2D(ch, 3, padding="same")(x)
        x = BatchNormalization()(x)
        x = Activation("relu")(x)
    return x

def create_network():
    input = Input((32,32,3))
    x = create_block(input, 64, 3)
    x = AveragePooling2D(2)(x)
    x = create_block(x, 128, 3)
    x = AveragePooling2D(2)(x)
    x = create_block(x, 256, 3)
    x = GlobalAveragePooling2D()(x)
    x = Dense(10, activation="softmax")(x)

    return Model(input, x)

class RICAPGenerator(ImageDataGenerator):
    def __init__(self, use_batchwise_random, beta, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.use_batchwise_random = use_batchwise_random
        self.beta = beta

    def flow(self, *args, **kwargs):
        for X_batch, y_batch in super().flow(*args, **kwargs):
            if self.beta > 0:
                X_ricap, y_ricap = ricap(X_batch, y_batch, self.beta, self.use_batchwise_random)
            else:
                X_ricap, y_ricap = X_batch, y_batch # No RICAP
            yield X_ricap, y_ricap

def train(use_same_random_on_batch, beta):
    (X_train, y_train), (X_test, y_test) = cifar10.load_data()
    y_train, y_test = to_categorical(y_train), to_categorical(y_test)

    batch_size = 2048
    train_generator = RICAPGenerator(rescale=1.0/255.0, use_batchwise_random=use_same_random_on_batch, beta=beta).flow(
        X_train, y_train, batch_size=batch_size)
    test_generator = ImageDataGenerator(rescale=1.0/255.0).flow(
        X_test, y_test, batch_size=batch_size)

    model = create_network()
    model.compile(tf.train.AdamOptimizer(), "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)

    if not os.path.exists("static_result"):
        os.mkdir("static_result")

    hist = History()
    model.fit_generator(train_generator, steps_per_epoch=X_train.shape[0]//batch_size,
                        validation_data=test_generator, validation_steps=X_test.shape[0]//batch_size,
                        epochs=150, callbacks=[hist])

    history = hist.history
    with open(f"static_result/{use_same_random_on_batch}_{beta}.dat", "wb") as fp:
        pickle.dump(history, fp)

if __name__ == "__main__":
    for flag in [True, False]:
        train(False, 0)
        for beta in [0.1, 0.3, 0.5, 0.75, 1]:
            print(flag, beta, "Starts")
            train(flag, beta)

    with zipfile.ZipFile("static_result.zip", "w") as zip:
        for f in glob.glob("static_result/*.dat"):
            zip.write(f)