Peele Net CIFAR-10

cifar10.py

from pelee_net import PeleeNet
import tensorflow as tf
import tensorflow.keras as keras
import numpy as np
from PIL import Image
import pickle
import os
from tensorflow.contrib.tpu.python.tpu import keras_support

def generator(X, y, batch_size, use_augmentation, shuffle, scale):
    if use_augmentation:
        base_gen = keras.preprocessing.image.ImageDataGenerator(
            horizontal_flip=True,
            width_shift_range=4.0/32.0,
            height_shift_range=4.0/32.0)
    else:
        base_gen = keras.preprocessing.image.ImageDataGenerator()
    for X_base, y_base in base_gen.flow(X, y, batch_size=batch_size, shuffle=shuffle):
        if scale != 1:
            X_batch = np.zeros((X_base.shape[0], X_base.shape[1]*scale,
                                X_base.shape[2]*scale, X_base.shape[3]), np.float32)
            for i in range(X_base.shape[0]):
                with Image.fromarray(X_base[i].astype(np.uint8)) as img:
                    img = img.resize((X_base.shape[1]*scale, X_base.shape[2]*scale), Image.LANCZOS)
                    X_batch[i] = np.asarray(img, np.float32) / 255.0
        else:
            X_batch = X_base / 255.0
        yield X_batch, y_base

def lr_scheduler(epoch):
    x = 0.4
    if epoch >= 70: x /= 5.0
    if epoch >= 120: x /= 5.0
    if epoch >= 170: x /= 5.0
    return x

def train(use_augmentation, use_stem_block):
    tf.logging.set_verbosity(tf.logging.FATAL)
    (X_train, y_train), (X_test, y_test) = keras.datasets.cifar10.load_data()
    y_train = keras.utils.to_categorical(y_train)
    y_test = keras.utils.to_categorical(y_test)

    # generator
    batch_size = 512
    scale = 7 if use_stem_block else 1
    train_gen = generator(X_train, y_train, batch_size=batch_size,
                          use_augmentation=use_augmentation, shuffle=True, scale=scale)
    test_gen = generator(X_test, y_test, batch_size=1000,
                         use_augmentation=False, shuffle=False, scale=scale)
    
    # network
    input_shape = (224,224,3) if use_stem_block else (32,32,3)
    model = PeleeNet(input_shape=input_shape, use_stem_block=use_stem_block, n_classes=10)
    model.compile(keras.optimizers.SGD(0.4, 0.9), "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)

    scheduler = keras.callbacks.LearningRateScheduler(lr_scheduler)
    hist = keras.callbacks.History()

    model.fit_generator(train_gen, steps_per_epoch=X_train.shape[0]//batch_size,
                        validation_data=test_gen, validation_steps=X_test.shape[0]//1000,
                        callbacks=[scheduler, hist], epochs=1, max_queue_size=1)
    history = hist.history
    with open(f"pelee_aug_{use_augmentation}_stem_{use_stem_block}.pkl", "wb") as fp:
        pickle.dump(history, fp)

if __name__ == "__main__":
    train(True, True)

pelee_net.py

import tensorflow.keras as keras
import tensorflow.keras.layers as layers

def conv_bn_relu(input_tensor, ch, kernel, padding="same", strides=1):
    x = layers.Conv2D(ch, kernel, padding=padding, strides=strides, kernel_regularizer=keras.regularizers.l2(5e-4))(input_tensor)
    x = layers.BatchNormalization()(x)
    return layers.Activation("relu")(x)

def stem_block(input_tensor):
    x = conv_bn_relu(input_tensor, 32, 3, strides=2)
    branch1 = conv_bn_relu(x, 16, 1)
    branch1 = conv_bn_relu(branch1, 32, 3, strides=2)
    branch2 = layers.MaxPool2D(2)(x)
    x = layers.Concatenate()([branch1, branch2])
    return conv_bn_relu(x, 32, 1)

def dense_block(input_tensor, num_layers, growth_rate, bottleneck_width):
    x = input_tensor
    growth_rate = int(growth_rate / 2)

    for i in range(num_layers):
        inter_channel = int(growth_rate*bottleneck_width/4) * 4
        branch1 = conv_bn_relu(x, inter_channel, 1)
        branch1 = conv_bn_relu(branch1, growth_rate, 3)

        branch2 = conv_bn_relu(x, inter_channel, 1)
        branch2 = conv_bn_relu(branch2, growth_rate, 3)
        branch2 = conv_bn_relu(branch2, growth_rate, 3)
        x = layers.Concatenate()([x, branch1, branch2])
    return x

def transition_layer(input_tensor, k, use_pooling=True):
    x = conv_bn_relu(input_tensor, k, 1)
    if use_pooling:
        return layers.AveragePooling2D(2)(x)
    else:
        return x

def PeleeNet(input_shape=(224,224,3), use_stem_block=True, n_classes=1000):
    n_dense_layers = [3,4,8,6]
    bottleneck_width = [1,2,4,4]
    out_layers = [128,256,512,704]
    growth_rate = 32

    input = layers.Input(input_shape)
    x = stem_block(input) if use_stem_block else input
    for i in range(4):
        x = dense_block(x, n_dense_layers[i], growth_rate, bottleneck_width[i])
        use_pooling = i < 3
        x = transition_layer(x, out_layers[i], use_pooling=use_pooling)
    x = layers.GlobalAveragePooling2D()(x)
    x = layers.Dense(n_classes, activation="softmax")(x)
    return keras.models.Model(input, x)