count.py

from random import random
import torch
import torchvision.models as models
import torch.nn as nn
from torchcontrib.optim import SWA
from torchsummary import summary
from torch.autograd import Variable
import torch.backends.cudnn as cudnn
from collections import OrderedDict
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from networkx.drawing.nx_agraph import graphviz_layout, to_agraph
import pygraphviz as pgv
import json
import copy
import numpy as np
import pickle
import os
import argparse


from modelsummary import torch_summary
from utils import Utils
from calc_tam_layer import *
from operations import * # operations (layers) in pytorch
from zero_cost_proxy import *
from processify import processify
from prunning import *
#import torch.multiprocessing as mp
#mp.set_start_method('spawn')


import time
import datetime
import sys
import traceback
import gc
from operator import itemgetter
#from memory_profiler import profile




#----------------------------------------------------------------------------#
''' ########################### '''

#----------------------------------------------------------------------------#

def generateDictStateTransitionGraph(modelSummary, printDictGraph=False):
    """Generate Markov Chain

    Arguments:
        modelSummary {[dict]} -- [Information about layer transition - Summary]

    Keyword Arguments:
        printDictGraph {bool} -- [If True, print markov chain at the end] (default: {False})

    Returns:
        [dict of dicts] -- [Markov Chain]
    """    
    
    print("#### Generating State Transition Graph [...] ####")

    # Count the frequencies
    stateTransitionFreq = {}
    predecessorLayer = "START"
    for index, row in modelSummary.iterrows():
        # Get layer name
        layerName = row['Layer'].split("-")[0]
        # If dict does not have last layer, add empty dict
        if predecessorLayer not in stateTransitionFreq.keys(): 
            stateTransitionFreq[predecessorLayer] = {}

        # Add layer transition
        if layerName not in stateTransitionFreq[predecessorLayer]:
            stateTransitionFreq[predecessorLayer][layerName] = 1
        else:
            stateTransitionFreq[predecessorLayer][layerName] = stateTransitionFreq[predecessorLayer][layerName]+1
        predecessorLayer = layerName

    # To add STOP node (previous->STOP)
    if predecessorLayer not in stateTransitionFreq.keys():
        stateTransitionFreq[predecessorLayer] = {}
    if 'END' not in stateTransitionFreq[predecessorLayer]:
        stateTransitionFreq[predecessorLayer]['END'] = 1
    else:
        stateTransitionFreq[predecessorLayer]['END'] = stateTransitionFreq[predecessorLayer]['END']+1

    if (printDictGraph):
        print("#-----------------------------#\n")
        print("---- Graph in the form of Dict ----")
        print (stateTransitionFreq)

    return (stateTransitionFreq)


def getLayerComponents(layerName):
    if ("Conv" in layerName):
        return ['out_channels', 'kernel_size']
    if (("AvgPool" in layerName) or ("MaxPool" in layerName)) and not ("Adaptive" in layerName):
        return ['kernel_size']
    if ("Linear" in layerName):
        return ['out_features']
    if ("Dropout" in layerName):
        return ['p']
    if ("Adaptive" in layerName):
        return ['output_size']
    if ("Bottleneck" in layerName):
        return ['out_channels']
    else:
        return ['None']


def generateHiddenMarkovStates(modelSummary, printHiddenStates=False):
    """Generate hidden states for each possible node in the markov chain, e.g.:
    conv2d -> (1*1:1,3*3:1,5*5:3,7*7:5)

    Returns:
        [dict of dict of dict] -- [Hidden Markov States that represent each layer and the frequency of components on it]
    """
    print("#### Generating Hidden Markov States [...] ####")

    hiddenMarkovStates = {}
    try:
        
        for index, row in modelSummary.iterrows():
            # Get layer name
            layerName = row['Layer'].split("-")[0]

            # If layer is not on the dict yet
            if layerName not in hiddenMarkovStates.keys(): 
                hiddenMarkovStates[layerName] = {}
            
            layerComponents = getLayerComponents(layerName)
            # Some layers don't require components (e.g. -> ReLU, ELU, ..)
            if layerComponents[0] == "None":
                #print (layerName)
                hiddenMarkovStates[layerName] = None
                continue
            for component in layerComponents: # Iterate over all possible components for given layer
                # Add component if not present
                if component not in hiddenMarkovStates[layerName]:
                    hiddenMarkovStates[layerName][component] = {}

                value = row[component] # Get value of the component
                # Check if value of the component is present in the dict, e.g: kernel_size: (3,3)
                #                                                               component   value
                if type(value) is list:
                    value = tuple(value)
                if str(value) not in hiddenMarkovStates[layerName][component].keys():
                    hiddenMarkovStates[layerName][component][value] = 1
                else:
                    hiddenMarkovStates[layerName][component][value] = hiddenMarkovStates[layerName][component][value]+1
    except Exception:
        print(traceback.format_exc())

    #print (type(list(hiddenMarkovStates['Conv2d']['kernel_size'].keys())[0]))
    return (hiddenMarkovStates)


#----------------------------------------------------------------------------#
def generateNetworkxDiGraph(stateTransitionFreq):
    '''
    Generate a Direct Graph in NetworkX from a ordereddict of orderedicts
    that contains the state transitions (and the number of occurences)
    '''
    ## Generate NetworkX graph and plot it
    G = nx.DiGraph(directed=True)
    G.edges.data('weight', default=1)
    # Add edges with weights
    for key in stateTransitionFreq:
        innerFrequencies = sum(stateTransitionFreq[key].values())
        #print(f'{key}  ->  {stateTransitionFreq[key]}   : {innerFrequencies}')
        G.add_node(key)
        for innerKey in stateTransitionFreq[key]:
            G.add_edge(key, innerKey, weight="%.2f"%(stateTransitionFreq[key][innerKey]/innerFrequencies))
            #print (f'i:{innerKey}, iv:{stateTransitionFreq[key][innerKey]}')
    return (G)


def generateGraphVizRepresentation(G, filename="filename.png"):
    print("\n#### Generating GraphViz Representation [...] ####")
    #nx.set_node_attributes(G, {k: {'label': labels[k]} for k in labels.keys()})
    nx.set_edge_attributes(G, {(e[0], e[1]): {'label': e[2]['weight']} for e in G.edges(data=True)})
    A = nx.drawing.nx_agraph.to_agraph(G)
    #print(A) # Print the graph structure
    A.layout('dot')
    A.draw(filename)

#----------------------------------------------------------------------------#
#-----------------------------Model Storing----------------------------------#
def writeDictPickle(modelDict, filename="./graphs/test.json"):
    print("\n#### Storing Dict as Pickle [...] ####")
    with open(filename, 'wb') as f:
        pickle.dump(modelDict, f)

def loadGraphDictPickle(filename="./graphs/test.json"):
    print("\n#### Reading Graph Dict Pickle [...] ####")
    with open(filename, 'rb') as f:
        loaded_dictionaries = pickle.load(f)
    return (loaded_dictionaries)

def writeGraphDict(modelDict, filename="./graphs/test.json"):
    print("\n#### Storing Dict as JSON [...] ####")
    with open(filename, 'w') as f:
        f.write(json.dumps(modelDict))

def loadGraphDict(filename="./graphs/test.json"):
    print("\n#### Reading Dict [...] ####")
    with open(filename, 'r') as read_file:
        loaded_dictionaries = json.loads(read_file.read())
    return (loaded_dictionaries)
#----------------------------------------------------------------------------#
#-----------------------------Create plots-----------------------------------#

def linePlot(xAxis, yAxis, labels, xLabel = "Epoch", yLabel ="Accuracy (%)", title=None, savePlot=False, fileName = "plot.pdf"):
    colors = ["blue", "red", "brown", "orange", "grey", "olive", "cyan", "purple", "pink"]
    plt.clf()

    if xAxis is not None:
        for index, x, y in enumerate(yAxis):
            plt.plot(xAxis, y, label=labels[index], color=colors[index])
    else:
        for index, y in enumerate(yAxis):
            plt.plot(y, label=labels[index], color=colors[index])

    plt.xlabel(xLabel)
    plt.ylabel(yLabel)
    if title is not None:
        plt.title(title)
    plt.legend()
    #plt.show()

    if savePlot:
        plt.savefig(fileName, bbox_inches='tight')


#----------------------------------------------------------------------------#
def readPopFromJSONs(path="./graphs/", args=None):
    torchModels = {}

    jsonFiles = [pos_json for pos_json in os.listdir(path) if pos_json.endswith('.json')]
    for fileName in jsonFiles:
        modelName = fileName[:-5]
        torchModels[modelName] = OrderedDict()
        torchModels[modelName] = loadGraphDictPickle(path+fileName)

        if args.mixed_training: #combination
            torchModels[modelName]['fitness'] = args.mixed_fitness_lambda*torchModels[modelName]['score'] + (1-args.mixed_fitness_lambda)*torchModels[modelName]['valacc']
        elif args.without_training: # without train fitness
            torchModels[modelName]['fitness'] = torchModels[modelName]['score']
        else: # regular partial fitness #if args.without_training == False and args.mixed_fitness==False: 
            torchModels[modelName]['fitness'] = torchModels[modelName]['valacc']

    return (torchModels)


def generateSummaryAndGraph(model, img_shape = (3,244,244), automatedModel=True, input_image=0):
    try:
        modelSummary = torch_summary(img_shape, model, automatedModel, input_image)
        #print(modelSummary)
    except Exception as e:
        print (f'Exception generating model summary: {e}')
        return (None)
    
    # Because not every layer has the same columns (some may appear only once, others all the time)
    columns = set(['Layer'])
    # Add Column named - layer
    for key, value in modelSummary.items():
        colsAux = set(value.keys())
        columns = columns.union(colsAux)
    #print (columns)

    # Pass all values to a list in order to create a pandas df
    values = []
    for key,value in modelSummary.items():
        value['Layer'] = key
        values.append(value)
    df = pd.DataFrame(values, columns=columns, dtype=object) #[value for key,value in modelSummary.items()], columns = ms.keys()
    df = df.reindex(columns=(['Layer'] + list([a for a in df.columns if a != 'Layer']) )) #Layer should be the first column

    #print(df)
    #print (f'df: {type(df)}')

    stateTransitionFreq = generateDictStateTransitionGraph(df)
    hiddenMarkovStates  = generateHiddenMarkovStates(df)


    return stateTransitionFreq, hiddenMarkovStates # MarkovChain; Hidden States


def generateGraphsOfTorchVisionModels(utils, evaluateFitnes = True, data_dir = None, dataset = 'cifar10', datasetType='partial', 
        generateGraphVisualRepresentation=True, storeModelToFile=True, args=None, path="./models_pytorch/cifar10/partial/",
        trainloader=None, valloader=None):

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Get all pytorch models name
    model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

    torchModels = {}
    for modelName in model_names:
        print(f"\n---- {modelName} [...] ----")
        print(modelName)

        # Get pytorch model

        try:
            #model = models.__dict__[modelName](pretrained=True).to(device)
            model = models.__dict__[modelName]().to(device)
        except:
            continue
        # Generate the stateTransition graph of the model and get the graph
        try:
            graph, hiddenStates = generateSummaryAndGraph(model, img_shape=(3,244,244), automatedModel=False)
            if graph is None: # some pytorch models can eventually produce errors
                continue
        except:
            print ("Model: "+modelName+", error creating Markov Chain.")
            continue

        torchModels[modelName] = OrderedDict()

        #stateTransitionFreq = generateDictStateTransitionGraph(modelSummary)
        torchModels[modelName]['graph'] = graph
        torchModels[modelName]['hiddenStates'] = hiddenStates
        

        # Get fitness of pytorch models
        fitness, valacc, score = 0, 0, 0
        # Get Criterion
        if (evaluateFitnes):
            fitness, valacc, score = getFitness(args, utils, model, trainloader, valloader, device, generated=False)
            '''
            criterion = utils.getCrossEntropyLoss()
            optimizer = utils.getSGDOptimizer(model)
            try:
                # Get dataloaders for the dataset
                trainloader, valloader, _ = utils.get_train_dataset(args, data_dir=data_dir, \
                                                            dataset=dataset, datasetType=datasetType, resize=True)

                print("#### Starting Training and Fitness Extraction [...] ####")
                for i in range(0,args.search_epochs): # Train for n epochs, evaluate
                    model, acc, _ = utils.trainNormalModels(model, trainloader, criterion, optimizer)
                    valacc, _ = utils.evaluateNormalModels(model, valloader, criterion)
            except Exception as e:
                print(e)
        #print (valacc)
        torchModels[modelName]['fitness'] = valacc
        print(f'Final Val Accuracy: {valacc}')


        # Get NAS without training score
        try:
            if args.without_training==True:
                model = models.__dict__[modelName]()
                model = utils.change_CNN_classifier(model, modelName, 1) #change for 1 class output
                model = model.to(device)
                bs = args.batch_size
                args.batch_size = 8 # change batch_size for score calculation
                # Get dataloader with required batch_size 
                trainloader, _, _ = utils.get_train_dataset(args, data_dir=data_dir, \
                                                        dataset=dataset, datasetType=datasetType, resize=True)
                torchModels[modelName]['score'] = zero_cost_proxy(trainloader, model, device, args.batch_size)
                print(f'Score: {torchModels[modelName]["score"]}')

                args.batch_size = bs # change batch_size for the original
        except Exception as e:
            print(e)
            pass
        '''
        torchModels[modelName]['valacc'] = valacc
        torchModels[modelName]['score'] = score
        torchModels[modelName]['fitness'] = fitness

        #torchModels[modelName]['model'] = model #copy.deepcopy(model.state_dict()) # Store the model
                                                # to rebuild it if needed

        if(generateGraphVisualRepresentation):
            modelGraph = generateNetworkxDiGraph(graph)
            generateGraphVizRepresentation(modelGraph, filename="./graphs/"+modelName+".png")
        
        if(storeModelToFile):
            writeDictPickle(torchModels[modelName], filename=path+modelName+".json")

        # The model itself is no longer needed
        del model
        #torch.cuda.empty_cache()
    return (torchModels)


def getFitness(args, utils, model, trainloader, valloader, device, generated=True):
    criterion = utils.getCrossEntropyLoss()
    optimizer = utils.getSGDOptimizer(model)

    fitness, val, score =0,0,0 #np.random.randint(100)
    # Acquire fitness
    try:
        # Acquire score that represents model fitness without requiring training
        if args.without_training or args.mixed_training:
            print("### Without Training [...] ###")
            score = zero_cost_proxy(trainloader, model, device, args.batch_size, desired_size=256, generated=generated)
    except Exception as e:
        score = 0
        print(e)
        if "out of memory" not in str(e):
            print(e)
            #print(modelLayersStringList)
            print(model)
        utils.clean_memory()
    try:
        # Train for some epochs to acquire validation accuracy of the generated model
        if args.mixed_training or (args.without_training == False and args.mixed_training == False):
            print("### Regular Training [...] ###")
            optimizer = utils.getSGDOptimizer(model)
            if generated: #the network is created using a modulelist
                for i in range(0, args.search_epochs): # Train for n epochs, evaluate
                    model, _, _, _ = utils.train(model, trainloader, criterion, optimizer)
                    top1, _, _ = utils.evaluate(model, valloader, criterion)
                    if (top1.avg > fitness):
                        val = top1.avg # store the best val accuracy (fitness)
            else: 
                for i in range(0, args.search_epochs): # Train for n epochs, evaluate
                    model, _, _ = utils.trainNormalModels(model, trainloader, criterion, optimizer)
                    valacc, _ = utils.evaluateNormalModels(model, valloader, criterion)
                    if (valacc > fitness):
                        val = valacc # store the best val accuracy (fitness)
    except Exception as e:
        val = 0
        print(e)
        if "out of memory" not in str(e):
            #print(modelLayersStringList)
            print(model)
        utils.clean_memory()

    if args.mixed_training: #combination
        fitness = args.mixed_fitness_lambda*score + (1-args.mixed_fitness_lambda)*val
    elif args.without_training: # without train fitness
        fitness = score
    else: # regular partial fitness #if args.without_training == False and args.mixed_fitness==False: 
        fitness = val

    return fitness, val, score


#----------------------------------------------------------------------------#
def checkIfLinearInProximity(list):
    for layer in list:
        if('Linear' in layer):
            return (True)
        if ("ReLU" in layer) or ("ELU" in layer) or ("SELU" in layer) or ("ReLU6" in layer) or ("Dropout" in layer):
            continue
        elif ("Conv2d" in layer) or ("AvgPool" in layer) or ("MaxPool" in layer) or ("BatchNorm2d" in layer):
            return (False)
    return (True) #is the last layer

def checkIfLinearInProximityReverse(list):
    list.reverse()
    for layer in list:
        if('Linear' in layer):
            return (True)
        if ("ReLU" in layer) or ("ELU" in layer) or ("SELU" in layer) or ("ReLU6" in layer) or ("Dropout" in layer):
            continue
        elif ("Conv2d" in layer) or ("AvgPool" in layer) or ("MaxPool" in layer) or ("BatchNorm2d" in layer):
            return (False)
    return (True) #is the last layer

def transformNetworkIntoPytorch(modelList, input_shape=[3,32,32], maxneurons=6000, n_classes=10):
    img_size = input_shape[1:]
    input_channels = input_shape[0]
    last_conv_output = input_shape[0]

    layers = nn.ModuleList()
    # Transform strings into pytorch layers
    #print (modelList)
    for idx, layerInformation in enumerate(modelList):
        stringLayer = layerInformation[0]
        components = layerInformation[1]

        if ("Conv" in stringLayer) and not ("Transpose" in stringLayer): #TODO: padding should be fixed using the calc_tam_layers, so the images does not get cropped
            out_channels = components['out_channels']
            kernel_size  = components['kernel_size']
            padding = Conv_Padding(kernel_size=kernel_size, img_sizes=img_size)
            layer = OPS[stringLayer](Cin = last_conv_output, Cout=out_channels, kernel_size=kernel_size, padding=padding)
            #layer = transformLayerNameIntoPytorch(stringLayer, in_channels=last_conv_output, out_channels=out_channels, kernel_size=kernel_size)
            last_conv_output = out_channels

        elif "Linear" in stringLayer:
            try:
                #if checkIfLinearInProximityReverse(modelList[:idx-1]):
                aux = modelList[:idx]
                aux.reverse()
                if checkIfLinearInProximity(aux):
                    in_features = out_features
                else:
                    in_features = Flatten(img_size = img_size, channels = last_conv_output)
            except:
                in_features = Flatten(img_size = img_size, channels = last_conv_output)
            keepFlatten = checkIfLinearInProximity(modelList[idx+1:])
            
            out_features = components['out_features'] if keepFlatten == True else in_features

            layer = OPS[stringLayer](in_features=in_features, out_features=out_features, unflatten=(not keepFlatten))

        elif "BatchNorm" in stringLayer:
            #layer = transformLayerNameIntoPytorch(stringLayer, num_features=last_conv_output)
            layer = OPS[stringLayer](last_conv_output)

        elif "Adaptive" in stringLayer:
            #print (components)
            out_features = components['output_size'] if components['output_size'] is not None else img_size
            #print (out_features)

            # if image size would be less than (3,..), 
            if (any(i < 3 for i in out_features)):
                out_features = img_size
            img_size = out_features
            
            layer = OPS[stringLayer](out_features=out_features)

        elif (("AvgPool" in stringLayer) or ("MaxPool" in stringLayer)) and not ("Adaptive" in stringLayer):
            #print ((components['kernel_size']))
            if isinstance(components['kernel_size'], int):
                kernel_size = components['kernel_size'] if all(i > components['kernel_size'] for i in img_size) else 1
                img_size = [int(elem/kernel_size) for elem in img_size] 
            else:
                kernel_size = components['kernel_size'] if all(i > c for c,i in zip(components['kernel_size'], img_size)) else 1
                # reduze img size accordingly
                img_size = [int(elem/k) for k,elem in zip(kernel_size, img_size)] 
            layer = OPS[stringLayer](kernel_size=kernel_size)                

        elif ("Dropout" in stringLayer): 
            dropout_prob = components['p']
            layer = OPS[stringLayer](p=dropout_prob)

        elif ("ReLU" in stringLayer) or ("ELU" in stringLayer) or ("SELU" in stringLayer):
            layer = OPS[stringLayer](False)

        elif ("Bottleneck" in stringLayer):
            out_channels = components['out_channels']
            stride = components['stride']
            # If the bottleneck does not downsample, out_channels must be equal to the input
            if stride == 1:
                out_channels, components['out_channels'] = int(last_conv_output/Bottleneck.expansion), int(last_conv_output/Bottleneck.expansion)
            layer = OPS[stringLayer](last_conv_output, out_channels, stride, None, 1, 64, 1, None)
            last_conv_output = layer.output
            img_size = [int(elem/stride) for elem in img_size] 

        else: # It will return None, for debug purposes
            print(stringLayer)
            layer = OPS[stringLayer]

        layers += [layer]

    iterator = 1
    while 1:
        layer = layers[-iterator]
        if isinstance(layer, LinearLayer):
            #print(f'img_size:{img_size}')
            in_features = layer.op.in_features
            #print(f'in_features:{in_features}')

            # Check if any of the last linear layers has too many in_features
            # That may cause overfitting
            if (in_features > maxneurons): # if too many neurons, try other model
                #print("oleeeeeeeeeee")
                return None
                
        elif not (isinstance(layer,nn.Dropout) or
            isinstance(layer,nn.ELU) or
            isinstance(layer,nn.ReLU) or
            isinstance(layer,nn.ReLU6) or
            isinstance(layer,nn.SELU)):
            break # there is no more MLP
        iterator = iterator+1

    # Adjust last linear layer for output_size = n_classes
    if isinstance(layers[-1], LinearLayer):
        if layers[-1].op.out_features != n_classes:
            layers[-1] = LinearLayer(in_features=layers[-1].op.in_features, out_features=n_classes)

    #model = nn.Sequential(*layers) #nn.ModuleList
    #print(layers)
    return (layers)

#----------------------------------------------------------------------------#
#-------------------------------Evolution------------------------------------#
def modelWeightedRouletteSelection(models, lastLayer):
    #Copy the dict, because dict parameters are references
    modelsDict = models.copy()
    ###print (f'Copied model: {modelsDict.keys()}')
    for k, v in list(modelsDict.items()):
        #print(f'key={k},  {v}')
        # Remove models that don't have HMC, the last layer or any fitness
        try:
            if ('graph' not in v) or not (lastLayer in v['graph']) or (v['fitness'] == 0):
                del modelsDict[k]
        except:
            del modelsDict[k]
    sortedPopulation = sorted(modelsDict.items(), key=lambda x: x[1]['fitness'])
    sortedPopulationNames = [i[0] for i in sortedPopulation] # Get only the names

    # Get rank by fitness of the previous population
    for key in modelsDict.keys():
        modelsDict[key]['rank'] = sortedPopulationNames.index(key)

    ###print (f'Copied model after del: {modelsDict.keys()}')
    max = sum(models[key]['rank'] for key in modelsDict) #'fitness'

    pick = np.random.uniform(0, max)
    current = 0
    for key in modelsDict:
        current += modelsDict[key]['rank'] #'fitness'
        if current > pick:
            return key

def layerSelectionFromModelGraph(graph, layer):
    max = sum(graph.values())
    pick = np.random.uniform(0, max)
    current = 0
    for key in graph:
        ###print(f'Key:{key}, value:{graph[key]}')
        current += graph[key]
        if current > pick:
            return key

def componentSelectionFromModelLayer(hiddenStates, component):
    """[summary]

    Arguments:
        hiddenStates {[dict]} -- [Has all the values of frequency for a given layer of the model]
        component {[String]} -- [Which component is to be used]
    """
    max = sum(hiddenStates[component].values())
    pick = np.random.uniform(0, max)
    current = 0
    for key in hiddenStates[component]:
        ###print(f'Key:{key}, value:{graph[key]}')
        current += hiddenStates[component][key]
        if current > pick:
            return key


def generateNewModel(models, img_size, n_classes=10, prob_residual = 50, prob_residual_downsample = 30):
    # Sample layers from start to end using roulette wheel selecetion
    # Based on models fitness
    lastLayer = "START"
    
    # Select parent model to generate a child new model
    #modelSelected = modelWeightedRouletteSelection(models, lastLayer)
    #print(modelSelected)

    newModelLayers=[]
    while(lastLayer != "END"):

        # Get model that will be encharged of the new layer
        # Select 1 model per layer
        modelSelected = modelWeightedRouletteSelection(models, lastLayer)
        if(modelSelected is None): # No model has that layer (defensive programming)
            break
        
        ###print(modelSelected)
        # All possible layers on the sampled model, from LastLayer: P(x+1|x), where x=LastLayer
        try:
            possibleChoises = models[modelSelected]['graph'][lastLayer]
        except:
            return None
        
        # Select the next layer
        if not args.random_search_layer:
            sampledLayer = layerSelectionFromModelGraph(possibleChoises, lastLayer)
        else:

            possibleChoises_keys = list(possibleChoises.keys())
            sampledLayer = possibleChoises_keys[np.random.randint(0,len(possibleChoises_keys))]

        # If sampled layer is the final one, end the generation
        # So that "END" layer does not get appended to the list
        if(sampledLayer=="END"):
            break

        # Get all possible components of a layer (kernel_size, output...)
        possibleComponents = getLayerComponents(sampledLayer)

        components = {}
        for component in possibleComponents:
            if ('None' in component):
                continue
            # Select component based on the probabilities
            if not args.random_search_layer:
                components[component] = componentSelectionFromModelLayer(models[modelSelected]['hiddenStates'][sampledLayer], component)
            else:
                possibleChoises = list(models[modelSelected]['hiddenStates'][sampledLayer][component])
                components[component] = possibleChoises[np.random.randint(0,len(possibleChoises))]

        if 'Conv' in sampledLayer:
            channels = components['out_channels']
        elif "Adaptive" in sampledLayer:
            img_size = components['output_size']
            #print(img_size)
        elif "Pool" in sampledLayer and "Adaptive" not in sampledLayer:
            if isinstance(components['kernel_size'], int):
                kernel_size = components['kernel_size'] if all(i > components['kernel_size'] for i in img_size) else 1
                img_size = [int(elem/kernel_size) for elem in img_size]
            else:
                kernel_size = components['kernel_size'] if all(i > c for c,i in zip(components['kernel_size'], img_size)) else 1
                # reduze img size accordingly
                img_size = [int(elem/k) for k,elem in zip(kernel_size, img_size)] 

        lastLayer = sampledLayer

        if 'Conv' in sampledLayer:
            # Probably, transform conv layer in residual connection
            pick = np.random.uniform(0, 100)
            if pick <= prob_residual:
                components['stride'] = 1
                pick = np.random.uniform(0, 100)
                # Probability of the bottleneck having downsampling
                if pick <= prob_residual_downsample:
                    # If the image size is already under 3 pixels in any dimension, does not downsample
                    if not (any(i < 3 for i in img_size)):
                        components['stride'] = 2
                        #print("OIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII")
                        channels = components['out_channels'] * Bottleneck.expansion
                        
                # Calculate new img_size based on the stride of the bottleneck
                img_size = [int(Conv_Output(3, elem, components['stride'],1)) for elem in img_size]

                newModelLayers.append(('Bottleneck',components))
                lastLayer = 'ReLU'
            else:
                newModelLayers.append((sampledLayer,components))
        else:
            # Add sampled layer to the end of the model
            # Sampled layer becomes the x in P(x+1|x)
            newModelLayers.append((sampledLayer,components))


    #print(newModelLayers[-1])
    # Some generated networks might not end in linear layers (e.g, if squeezenet is the last parent)
    try:
        if 'Linear' not in newModelLayers[-1][0]:
            if 'ReLU' not in newModelLayers[-1][0]:
                newModelLayers.append(('ReLU',{}))
            if 'Dropout' not in newModelLayers[-1][0]:
                modelSelected = modelWeightedRouletteSelection(models, 'Dropout')
                components = {}
                if modelSelected is None:
                    components['p'] = 0.5 # Default for pytorch
                else:
                    components['p'] = componentSelectionFromModelLayer(models[modelSelected]['hiddenStates']['Dropout'], 'p')
                newModelLayers.append(('Dropout',components))
            newModelLayers.append(('Linear', {'out_features': n_classes}))
    except:
        return None
    
    # Adjust last linear layer for output_size = n_classes
    if 'Linear' in newModelLayers[-1][0] and newModelLayers[-1][1] != n_classes:
        #newModelLayers[-1] = ('Linear', {'out_features': n_classes}) # Original
        newModelLayers[-1] = ('Linear', {'out_features': n_classes}) 
    #print(newModelLayers[-1])

    return newModelLayers





def train_network(model, utils, args, train_loader, val_loader):
    # Get Criterion
    criterion = utils.getCrossEntropyLoss()
    # Train the model for n epochs (normal training)
    # Store model in best val acc, best val loss
    best_model_wts = copy.deepcopy(model.state_dict())
    best_model_lowest_loss = []#copy.deepcopy(model.state_dict())
    best_acc_val = [0.0, 0] # accuracy, epoch
    best_loss_val = [20.0, 0] # loss epoch
    # Store each epoch evolution
    all_train_acc  = []
    all_train_loss = []
    all_val_acc    = []
    all_val_loss   = []

    # Create SGD Optimizer for this model
    optimizer = utils.getSGDOptimizer(model, learningRate=args.learning_rate,
                                    momentum=args.momentum, weight_decay=args.weight_decay)
    
    # Implemented SWA
    #optimizer = SWA(optimizer, swa_start=50, swa_freq=10, swa_lr=learning_rate)
    
    # Scheduler
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs)#, eta_min=0.001)

    startTimeFinalTrain = time.time() # Store time that it takes to train final model
    print("### Start Regular Training ###")
    for epoch in range(0,args.epochs):
        print ("-----")
        
        # Train
        model, top1, top5, train_loss = utils.train(model, train_loader, criterion, optimizer, epoch)
        all_train_acc.append(top1.avg)
        all_train_loss.append(train_loss.avg)
        
        # Learning Rate Decay Scheduler
        scheduler.step()

        # Validation
        val_top1, val_top5, val_loss = utils.evaluate(model, val_loader, criterion, epoch)
        all_val_acc.append(val_top1.avg)
        all_val_loss.append(val_loss.avg)

        if val_top1.avg > best_acc_val[0]: # store best model so far, for later, based on best val acc
            best_model_wts = copy.deepcopy(model.state_dict())
            best_acc_val[0], best_acc_val[1] = val_top1.avg, epoch
        #if val_loss.avg < best_loss_val[0]: # store best model according to loss
        #    best_model_lowest_loss = copy.deepcopy(model.state_dict())
        #    best_loss_val[0], best_loss_val[1] = val_loss.avg, epoch

        if args.prunning:
            model = unstructured_prune(model, threshold=args.prunning_threshold)

        #print (f'Epoch {epoch:3} -> Train Accuracy: {top1.avg}, Loss: {train_loss},   Val Accuracy: {val_acc}')
        print('Epoch {epoch:3} : ' 
            'Acc@1: {top1.avg:.3f}\t\t'
            'Acc@5: {top5.avg:.3f}\t\t'
            'Val Acc@1: {val_top1.avg:.4f}\t\t' 
            'Val Acc@5: {val_top5.avg:.4f}\t\t' 
            'Train Loss: {train_loss.avg:.4f}\t\t' 
            'Val Loss: {val_loss.avg:.4f}\t\t' 
            .format(
                epoch=epoch, top1=top1, top5=top5, 
                val_top1=val_top1, val_top5=val_top5, 
                train_loss=train_loss, val_loss=val_loss))

    endTimeFinalTrain = time.time() - startTimeFinalTrain # Store time that it takes to train final model

    #optimizer.swap_swa_sgd()
    
    return model, best_model_wts, best_model_lowest_loss, all_train_acc, all_train_loss, all_val_acc, all_val_loss, best_acc_val, best_loss_val, endTimeFinalTrain


#----------------------------------------------------------------------------#
#----------------------------------MAIN--------------------------------------#
def main(args, save_to = "./experiments/"):
    # Parameters
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    data_dir, batch_size, generations, populationSize, elitism, probResidual, probResidualDownsample, epochs, \
        learning_rate, weight_decay, maxneurons, epochsTrainSearch, dataset, datasetType, cutout, cutout_length, \
            auto_augment \
            = \
            args.data_dir, args.batch_size, args.generations, args.population, args.elitism, args.prob_residual, \
            args.prob_residual_downsample, args.epochs, args.learning_rate, args.weight_decay, args.max_neurons, \
            args.search_epochs, args.dataset.lower(), args.dataset_type.lower(), args.cutout, args.cutout_length, \
            args.auto_augment

    initPopDefined = args.init_pop
    modelsPath     = './models_pytorch/'
    if not os.path.exists(modelsPath):
            os.makedirs(modelsPath)
    modelsPath = os.path.join(modelsPath, dataset+'/')
    modelsPath = os.path.join(modelsPath, datasetType+'/') if "Fly" not in datasetType.lower() else os.path.join(modelsPath, 'partial/')

    # Store time for everything
    startTime      = time.time()

    # Create utils object (for train, optimizers, ...)
    utils = Utils(batch_size, device)

    # Get initial population
    # generate dictionary of models - initial population based on human-designed models
    if initPopDefined == True:
        # Read population from json
        population = readPopFromJSONs(path=modelsPath, args=args)
        if args.random_search_model:
            for key, _ in population.items():
                population[key]['fitness'] = 10
    else: 
        trainloader, valloader, n_classes = utils.get_train_dataset(args, data_dir=data_dir, \
                                                        dataset=dataset, datasetType=datasetType, resize=True)
        # Generate search space using pytorch models
        startTimeGenerateSearchSpace = time.time()
        
        if not os.path.exists(modelsPath):
            os.makedirs(modelsPath)
        population = generateGraphsOfTorchVisionModels(utils, True, data_dir, dataset, datasetType, True, 
                storeModelToFile=True, args=args, path=modelsPath, trainloader=trainloader, valloader=valloader)
        endTimeGenerateSearchSpace = time.time()
        
        print(f'Time taken to generate complete spearch space: {endTimeGenerateSearchSpace-startTimeGenerateSearchSpace}')
        exit()
        


    # Get dataloaders for the dataset
    ##trainloader, valloader, n_classes = data_loader_cifar10.get_train_valid_loader("./data/", batch_size,
    ##            datasetType=datasetType)
    trainloader, valloader, n_classes = utils.get_train_dataset(args, data_dir=data_dir, \
                                                            dataset=dataset, datasetType=datasetType)
    #testloader = utils.get_test_dataset(args, data_dir="./data/", dataset=dataset)

    # Get image shape for this dataset/problem
    for train_images, _ in trainloader:  
        sample_image = train_images[0]
        img_size = sample_image.shape # image size of the given problem
        break    
    print(f'Image size: {img_size}')


    # Get Criterion
    criterion = utils.getCrossEntropyLoss()

    startTimeSearch = time.time() # Store time that takes to perform search
    generationInfos = {} # store information about each generation
    for gen in range(0,generations):
        #trainloader, valloader, n_classes = utils.get_train_dataset(args, data_dir=data_dir, \
        #                                                dataset=dataset, datasetType=datasetType)
        print(f"### Starting Generation {gen} [...] ###")

        # newPopulation will hold the new individuals (generated models)
        newPopulation = {}

        # Elitism : the best model from last gen continues
        #sortedPopulation -> tuples: (name, MODELINFO)
        sortedPopulation = sorted(population.items(), key=lambda x: x[1]['fitness'], reverse=True)
        del population
        utils.clean_memory()

        # avoid cases where the initial generation (search space) is much lower than the populationsize
        elitism = args.elitism if args.elitism < len(sortedPopulation) else len(sortedPopulation)
        
        # Elitism for generation 0, only 'search space models' are available
        if gen == 0:
            for i in range(elitism):
                newPopulation[i] = OrderedDict()
                newPopulation[i]['graph'] = sortedPopulation[i][1]['graph']
                newPopulation[i]['hiddenStates'] = sortedPopulation[i][1]['hiddenStates']
                newPopulation[i]['fitness'] = sortedPopulation[i][1]['fitness']
        # Elitism for the rest of the generations
        else:
            iterator, countElitism, flagInitialModels = 0, 0, 1 #flaginitialmodels to 1, does not allow any of the search space models to be passed to next gens
            while(True):
                if countElitism == elitism:
                    break
                #print(sortedPopulation)
                if 'model' not in sortedPopulation[iterator][1] and flagInitialModels == 1:
                    iterator = iterator + 1
                    continue

                newPopulation[countElitism] = OrderedDict()
                newPopulation[countElitism]['graph']        = sortedPopulation[iterator][1]['graph']
                newPopulation[countElitism]['hiddenStates'] = sortedPopulation[iterator][1]['hiddenStates']
                newPopulation[countElitism]['fitness']      = sortedPopulation[iterator][1]['fitness']
                if 'model' in sortedPopulation[iterator][1]: # initial pytorch models are ditched for memory optimization
                    newPopulation[countElitism]['model']    = sortedPopulation[iterator][1]['model']
                else:
                    flagInitialModels = 1

                iterator = iterator + 1
                countElitism = countElitism + 1

        print(f'Dictionary sizes ->  newPopulation:{utils.get_size(newPopulation)} \t sortedPopulation:{utils.get_size(sortedPopulation)}')

        # transform list of Tuple(int, orderedict) into ordereddict for easier management
        sortedPopulation = OrderedDict(sortedPopulation)

        # Generate the rest of the population
        for individual in range(elitism,populationSize): #[elitism,popsize[ because the inicial indexes are the best from last generation
            # Generate individual based on the graphs from last generation
            modelLayersStringList, model = None, None # Initialize model variables
            graph, hiddenStates = None, None # Initilize HMC variables
            print("Generating new model [...]")
            while modelLayersStringList is None or model is None: # Generate models until one is ok
                # Generate a new model
                try:
                    modelLayersStringList = generateNewModel(sortedPopulation, img_size[1:], n_classes=n_classes, prob_residual = probResidual, \
                                                            prob_residual_downsample = probResidualDownsample)
                except:
                    #traceback.print_exc()
                    modelLayersStringList = None
                    utils.clean_memory()

                if modelLayersStringList is None:
                    continue

                lengthModel = 0
                for layer, components in modelLayersStringList:
                    if 'Bottleneck' in layer:
                        if components['stride'] == 2:
                            lengthModel += 8
                        else:
                            lengthModel += 7
                    else:
                        lengthModel += 1
                if lengthModel > 500: # Trying to avoid RAM segmentation fault
                    modelLayersStringList = None
                    continue

                # Transform model into pytorch
                try:
                    num_classes=n_classes
                    if args.without_training and args.mixed_training == False:
                        num_classes = 1
                    model = processify(transformNetworkIntoPytorch, modelLayersStringList, input_shape=img_size, 
                                                        maxneurons=maxneurons, n_classes=num_classes)
                    #model = transformNetworkIntoPytorch(modelLayersStringList, input_shape=img_size, 
                    #                                    maxneurons=maxneurons, nClasses=num_classes)
                    if model is None:
                        continue
                    model = model.to(device)
                except Exception as e:
                    #traceback.print_exc()
                    print(e)
                    model = None #too big for memory
                    #exit()
                    utils.clean_memory()
                    if 'broken pipe' in str(e.__str__).lower():
                        torch.cuda.empty_cache()
                    continue

                # Generate Hidden-Markov Chain
                try:
                    graph, hiddenStates = generateSummaryAndGraph(model, img_shape=img_size, 
                                                        automatedModel=True, input_image=sample_image)
                except Exception as e:
                    #print (model)
                    #print(modelLayersStringList)
                    #if "NoneType" in str(e):
                    print(f'Individual {individual} not capable of generating Hidden Markov Chain')
                    print (e)
                    del model
                    model = None
                    utils.clean_memory()
                    continue

            #print (modelLayersStringList)

            # Add individual to the pool
            newPopulation[individual] = OrderedDict()

            fitness = 0 #np.random.randint(100)

            # Acquire fitness
            print("#### Starting Fitness Extraction [...] ####")
            try:
                if args.random_search_model:
                    fitness = 10 #random, equal to all models
                else: 
                    fitness, _,_ = getFitness(args, utils, model, trainloader, valloader, device)
            except:
                fitness = 0
            #print (fitness)
            '''
            try:
                # Acquire score that represents model fitness without requiring training
                if args.without_training or args.mixed_training:
                    print("### Without Training [...] ###")
                    score = zero_cost_proxy(trainloader, model, device, batch_size, desired_size=256, generated=True)
            except Exception as e:
                score = 0
                print(e)
                if "out of memory" not in str(e):
                    print(modelLayersStringList)
                    print(model)
                utils.clean_memory()
            try:
                # Train for some epochs to acquire validation accuracy of the generated model
                if args.without_training == False or args.mixed_training:
                    print("### Regular Training [...] ###")
                    optimizer = utils.getSGDOptimizer(model)
                    for i in range(0, args.search_epochs): # Train for n epochs, evaluate
                        model, _, _, _ = utils.train(model, trainloader, criterion, optimizer)
                        top1, _, _ = utils.evaluate(model, valloader, criterion)
                        if (top1.avg > fitness):
                            val = top1.avg # store the best val accuracy (fitness)
            except Exception as e:
                val = 0
                print(e)
                if "out of memory" not in str(e):
                    print(modelLayersStringList)
                    print(model)
                utils.clean_memory()
            


            if args.mixed_training: #combination
                fitness = args.mixed_fitness_lambda*score + (1-args.mixed_fitness_lambda)*val
            elif args.without_training: # without train fitness
                fitness = score
            else: # regular partial fitness #if args.without_training == False and args.mixed_fitness==False: 
                fitness = val
            '''
            #print(f'Fitness:{fitness}')
            # Fitness is the accuracy in validation with 1 epoch train or score from NASwithout training
            newPopulation[individual]['fitness'] = fitness 

            # Add the hidden markov chain
            newPopulation[individual]['graph'] = graph
            newPopulation[individual]['hiddenStates'] = hiddenStates

            newPopulation[individual]['model'] = modelLayersStringList
            print("\n")


        #print(f'Individual:{individual}, graph:{graph}, fitness:{valacc}')
        individuals, sumFitnesses = 0, 0
        print(f'Gen:{gen} Stats:')
        for k,v in newPopulation.items():
            individuals += 1
            sumFitnesses +=v["fitness"]
            print(f'Individual:{k} - fitness:{v["fitness"]}')
        print(f'Mean Fitness:{sumFitnesses/individuals:.3f}')
        print("-- End Generation --\n")
        
        
        #print(h.heap())

        # to store the information of the best model each epoch
        sortedPopulation = sorted(newPopulation.items(), key=lambda x: x[1]['fitness'], reverse=True)
        generationInfos[gen] = {}
        generationInfos[gen]['numberIndvs'] = individuals
        generationInfos[gen]['meanFitness'] = sumFitnesses/individuals  
        generationInfos[gen]['bestModelFitness'] = sortedPopulation[0][1]['fitness']
        #generationInfos[gen]['graph'] = sortedPopulation[0][1]['graph']
        #if 'model' in sortedPopulation[0][1]: # initial pytorch models are ditched for memory optimization
        #    generationInfos[gen]['model'] = sortedPopulation[0][1]['model']

        population = copy.deepcopy(newPopulation)

        try:
            # Free memory
            del model
            del modelLayersStringList
            del newPopulation
            # Free Torch memory & Call Garbage collector
            utils.clean_memory()
        except:
            pass

    endTimeSearch = time.time() - startTimeSearch # Store final time after search



    print("#### Starting Best Model Selection [...] ####")
    trainloader, valloader, n_classes = utils.get_train_dataset(args, data_dir=data_dir, \
                                                        dataset=dataset, datasetType=datasetType)#datasetType)
    # Get best NAS generated model
    sortedPopulation = sorted(population.items(), key=lambda x: x[1]['fitness'], reverse=True)
    if args.without_training or (args.mixed_training and args.search_epochs < args.best_model_search_epochs) or (not args.without_training and not args.mixed_training and args.search_epochs < args.best_model_search_epochs):
        newPopulation = OrderedDict()
        iterations = len(sortedPopulation) if  len(sortedPopulation) < args.final_indv_without_train else args.final_indv_without_train
        for i in range(0, iterations): # default 10
            network = sortedPopulation[i][1]
            try: 
                model = transformNetworkIntoPytorch(network['model'], input_shape=img_size, maxneurons=maxneurons, n_classes=n_classes)
                model = model.to(device)
            except Exception as e: # A model from search space might happen to get till the end over the elitism, remove such choice (Defense Programming)
                print(e)
                args.final_indv_without_train = args.final_indv_without_train+1
                continue
            
            optimizer = utils.getSGDOptimizer(model, args.learning_rate)
            print("#### Starting Fitness Extraction [...] ####")
            fitness = 0#network['fitness']
            try:
                for _ in range(0, args.best_model_search_epochs): # Train for n epochs, evaluate
                    model, _, _, _ = utils.train(model, trainloader, criterion, optimizer)
                    top1, _, _ = utils.evaluate(model, valloader, criterion)
                    print(top1.avg)
                    if (top1.avg > fitness):
                        fitness = top1.avg # store the best val accuracy (fitness)
            except Exception as e: #out of memory model
                print(traceback.format_exc())
                print(e)
                pass

            del model
            utils.clean_memory()

            newPopulation[i] = OrderedDict()
            newPopulation[i]['graph'] = network['graph']
            newPopulation[i]['hiddenStates'] = network['hiddenStates']
            newPopulation[i]['model'] = network['model']
            newPopulation[i]['fitness'] = fitness

        sortedPopulation = (sorted(newPopulation.items(), key=lambda x: x[1]['fitness'], reverse=True))

        print (f'Top {iterations} models after {epochsTrainSearch} epochs trained [...]')
        for k,v in sortedPopulation:
            print(f'Individual:{k} - fitness:{v["fitness"]}')


    # Select the best model!
    for model in sortedPopulation:
        bestModel = model[1]
        # If the best model is from the search space, continue until a generated model
        if 'model' in bestModel:
            break


    print ("\n### Final Model Training [...] ###")
    # Train model
    bestval_acc_each=[]
    for i in range(1): #(3):
        if i == 1:
            #args.auto_augment = True
            args.prunning = True 
            
        # Get final dataset
        # Final train, val,test loader
        try:
            trainloader, valloader, n_classes = utils.get_train_dataset(args, data_dir=data_dir, \
                                                                    dataset=dataset, datasetType='TrainEntire', shuffle=True, \
                                                                    auto_augment=args.auto_augment)
            print(n_classes)

            model = bestModel['model']
            #if i == 0:
            #    model = utils.add_bn(model)
            #elif i==1:
            #    model = utils.add_bn(model, False)

            # Transform most performant NAS model to pytorch network
            #model = transformNetworkIntoPytorch(bestModel['model'], input_shape=img_size, maxneurons=maxneurons).to(device)
            model = transformNetworkIntoPytorch(model, input_shape=img_size, maxneurons=maxneurons, n_classes=n_classes)
            print(model)
            model = model.to(device)

            model, best_model_wts, best_model_lowest_loss, all_train_acc, all_train_loss, all_val_acc, \
                all_val_loss, best_acc_val, best_loss_val, endTimeFinalTrain \
                = train_network(model, utils, args, trainloader, valloader)

            bestval_acc_each.append(best_acc_val[0])
        except Exception as e:
            print (e)
            continue
    print(model)

    utils.clean_memory()

    model.load_state_dict(best_model_wts)

    # Test model on new data
    testloader = utils.get_test_dataset(args, data_dir=data_dir, dataset=dataset)
    accuracy_test, inference_time = utils.test(model, testloader)

    print(f'trainAcc={all_train_acc}')
    print(f'trainLoss={all_train_loss}')
    print(f'valAcc={all_val_acc}')
    print(f'valLoss={all_val_loss}')
    print(f'Final Test Accuracy: {accuracy_test} : val accuracy {best_acc_val[0]} - train accuracy {all_train_acc[best_acc_val[1]]}')
    
    #model.load_state_dict(best_model_lowest_loss)
    #accuracy_test, inference_time = utils.test(model, testloader)
    #accuracy_test = all_val_acc[best_loss_val[1]]
    #print(f'Final Test Accuracy WITH LOWEST LOSS: {accuracy_test} : val accuracy {all_val_acc[best_loss_val[1]]} - train accuracy {all_train_acc[best_loss_val[1]]}')

    # Print best individual for each generation and mean fitness
    for k,v in generationInfos.items():
        print (f'Generation:{k} -> Best Individual:{generationInfos[k]["bestModelFitness"]}; Mean Fitness:{generationInfos[k]["meanFitness"]}')

    # Times
    print(f'Time taken to perform search: {endTimeSearch}')
    print(f'Time taken to train the final model from scratch: {endTimeFinalTrain}')
    print(f'Time taken for the whole process:{time.time()-startTime}')

    try:
        model.eval()
        print(f'Model parameters:{utils.count_parameters(model)}')
        print(f'Inference Time Mean: {np.mean(inference_time)}, STD:{np.std(inference_time)} ms')
    except:
        pass

    '''
    print(f'Validation accuracy of the model using different Autoaugment schemes:')
    print(f'Normal model train: {bestval_acc_each[0]}')
    print(f'Auto Augment: {bestval_acc_each[1]}')
    '''

    '''
    print(f'Validation accuracy of the model using different BN schemes:')
    print(f'Normal model : {bestval_acc_each[2]}')
    print(f'All layers BN: {bestval_acc_each[0]}')
    print(f'non-activ  BN: {bestval_acc_each[1]}')
    '''

    torch.save(model, save_to+"model.pth")


    modelGraph = generateNetworkxDiGraph(bestModel['graph'])
    year, month, day, hour, min = map(int, time.strftime("%Y %m %d %H %M").split())
    generateGraphVizRepresentation(modelGraph, filename=save_to+"graphviz.png")
    #writeGraphDict(stateTransitionFreq, filename="./graphs/experiments.json")

    # Generate Plots
    ## Best Indv per generation
    fitnesses = [d['bestModelFitness'] for d in generationInfos.values()]
    linePlot(None, [fitnesses], ["Best Individual"], xLabel="Generation", yLabel="Fitness (%)", title=None, savePlot=True, 
        fileName=save_to+".bestindv.pdf")
    ## Mean Fitness per generation    
    fitnesses = [d['meanFitness'] for d in generationInfos.values()]
    linePlot(None, [fitnesses], ["Mean Fitness"], xLabel="Generation", yLabel="Fitness (%)", title=None, savePlot=True, 
        fileName=save_to+".meanfitness.pdf")
    ## Final Model Train Accuracy
    linePlot(None, [all_train_acc, all_val_acc], ["Train", "Validation"], xLabel="Epoch", yLabel="Accuracy (%)", title=None, savePlot=True, 
        fileName=save_to+".accuracy.pdf")
    ## Final Model Train Loss
    linePlot(None, [all_train_loss, all_val_loss], ["Train", "Validation"], xLabel="Epoch", yLabel="Loss", title=None, savePlot=True, 
        fileName=save_to+".loss.pdf") #str(year)+"-"+str(month)+"-"+str(day)+":"+str(hour)+"_"+str(datasetType)+



if __name__ == "__main__":
    parser = argparse.ArgumentParser("UMNAS")
    parser.add_argument('--data_dir', type=str, default='./data/', help='location of the datasets')
    parser.add_argument('--output_path', type=str, default='./experiments/', help='directory to where all experiments will be dumped')
    parser.add_argument('--dataset', type=str, default='cifar10', help='name of the dataset folder')
    parser.add_argument('--dataset_type', type=str, default='partial', help='Type of the dataset (ex: Full, Partial, ..)')
    parser.add_argument('--batch_size', type=int, default=96, help='batch size')
    parser.add_argument('--learning_rate', type=float, default=0.025, help='init learning rate')
    parser.add_argument('--weight_decay', type=float, default=3e-4, help='weight decay')
    parser.add_argument('--momentum', type=float, default=0.9, help='momentum')
    parser.add_argument('--epochs', type=int, default=600, help='num of training epochs')
    parser.add_argument('--search_epochs', type=int, default=4, help='num of epochs to train searched models')
    parser.add_argument('--best_model_search_epochs', type=int, default=10, help='num of epochs to select the best model searched')
    parser.add_argument('--without_training', action='store_true', default=False, help='use NAS without training for search') 
    parser.add_argument('--mixed_training', action='store_true', default=False, help='acquire fitness with partial train +without training') 
    parser.add_argument('--final_indv_without_train', type=int, default=10, help='num indvs to train at the end of the search, using without train')
    parser.add_argument('--mixed_fitness_lambda', type=float, default=0.75, help='percentage of importance given to without train (1-lambda for val)')
    parser.add_argument('--seed', type=int, default=42, help='random seed')
    parser.add_argument('--population', type=int, default=25, help='population size')
    parser.add_argument('--generations', type=int, default=50, help='number of generations')
    parser.add_argument('--max_neurons', type=int, default=6000, help="number of max input neurons in final MLP layers")
    parser.add_argument('--prob_residual', type=int, default=50, help='probability of conv becomming residual block')
    parser.add_argument('--prob_residual_downsample', type=int, default=30, help='probability of residual block includes downsampling')
    
    parser.add_argument('--init_pop', action='store_false', default=True, help='init population defined') 
    parser.add_argument('--cutout', action='store_true', default=True, help='use cutout') 
    parser.add_argument('--auto_augment', action='store_true', default=False, help='use auto augment')
    parser.add_argument('--cutout_length', type=int, default=16, help='cutout length')
    parser.add_argument('--random_search_layer', action='store_true', default=False, help='apply random search on layer selection (baseline)')
    parser.add_argument('--random_search_model', action='store_true', default=False, help='apply random search on model selection (baseline)')

    parser.add_argument('--prunning', action='store_false', default=True, help='use prune in final training') #TODO true?
    parser.add_argument('--prunning_threshold', type=float, default=1e-2, help='threshold value for prunning (abs(x)>threshold)') #TODO true?
    

    parser.add_argument('--save', type=str, default='EXP', help='experiment name')
    parser.add_argument('--gamma', type=float, default=0.1, help='learning rate decay')
    parser.add_argument('--gpu', type=int, default=0, help='gpu device id')


    parser.add_argument('--elitism', type=int, default=int((0.15*parser.parse_args().population) if int(0.15*parser.parse_args().population > 0) else 1), help="number of elements to pass through elitism")


    args = parser.parse_args()

    '''
    np.random.seed(args.seed)
    torch.cuda.set_device(args.gpu)
    cudnn.benchmark = True
    torch.manual_seed(args.seed)
    cudnn.enabled=True
    torch.cuda.manual_seed(args.seed)
    '''
    
    # check if experiments folder already exists
    experiments_path="./experiments/"
    if not os.path.exists(experiments_path):
        os.makedirs(experiments_path)
    d = datetime.datetime.now()
    iterator = 0
    while (1):
        path = os.path.join(experiments_path, "exp"+str(iterator)+"_"+d.strftime("%d_%b")+"/")
        print(path)
        if (os.path.exists(path)):
            iterator += 1
            continue
        os.makedirs(path)
        break
    
    print(args)

    main(args, save_to = path)
    sys.exit()