generic_trainer.py

# coding: utf-8

__author__ = 'rishab-sharma'

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import keras
from keras.preprocessing import image
from keras.models import Model
from keras.layers import Dense, GlobalAveragePooling2D , Dropout
from keras import backend as K
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import TensorBoard , EarlyStopping
import argparse
from keras import metrics
from keras.applications import *
from keras.applications.inception_v3 import preprocess_input as pi_one
from logo.logo import logo
from PIL import ImageFile
import time
from keras.regularizers import l2
ImageFile.LOAD_TRUNCATED_IMAGES = True

obj = logo()
obj.printer()

ap = argparse.ArgumentParser()
ap.add_argument("-c", "--classes", required=True,
    help="Number of Classes in Your Data")
ap.add_argument("-f", "--freeze", default=172,#vgg19=21,vgg16=12,inceptionV3=172,249
    help="Freeze Limit for Fine Tuning")
ap.add_argument("-b", "--batch", default=32,
    help="Batch Size for Training")
ap.add_argument("-e1", "--epoch1", default=10,
    help="Number of Epochs for first round of Training")
ap.add_argument("-e2", "--epoch2", default=30,
    help="Number of Epochs for Second Round of Training")
ap.add_argument("-d", "--data", default='data',
    help="Data Directory")
ap.add_argument("-v", "--validation", default=100,
    help="Number of Validation Samples")
ap.add_argument("-ckpt", "--ckpt", default='models',
    help="Checkpoint Directory Directory")
ap.add_argument("-m", "--model", required=True,
    help="Model You Want to Fine Tune")
ap.add_argument("-ci", "--imbalance", required=True,
    help="Use Class Imbalance or Not (0/1)")
ap.add_argument("-pi", "--preprocess", required=True,
    help="Use Imagenet Preprocess or Not (0/1)")
args = vars(ap.parse_args())

model_name = str(args["model"])

list_model = {
        "Xception": Xception,
        "InceptionV3": InceptionV3,
        "InceptionResNetV2": InceptionResNetV2,
        "VGG16": VGG16,
        "VGG19": VGG19,
        "MobileNet": MobileNet,
	   "NASNet-m":NASNetMobile,
	   "NASNet-l":NASNetLarge
    }

############### DIVISION ###############

print("[INFO] Setting up certain hyperparameters")

epochs_1 = int(args["epoch1"])
epochs_2 = int(args["epoch2"])
freeze_limit = int(args["freeze"])
nb_classes = int(args["classes"])
batch = int(args["batch"])
ckpt = str(args["ckpt"])
climb = int(args["imbalance"])
preprocess = int(args["preprocess"])

if not os.path.exists("{}".format(ckpt)):
        os.makedirs("{}".format(ckpt))
        
print("epochs_1 = {} , epochs_2 = {} , freeze_limit = {} , classes = {} , batch = {}".format(
    epochs_1,epochs_2,freeze_limit,nb_classes,batch))

filepath=ckpt+'/Model@epoch.{epoch:02d}-{val_loss:.2f}.hdf5'
checkpointer = keras.callbacks.ModelCheckpoint(filepath, monitor='val_loss', verbose=1, save_best_only=False, mode='auto')
early_stopping = EarlyStopping(monitor='val_loss',patience=15, verbose=1, mode='auto')
board = TensorBoard(log_dir='./logs', histogram_freq=0, batch_size=32, write_graph=True,
 write_grads=False, write_images=False, embeddings_freq=0, embeddings_layer_names=None, embeddings_metadata=None)

############### DIVISION ###############
if preprocess == 1:
   pi = pi_one
elif preprocess == 0:
   pi = None

if model_name in ["VGG19","VGG16","NASNet-m"]:
   inp_size = 224
elif model_name == "InceptionV3":
   inp_size = 299
elif model_name == "NASNet-l":
   inp_size = 331
else:
    print("Unkown Input Size")

print("[INFO] Creating Data Generators")

datagen_train = ImageDataGenerator(
    cval=255,
    rescale=1. / 255,
    rotation_range=10,
    width_shift_range=0.2,
    height_shift_range=0.2,
    zoom_range = [0.93, 1.3],
    fill_mode='nearest',
    horizontal_flip=True,
    vertical_flip=False,
    channel_shift_range=50,
    preprocessing_function=pi)

datagen_val = ImageDataGenerator(
    cval=255,
    rescale=1. / 255,
    rotation_range=10,
    width_shift_range=0.2,
    height_shift_range=0.2,
    zoom_range = [0.93, 1.3],
    fill_mode='nearest',
    horizontal_flip=True,
    vertical_flip=False,
    channel_shift_range=50,
    preprocessing_function=pi)

datagen_train = datagen_train.flow_from_directory(
        '{}/train'.format(str(args["data"])),
        target_size=(inp_size, inp_size),
        batch_size=batch,
        class_mode='categorical')

datagen_val = datagen_val.flow_from_directory(
        '{}/val'.format(str(args["data"])),
        target_size=(inp_size, inp_size),
        batch_size=batch,
        class_mode='categorical')

############### DIVISION ###############

print("[INFO] Class to Index Map")
print(datagen_train.class_indices)

############### DIVISION ###############

def get_class_weights(train_class_list):
        counts = []
        for i in range(len(train_class_list)):
            counts.append(len(os.listdir("{}/train/{}".format(str(args["data"]),train_class_list[i]))))
        majority = max(counts)
        return {cls_: float(majority / count) for cls_, count in enumerate(counts)}

############### DIVISION ###############

class_weights = get_class_weights(os.listdir('{}/train'.format(str(args["data"]))))

if climb == 0:
	class_weights = None
else:
	class_weights = class_weights
############### DIVISION ###############

nb_train_samples =  sum([len(files) for r, d, files in os.walk('{}/train'.format(str(args["data"])))])

nb_validation_samples =  sum([len(files) for r, d, files in os.walk('{}/val'.format(str(args["data"])))])

print("Number of Training Data = {}".format(nb_train_samples))

print("Number of Validation Data = {}".format(nb_validation_samples))

############### DIVISION ###############


print("[INFO] Creating Model Architecture")

base_model = list_model[model_name](weights='imagenet', include_top=False , input_shape = (inp_size,inp_size,3))
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu')(x) #512
x = Dropout(0.3)(x) # New Addition
predictions = Dense(nb_classes, kernel_regularizer=l2(.0005),activation='softmax')(x)

model = Model(inputs=base_model.input, outputs=predictions)

############### DIVISION ###############

print("[INFO] Freezing All Base Model Layers")

for layer in base_model.layers:
    layer.trainable = False

############### DIVISION ###############


model.compile(optimizer='adam', loss='categorical_crossentropy',metrics=[metrics.mae, metrics.categorical_accuracy])


############### DIVISION ###############


print("[INFO] Initiating Training")

t1 = time.time()

model.fit_generator(datagen_train,
                    steps_per_epoch=nb_train_samples//batch,
                    epochs= epochs_1,
                    verbose=1,
                    validation_data=datagen_val,
                    callbacks=[board , early_stopping],
                    validation_steps=nb_validation_samples // batch,
                    class_weight = class_weights)


############### DIVISION ###############


print("[INFO] Describing the Model")

for i, layer in enumerate(base_model.layers):
   print(i, layer.name)

############### DIVISION ###############


print("[INFO] Unfreezing {} Layers For BottleNeck Features".format(freeze_limit))

for layer in model.layers[:freeze_limit]:
   layer.trainable = False
for layer in model.layers[freeze_limit:]:
   layer.trainable = True


############### DIVISION ###############


from keras.optimizers import SGD
model.compile(optimizer=SGD(lr=0.001, momentum=0.9), loss='categorical_crossentropy',metrics=[metrics.mae, metrics.categorical_accuracy])

print("[INFO] Intiating Fine Tuning")

model.fit_generator(datagen_train,
                    steps_per_epoch=nb_train_samples//batch,
                    epochs= epochs_2,
                    verbose=1,
                    validation_data=datagen_val,
                    callbacks=[checkpointer , board , early_stopping],
                    validation_steps=nb_validation_samples // batch,
                    class_weight = class_weights)

t2 = time.time()

print("Total Time Taken in Training = {}".format(t2-t1))

############### THE END ###############