!pip install -q six numpy scipy matplotlib scikit-image opencv-python imageio
!pip install -q keras imgaug
!pip install -q tensorboardcolab
!pip install -q git+https://www.github.com/keras-team/keras-contrib.git

import os,sys,math
from matplotlib import pyplot as plt
%matplotlib inline
import numpy as np
import pandas as pd

import keras
from keras.models import Model, Sequential, load_model
from keras.layers import Dropout, Flatten, GlobalAveragePooling2D, AveragePooling2D, merge, Activation, Conv2D, MaxPooling2D, Input
from keras.layers import BatchNormalization, LeakyReLU, SeparableConv2D, Lambda, Concatenate,Softmax
from keras.optimizers import SGD, Adam
from keras.preprocessing.image import ImageDataGenerator
from keras.utils import np_utils
from keras.callbacks import ModelCheckpoint, LearningRateScheduler, ReduceLROnPlateau, CSVLogger, EarlyStopping

from keras_contrib.callbacks import CyclicLR
import imgaug as ia
from imgaug import augmenters as iaa
from tensorboardcolab import TensorBoardColab, TensorBoardColabCallback

# this part will prevent tensorflow to allocate all the avaliable GPU Memory
# backend
import tensorflow as tf
from keras import backend as k

# Don't pre-allocate memory; allocate as-needed
config = tf.ConfigProto()
config.gpu_options.allow_growth = True

# Create a session with the above options specified.
k.tensorflow_backend.set_session(tf.Session(config=config))

import google
colab_dir='./'
weights_dir='./'
file_name='EIP_CIFAR_10'
if hasattr(google,'colab'):
    from google.colab import drive
    drive.mount('/content/gdrive')
    colab_dir='/content/gdrive/My Drive/Colab Notebooks/'
    weights_dir=colab_dir+'weights/'
model_file=colab_dir+file_name+'.h5'

def create_dataset(image_dim=64,augment=True,batch_size=256):
  file_exists = os.path.exists("tiny-imagenet-200.zip")
  if file_exists:
    print("Files present. Skipping.")
  else:
    print(os.popen('wget http://cs231n.stanford.edu/tiny-imagenet-200.zip').read())
    print(os.popen('unzip -qq tiny-imagenet-200.zip').read())
    
  val_data = pd.read_csv('./tiny-imagenet-200/val/val_annotations.txt', sep='\t', header=None, names=['File', 'Class', 'X', 'Y', 'H', 'W'])
  val_data.drop(['X', 'Y', 'H', 'W'], axis=1, inplace=True)
  train_datagen = None
  
  if augment:
    train_datagen = ImageDataGenerator(rescale= 1./255,fill_mode='nearest',
      zoom_range = 0.2,horizontal_flip=True,shear_range=0.2,
      rotation_range=40,width_shift_range=0.2,height_shift_range=0.2,
      featurewise_std_normalization=True)
  else:
    train_datagen = ImageDataGenerator(rescale= 1./255,featurewise_std_normalization=True)
  
  valid_datagen = ImageDataGenerator(rescale=1./255)
  train_generator = train_datagen.flow_from_directory( r'./tiny-imagenet-200/train/', target_size=(image_dim, image_dim), color_mode='rgb',
                                                    batch_size=batch_size, class_mode='categorical', shuffle=True, interpolation='bicubic')
  
  validation_generator = valid_datagen.flow_from_dataframe(val_data, directory='./tiny-imagenet-200/val/images/', x_col='File', y_col='Class', target_size=(image_dim, image_dim),
                                                    color_mode='rgb', class_mode='categorical', batch_size=batch_size, shuffle=True,interpolation='bicubic')
  #x_batch, y_val_batch = next(train_generator)
  #print('Training batch shape:',x_batch.shape)
  
  return train_generator, validation_generator#, x_batch, y_val_batch

#Augmentation and resizing
#augment and then concat samples with original
def augment(dataset,flip=0.5,blur=1.0):
    ia.seed(1)
    seq = iaa.Sequential([iaa.Crop(px=(0, 16)),iaa.Fliplr(0.5),iaa.GaussianBlur(sigma=(0, 3.0))], random_order=True)   
    return seq.augment_images(dataset)

def augmenter(X,y,start=0,end=1,append=False):
    ln=len(X)
    print('Before augmentation:',X.shape,y.shape)
    start=int(start*ln)
    end=int(end*ln)
    if append:
      new_X=augment(X)[start:end]
      new_y=y[start:end]
      X=np.concatenate((X,new_X))
      y=np.concatenate((y,new_y))
    else:
      X=augment(X)
    print('After augmentation:',X.shape,y.shape)
    return (X,y)

#26x26 is almost half of 32x32. 22x22 maybe too small even though its exact half.
def resize_imgs(imgs,shape=(26,26)):
    seq = iaa.Sequential([iaa.Scale({"height": shape[0], "width": shape[1]})])
    return seq.augment_images(imgs)

modval=50
max_lr=0.1

def cosine_anneal_schedule(t,lr):
  global modval,max_lr
  modval=int(modval)
  cos_inner = (np.pi * (t % modval))/ modval
  cos_out = np.cos(cos_inner) + 1
  return float(max_lr / 2 * cos_out)

def load_prev_model(model_small=None):
#     print(os.popen('du -sh '+colab_dir+'*').read())
    last_best=os.popen('du -sh '+colab_dir+r'*|sed -r "s/^[0-9\.]+[MK]?\s(.*)$/\1/g;"|tail -1').read().strip()
    model_prev=None
    try:
        print('Attempting to load last best model from file - ',end='')
        last_best_fn=last_best.split('/')[-1]
        last_best_epoch=last_best_fn.split('.')[1].split('-')[0]
        model_prev=load_model(last_best)
        print('Sucess\nLoaded '+last_best_fn, 'epoch :', last_best_epoch)
    except Exception as e:
        print('Failed!\n',e)
        try:
            print('Attempting to load last saved model from file - ',end='')
            model_prev = load_model(model_file)
            print('Sucess\nLoaded model from file',model_file)
        except Exception as e:
    #         print(str(e), 'at line ', sys.exc_info()[2].tb_lineno)
            print('Failed!\n',e)
            try:
                print('Attempting to load in memory, small model - ',end='')
                if len(model_small.layers)>1:
                    model_prev=model_small
                print('Sucess')
            except Exception as e:
                print('Failed!\n',e)
    return model_prev
  
def copy_weights(model_to,model_from):
    #model_to.set_weights(model_from.get_weights())
    s,err=0,0
    print('Trying to copy weights')
    try:
        if len(model_to.layers) >1 and len(model_from.layers) >1:
            pass
    except Exception as e:
        print('Inavlid models',model_to,model_from)
        return
    for new_layer, layer in zip(model_to.layers[1:], model_from.layers[1:]):
        s+=1
        try:
            new_layer.set_weights(layer.get_weights())
        except Exception as e:
            pass
    print('Done: errors:',err)
      
#new model with larger i/p layer
def larger_model(src_model,shape=(32,32,3)):
    new_input=Input(shape)
    src_model.layers.pop(0)
    new_output=src_model(new_input)
    new_model=Model(new_input,new_output)
    return new_model

# Load CIFAR10 Data
def load_data(resize=False,append=False,shape=(26,26),train_augment=True,test_augment=False):
    (x_train, y_train), (x_test, y_test) = cifar10.load_data()
    
    y_train = keras.utils.to_categorical(y_train, num_classes)
    y_test = keras.utils.to_categorical(y_test, num_classes)
    
    if resize:
        x_train=resize_imgs(x_train,shape)
        x_test=resize_imgs(x_test,shape)
    if train_augment:
      (x_train, y_train) = augmenter(x_train, y_train,end=1,append=append)
    if test_augment:
        (x_test, y_test) = augmenter(x_test, y_test,end=1,append=append)
    return (x_train, y_train,x_test, y_test)
#create a dnn model
def create_model(input_shape,num_layers,input_conv_filters=12,input=None):
    print('Creating model with input shape',input_shape)
    if input is None:
        input = Input(input_shape)
    First_Conv2D = Conv2D(input_conv_filters, (3,3), use_bias=False ,padding='same')(input)
    hidden_dense_blocks = dense_units_chain(n_dense_blocks,First_Conv2D,num_filter,dropout_rate,num_layers)
    Last_Block = add_denseblock(hidden_dense_blocks, num_filter, dropout_rate)
    output = output_layer(Last_Block)
    model = Model(inputs=[input], outputs=[output])
    return model

sc=False
after_transistion = True
# Dense Block
def add_denseblock(input,num_layers, growth_rate = 12, dropout_rate = 0.2, bn_layers=48, middle_layers=24):
    
    temp = input
    
    #if len(Block_outputs)>1:
    #  bv = Block_outputs[-2]
    #  #temp = Concatenate(axis=-1)([Block_outputs[-2],[temp]])
    #  print(bv)
    
    for _ in range(num_layers):
        ######################################################################################################################################
        BatchNorm = BatchNormalization()(temp)
        relu = Activation('relu')(BatchNorm)        
        Conv2D_1_1_bottleneck = Conv2D(bn_layers, (1,1), use_bias=False ,padding='same',kernel_initializer='he_normal')(relu) #bottleneck  <===
        ###################################################################################################################################### |
        #BatchNorm = BatchNormalization()(Conv2D_1_1_bottleneck)
        #relu = Activation('relu')(BatchNorm)        
        #Conv2D_1_1_bottleneck = Conv2D(bn_layers, (3,3), use_bias=False ,padding='same',kernel_initializer='he_normal')(relu) #bottleneck  <===
        ###################################################################################################################################### |
        #BatchNorm = BatchNormalization()(Conv2D_1_1_bottleneck) #extra layer                                                                 # |
        #relu = Activation('relu')(BatchNorm) #extra layer                                                                                    # |
        #Conv2D_1_1 = Conv2D(middle_layers, (1,1), use_bias=False ,padding='same',kernel_initializer='he_normal')(relu) #extra layer          # |
        ###################################################################################################################################### |
        BatchNorm = BatchNormalization()(Conv2D_1_1_bottleneck) #bottleneck   <================================================================
        relu = Activation('relu')(BatchNorm)
        Conv2D_3_3 = Conv2D(growth_rate, (3,3), use_bias=False ,padding='same',kernel_initializer='he_normal')(relu)
        ######################################################################################################################################|
        
        if dropout_rate>0:
          Conv2D_3_3 = Dropout(dropout_rate)(Conv2D_3_3)
        concat = Concatenate(axis=-1)([temp,Conv2D_3_3])
        
        temp = concat
    #if not after_transistion:
    #  Block_outputs.append(temp)
    return temp

def add_transition(input, num_filter = 12, dropout_rate = 0.2):
    BatchNorm = BatchNormalization()(input)
    relu = Activation('relu')(BatchNorm)
    Conv2D_BottleNeck = Conv2D(int(num_filter), (1,1), use_bias=False ,padding='same')(relu)
    if dropout_rate>0:
      Conv2D_BottleNeck = Dropout(dropout_rate)(Conv2D_BottleNeck)
    avg = AveragePooling2D(pool_size=(2,2),padding='same')(Conv2D_BottleNeck)
    #if after_transistion:
    #  Block_outputs.append(avg)
    return avg

  
def output_layer(input,num_classes,mf=2,nl=2):
    inp = input
    for i in range(nl,-1,-1):
      classes = int(num_classes*(mf**i))
      BatchNorm = BatchNormalization()(inp)
      relu = Activation('relu')(BatchNorm)
      conv1_1 = Conv2D(classes,(1,1),use_bias=False,kernel_initializer='he_normal',padding='same')(relu)
      conv3_3 = Conv2D(classes,(3,3),use_bias=False,kernel_initializer='he_normal',padding='same')(conv1_1)
      inp = conv3_3
    
    GlobalAveragePooling = GlobalAveragePooling2D(data_format='channels_last')(inp)
    sf = Softmax()(GlobalAveragePooling)
    return sf

def custom_model(input_shape,initial_layers,num_layers,db_arch,dropout_rate=0,trans_2k=False,n_blocks=3,omf=2,onl=2):
  Block,Transition = None,None
  bn_layers,middle_layers,growth_rate = db_arch
  
  input = Input(shape=input_shape)
  First_Conv2D = Conv2D(initial_layers, (3,3), use_bias=False ,padding='same')(input)
  num_filter=initial_layers
  
  input_layer = First_Conv2D
  for i in range(n_blocks):
    Block = add_denseblock(input_layer,num_layers, growth_rate, dropout_rate,bn_layers,middle_layers)
    
    num_filter+=growth_rate*num_layers
    num_filter=(growth_rate*2) if trans_2k else num_filter//2    
    
    Transition = add_transition(Block, num_filter, dropout_rate)    
    input_layer = Transition
  
  #Last_Block = add_denseblock(Transition, num_layers, growth_rate, dropout_rate,bn_layers,middle_layers)
  output = output_layer(Block,num_classes,omf,onl)
  model = Model(inputs=[input], outputs=[output])
  return model

print('====================HYPER PARAMETERS====================')

num_filter = 0 # add it up as you go
compression = 0.5
dropout_rate =0#0.2#.5
num_classes = 200
train_ds_size=100000
validation_ds_size=10000

modval=20
max_lr=0.3
image_dim = 32
batch_size = 64
inp_shape = (image_dim,image_dim,3)

n_blocks=3
num_layers=10#26
initial_layers=24
db_archs=[[128,64,32],[96,48,24],[48,24,12]]
dbi = 1
db_arch=db_archs[dbi]

hist=[]
epochs=initial_epoch=0
Block_outputs = []

model = custom_model(inp_shape,initial_layers,num_layers,db_arch,n_blocks=n_blocks,trans_2k=False,omf=2,onl=0)

tbc=TensorBoardColabCallback(TensorBoardColab())
hist=[]

model.compile(loss='categorical_crossentropy',metrics=['accuracy'],optimizer=SGD(lr=0.1, decay=1e-6, momentum=0.9, nesterov=True))
model.summary()

# from keras.datasets import cifar10
# x_train, y_train, x_test, y_test = load_data(resize=False,append=False,shape=(image_dim,image_dim))

#callbacks
batch_size=64#128
modval=8
max_lr=0.3
augment_data = False
train_generator, validation_generator = create_dataset(image_dim=image_dim,augment=augment_data,batch_size=batch_size)


model_checkpointer=ModelCheckpoint(weights_dir+'large_weights.{epoch:02d}-{val_acc:.2f}.h5', monitor='val_loss',verbose=1, save_best_only=True, save_weights_only=False, mode='min', period=5)
clr = CyclicLR(base_lr=0.1, max_lr=0.2,step_size=8*(int(train_ds_size/batch_size)))
els = EarlyStopping(monitor='val_loss', min_delta=0.001, patience=10, verbose=1, mode='auto', baseline=None, restore_best_weights=False)

use_cosine_annealing=False
lrs = ReduceLROnPlateau(monitor='val_loss', factor=0.2,patience=5, min_lr=0.001)
lrs=LearningRateScheduler(cosine_anneal_schedule) if use_cosine_annealing else clr
callbacks = [model_checkpointer,lrs,els,tbc]
initial_epoch=epochs
epochs=10
epochs+=initial_epoch

#h=model.fit(x_train, y_train, batch_size=batch_size, verbose=1,initial_epoch=initial_epoch,epochs=epochs, callbacks=callbacks, validation_data=(x_test,y_test))
h =model.fit_generator(train_generator,steps_per_epoch=int(train_ds_size/batch_size),verbose=1,epochs=epochs,callbacks=callbacks,validation_data=validation_generator,validation_steps=int(validation_ds_size/batch_size))
hist.append(h)
model.save(model_file)
print("Saved model to disk")

#callbacks
batch_size=64#128
modval=8
max_lr=0.3
augment_data = True
train_generator, validation_generator = create_dataset(image_dim=image_dim,augment=augment_data,batch_size=batch_size)


model_checkpointer=ModelCheckpoint(weights_dir+'large_weights.{epoch:02d}-{val_acc:.2f}.h5', monitor='val_loss',verbose=1, save_best_only=True, save_weights_only=False, mode='min', period=5)
clr = CyclicLR(base_lr=0.1, max_lr=0.2,step_size=8*(int(train_ds_size/batch_size)))
els = EarlyStopping(monitor='val_loss', min_delta=0.001, patience=10, verbose=1, mode='auto', baseline=None, restore_best_weights=False)

use_cosine_annealing=False
lrs = ReduceLROnPlateau(monitor='val_loss', factor=0.2,patience=5, min_lr=0.001)
lrs=LearningRateScheduler(cosine_anneal_schedule) if use_cosine_annealing else clr
callbacks = [model_checkpointer,lrs,els,tbc]
initial_epoch=epochs
epochs=20#24
epochs+=initial_epoch

#h=model.fit(x_train, y_train, batch_size=batch_size, verbose=1,initial_epoch=initial_epoch,epochs=epochs, callbacks=callbacks, validation_data=(x_test,y_test))
h =model.fit_generator(train_generator,steps_per_epoch=int(train_ds_size/batch_size),verbose=1,epochs=epochs,callbacks=callbacks,validation_data=validation_generator,validation_steps=int(validation_ds_size/batch_size))
hist.append(h)
model.save(model_file)
print("Saved model to disk")

print('====================END OF LARGER MODEL====================')