from src.hw_nats_fast_interface import HW_NATS_FastInterface from src.genetics import FastIndividual from src.genetics import Genetic, Population from typing import Iterable, Union, Text import numpy as np from collections import OrderedDict FreeREA_dict = { "n": 5, # tournament size "N": 25, # population size "mutation_prob": 1., # always mutates "recombination_prob": 1., # always recombines "P_parent1": 0.5, # fraction of child that comes from parent1 (on average) "n_mutations": 1, # how many loci to mutate at a time "loci_prob": None, # the probability of mutating a given locus (if None, uniform) } class GeneticSearch: def __init__(self, searchspace:HW_NATS_FastInterface, genetics_dict:dict=FreeREA_dict, init_population:Union[None, Iterable[FastIndividual]]=None, fitness_weights:Union[None, np.ndarray]=np.array([0.5, 0.5])): # instantiating a searchspace instance self.searchspace = searchspace # instatiating the dataset based on searchspace self.dataset = self.searchspace.dataset # instatiating the device based on searchspace self.target_device = self.searchspace.target_device # hardware aware scores changes based on whether or not one uses a given target device if self.target_device is None: self.hw_scores = ["flops", "params"] else: self.hw_scores = [f"{self.target_device}_energy"] # scores used to evaluate the architectures on downstream tasks self.classification_scores = ["naswot_score", "logsynflow_score", "skip_score"] self.genetics_dict = genetics_dict # weights used to combine classification performance with hardware performance. self.weights = fitness_weights # instantiating a population self.population = Population( searchspace=self.searchspace, init_population=True if init_population is None else init_population, n_individuals=self.genetics_dict["N"], normalization="dynamic" ) # initialize the object taking care of performing genetic operations self.genetic_operator = Genetic( genome=self.searchspace.all_ops, strategy="comma", # population evolution strategy tournament_size=self.genetics_dict["n"], searchspace=self.searchspace ) # preprocess population self.preprocess_population() def normalize_score(self, score_value:float, score_name:Text, type:Text="std")->float: """ Normalize the given score value using a specified normalization type. Args: score_value (float): The score value to be normalized. score_name (Text): The name of the score used for normalization. type (Text, optional): The type of normalization to be applied. Defaults to "std". Returns: float: The normalized score value. Raises: ValueError: If the specified normalization type is not available. Note: The available normalization types are: - "std": Standard score normalization using mean and standard deviation. """ if type == "std": score_mean = self.searchspace.get_score_mean(score_name) score_std = self.searchspace.get_score_std(score_name) return (score_value - score_mean) / score_std else: raise ValueError(f"Normalization type {type} not available!") def fitness_function(self, individual:FastIndividual)->FastIndividual: """ Directly overwrites the fitness attribute for a given individual. Args: individual (FastIndividual): Individual to score. # Returns: # FastIndividual: Individual, with fitness field. """ if individual.fitness is None: # None at initialization only scores = np.array([ self.normalize_score( score_value=self.searchspace.list_to_score(input_list=individual.genotype, score=score), score_name=score ) for score in self.classification_scores ]) hardware_performance = np.array([ self.normalize_score( score_value=self.searchspace.list_to_score(input_list=individual.genotype, score=score), score_name=score ) for score in self.hw_scores ]) # individual fitness is a convex combination of multiple scores network_score = (np.ones_like(scores) / len(scores)) @ scores network_hardware_performance = (np.ones_like(hardware_performance) / len(hardware_performance)) @ hardware_performance # in the hardware aware contest performance is in a direct tradeoff with hardware performance individual._fitness = np.array([network_score, -network_hardware_performance]) @ self.weights # return individual def preprocess_population(self): """ Applies scoring and fitness function to the whole population. This allows each individual to have the appropriate fields. """ # assign the fitness score self.assign_fitness() def perform_mutation( self, individual:FastIndividual, )->FastIndividual: """Performs mutation with respect to genetic ops parameters""" realization = np.random.random() if realization <= self.genetics_dict["mutation_prob"]: # do mutation mutant = self.genetic_operator.mutate( individual=individual, n_loci=self.genetics_dict["n_mutations"], genes_prob=self.genetics_dict["loci_prob"] ) return mutant else: # don't do mutation return individual def perform_recombination( self, parents:Iterable[FastIndividual], )->FastIndividual: """Performs recombination with respect to genetic ops parameters""" realization = np.random.random() if realization <= self.genetics_dict["recombination_prob"]: # do recombination child = self.genetic_operator.recombine( individuals=parents, P_parent1=self.genetics_dict["P_parent1"] ) return child else: # don't do recombination - simply return 1st parent return parents[0] def assign_fitness(self): """This function assigns to each invidual a given fitness score.""" # define a fitness function and compute fitness for each individual fitness_function = lambda individual: self.fitness_function(individual=individual) self.population.update_fitness(fitness_function=fitness_function) def obtain_parents(self, n_parents:int=2): # obtain tournament tournament = self.genetic_operator.tournament(population=self.population.individuals) # turn tournament into a local population parents = sorted(tournament, key = lambda individual: individual.fitness, reverse=True)[:n_parents] return parents def solve(self, max_generations:int=100, return_trajectory:bool=False)->Union[FastIndividual, float]: """ This function performs Regularized Evolutionary Algorithm (REA) with Training-Free metrics. Details on the whole procedure can be found in FreeREA (https://arxiv.org/pdf/2207.05135.pdf). Args: max_generations (int, optional): TODO - ADD DESCRIPTION. Defaults to 100. Returns: Union[FastIndividual, float]: Index-0 points to best individual object whereas Index-1 refers to its test accuracy. """ MAX_GENERATIONS = max_generations population, individuals = self.population, self.population.individuals bests = [] history = OrderedDict() for gen in range(MAX_GENERATIONS): # store the population history.update({self.searchspace.list_to_architecture(ind.genotype): ind for ind in population}) # save best individual bests.append(max(individuals, key=lambda ind: ind.fitness)) # perform ageing population.age() # obtain parents parents = self.obtain_parents() # obtain recombinant child child = self.perform_recombination(parents=parents) # mutate parents mutant1, mutant2 = [self.perform_mutation(parent) for parent in parents] # add mutants and child to population population.add_to_population([child, mutant1, mutant2]) # preprocess the new population - TODO: Implement a only-if-extremes-change strategy self.preprocess_population() # remove from population worst (from fitness perspective) individuals population.remove_from_population(attribute="fitness", n=2) # prune from population oldest individual population.remove_from_population(attribute="age", ascending=False) # overwrite population individuals = population.individuals best_individual = max(history.values(), key=lambda ind: ind._fitness) # appending in last position the actual best element bests.append(best_individual) test_accuracy = self.searchspace.list_to_accuracy(best_individual.genotype) if not return_trajectory: return (best_individual, test_accuracy) else: return (best_individual, test_accuracy, bests, len(history))