#| This file has a major part from the pymap_elites framework. #| Copyright 2019, INRIA #| Main contributor(s): #| Jean-Baptiste Mouret, jean-baptiste.mouret@inria.fr #| Eloise Dalin , eloise.dalin@inria.fr #| Pierre Desreumaux , pierre.desreumaux@inria.fr #| #| #| **Main paper**: Mouret JB, Clune J. Illuminating search spaces by #| mapping elites. arXiv preprint arXiv:1504.04909. 2015 Apr 20. #| #| This software is governed by the CeCILL license under French law #| and abiding by the rules of distribution of free software. You #| can use, modify and/ or redistribute the software under the terms #| of the CeCILL license as circulated by CEA, CNRS and INRIA at the #| following URL "http://www.cecill.info". #| #| As a counterpart to the access to the source code and rights to #| copy, modify and redistribute granted by the license, users are #| provided only with a limited warranty and the software's author, #| the holder of the economic rights, and the successive licensors #| have only limited liability. #| #| In this respect, the user's attention is drawn to the risks #| associated with loading, using, modifying and/or developing or #| reproducing the software by the user in light of its specific #| status of free software, that may mean that it is complicated to #| manipulate, and that also therefore means that it is reserved for #| developers and experienced professionals having in-depth computer #| knowledge. Users are therefore encouraged to load and test the #| software's suitability as regards their requirements in conditions #| enabling the security of their systems and/or data to be ensured #| and, more generally, to use and operate it in the same conditions #| as regards security. #| #| The fact that you are presently reading this means that you have #| had knowledge of the CeCILL license and that you accept its terms. import math import numpy as np import multiprocessing # from scipy.spatial import cKDTree : TODO -- faster? from sklearn.neighbors import KDTree from map_elites import common as cm from flopth import flopth from datasets import get_datasets from utils import read_python_file, extract_code, get_net_name, extract_code_section, get_class, csv_writer, get_max_curiosity, get_max_fitness, count_unique_components, count_parameters, get_network_width_depth_ratio from utils import Net from mutation_models import codegen_mutate, codex_mutate, replace_word_mutation from evaluations import is_trainable from train import train_net_on_cpu, train_net_on_gpu, forward_pass_on_cpu,forward_pass_on_gpu, detect_layers, transfer_weights import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim import torchvision import torchvision.transforms as transforms import numpy as np import os import random import hydra import ray from fvcore.nn import FlopCountAnalysis from conf.config import Config def __add_to_archive(s, centroid, archive, kdt, type_ind): niche_index = kdt.query([centroid], k=1)[1][0][0] niche = kdt.data[niche_index] n = cm.make_hashable(niche) s.centroid = n if n in archive: if type_ind == "network": if s.fitness < archive[n].fitness: archive[n] = s return True elif type_ind == "prompt": if s.fitness > archive[n].fitness: archive[n] = s return True return False else: archive[n] = s return True def to_specie(net,fit,desc,net_path): return cm.Species(net, desc, fit, net_path) # map-elites algorithm (CVT variant) @hydra.main(version_base="1.3.0", config_path="conf", config_name="config") def main(cfg: Config): """CVT MAP-Elites Vassiliades V, Chatzilygeroudis K, Mouret JB. Using centroidal voronoi tessellations to scale up the multidimensional archive of phenotypic elites algorithm. IEEE Transactions on Evolutionary Computation. 2017 Aug 3;22(4):623-30. Format of the logfile: evals archive_size max mean median 5%_percentile, 95%_percentile """ if cfg.DEVICE == "cuda": ray.init(num_gpus=cfg.NUM_GPUS_TOTAL) else: ray.init() mutation_fn = codex_mutate if cfg.MUTATION == "codex" else codegen_mutate prompts = ['"""Add a layer to improve the above network"""', '"""Delete a layer to improve the above network"""', '"""improve the above network"""', '"""improve the above network by reducing the size drastically"""', '"""improve the above network by increasing the size drastically"""', # Specific prompts below '"""Add fully connected layer to improve the above network"""', '"""Add convolutional layer to improve the above network"""', '"""Add pooling layer to improve the above network"""', '"""Add residual connection to improve the above network"""', '"""Add multiple residual connections to improve the above network"""', '"""Add dropout layer to improve the above network"""', '"""Add normalization layer to improve the above network"""', '"""Add recurrent layer to improve the above network"""', ] probabilities = [1.0 / len(prompts)] * len(prompts) prompt_score = {} for prompt in prompts: prompt_score[prompt] = 0 prompt_score["None"] = 0 prompt_to_int = {} int_to_prompt = {} for i,prompt in enumerate(prompts): prompt_to_int[prompt] = i int_to_prompt[i] = prompt if cfg.START_FROM_CHECKPOINT: best_net = read_python_file("initial_net.py") init_gen = 4 print(f"Starting from the following network:\n{best_net}") else: init_net = read_python_file(os.path.normpath(f"{cfg.SAVE_DIR}/initial_net.py")) init_gen = 0 print(f"Starting from the following network:\n{init_net}") # create the CVT params = cm.default_params c = cm.cvt(cfg.N_NICHES,cfg.DIM_MAP, params['cvt_samples'], params['cvt_use_cache']) kdt = KDTree(c, leaf_size=30, metric='euclidean') cm.__write_centroids(c) prompt_archive = {} # init archive (empty) nets_archive = {} n_evals = 0 # number of evaluations since the beginning b_evals = 0 # number evaluation since the last dump curios_net_path = os.path.normpath(f"{cfg.SAVE_DIR}/initial_net.py") c_net_str = read_python_file(curios_net_path) Net = get_class(curios_net_path) curios_net = Net() curios_prompt = prompts[0] # main loop temperature_start = 0.6 prev_best_score = np.inf prev_best_net_path = None prev_best_prompt = None exp_name = f"gen-nets_{cfg.MUTATION}_networks-{cfg.NUM_NETS}_temp-{cfg.TEMPERATURE}_net-training-epochs-{cfg.NET_TRAINING_EPOCHS}_niches-{cfg.N_NICHES}_infer-and-flops-as-bd" path_nets = f"{cfg.SAVE_DIR}/logs/{exp_name}" os.makedirs(path_nets,exist_ok=True) log_file = open(os.path.normpath(f"{path_nets}/cvt.dat"),"w") out_file = os.path.normpath(f"{path_nets}/exp_results.csv") csv_writer(["generations", "best_loss", "used_prompt"],out_file) for gen_i in range(init_gen, cfg.GENERATIONS): generated_nets = [] # random initialization (print(f"[Generation: {gen_i}]")) if (len(nets_archive.keys()) < cfg.RANDOM_INIT_NETS) and (len(prompt_archive.keys()) < (int(cfg.RANDOM_INIT_NETS*2))):#params['random_init'] * n_niches: for _i in range(cfg.ROLL_OUTS): selected_prompts = [] generated_net_results = [] selected_prompts = random.choices(prompts, weights=probabilities,k=cfg.INIT_NUM_NETS) for i in range(0, cfg.INIT_NUM_NETS): print(f"Selected prompt for generations: {selected_prompts[i]}") for prompt in selected_prompts: generated_net_results.append(mutation_fn.remote(cfg=cfg, prompt=init_net + "\n" + prompt, temperature=temperature_start)) for i,generated_net in enumerate(generated_net_results): generated_nets.append((ray.get(generated_net),selected_prompts[i],temperature_start)) else: # variation/selection loop evo_operator = random.choices(["mutation","crossover"], weights=[0.85,0.15])[0] print(f"Performing {evo_operator}") for _i in range(cfg.ROLL_OUTS): generated_net_results = [] if evo_operator == "mutation": if len(nets_archive.keys()) < 3: n_nets = 1 selection = 0 else: n_nets = 3 selection = random.randint(0, 2) curios_nets = [] curios_prompts = [] curios_temps = [] curios_net_paths = [] n_curios_nets = get_max_fitness(nets_archive.values(),n_nets) for n in n_curios_nets: curios_nets.append(n.x) curios_net_paths.append(n.net_path) n_prompts_temps = get_max_curiosity(prompt_archive.values(),n_nets) for pt in n_prompts_temps: curios_prompts.append(pt.desc[0]) curios_temps.append(pt.desc[1]) curios_net = curios_nets[selection] curios_temp = curios_temps[selection] curios_net_path = curios_net_paths[selection] curios_temp_ray = [] curios_prompt_ray = [] for i in range(cfg.NUM_NETS): # Temperature mutation if i == 1: curios_temp = curios_temps[selection] else: curios_temp += random.uniform(-0.1,0.1) if curios_temp > 1.0: curios_temp = 1.0 elif curios_temp < 0.1: curios_temp = 0.1 curios_prompt = curios_prompts[selection] curios_temp_ray.append(curios_temp) curios_prompt_ray.append(curios_prompt) print(f"prompt in mutation is {curios_prompt}") for i in range(cfg.NUM_NETS): generated_net_results.append(mutation_fn.remote(cfg=cfg, prompt=read_python_file(curios_net_path) + "\n" + int_to_prompt[int(curios_prompt_ray[i])], temperature = curios_temp_ray[i])) for i,generated_net in enumerate(generated_net_results): generated_nets.append((ray.get(generated_net),int_to_prompt[int(curios_prompt_ray[i])],curios_temp_ray[i])) elif evo_operator == "crossover": if len(nets_archive.keys()) < 1: print("No crossover performed") else: curios_nets_ray = [] curios_temps_ray = [] curios_nets = [] curios_prompts = [] curios_temps = [] curios_net_paths = [] if len(nets_archive.keys()) < 2: selection = 0 n_nets = 1 else: selection = random.randint(1, len(nets_archive.keys())-1)# change it to 1 n_nets = 2 n_curios_nets = get_max_fitness(nets_archive.values(),n_nets) for n in n_curios_nets: curios_nets.append(n.x) curios_net_paths.append(n.net_path) n_prompts_temps = get_max_curiosity(prompt_archive.values(),n_nets) for pt in n_prompts_temps: curios_prompts.append(pt.desc[0]) curios_temps.append(pt.desc[1]) for i in range(0, cfg.NUM_NETS): curios_net_str = read_python_file(curios_net_paths[0]) curios_net = curios_nets[0] curios_temp = curios_temps[0] curios_prompt = curios_prompts[0] curios_nets_ray.append(read_python_file(curios_net_paths[selection])) crossover_prompt = '"""Combine the above two neural networks and create a third neural network class that also inherits from nn.Module"""' for curios_2nd_net in curios_nets_ray: generated_net_results.append(mutation_fn.remote(cfg=cfg, prompt=curios_net_str + "\n" + curios_2nd_net + "\n" + crossover_prompt, temperature = curios_temp)) for i,generated_net in enumerate(generated_net_results): generated_nets.append((ray.get(generated_net),int_to_prompt[int(curios_prompt)],curios_temp)) net_class_name = None net_paths=[] training_prompts=[] training_nets=[] for i, k in enumerate(generated_nets): invalid_net = False generation, prompt, temperature = k net_path = os.path.normpath(f"{path_nets}/network_{gen_i}_{i}_{prompt[3:-3]}_{temperature}.py") if cfg.MUTATION == "codex": extract_code_section(generation['choices'][0]["text"], prompt, file_path=net_path) elif cfg.MUTATION.startswith("codegen"): extract_code_section(generation, prompt, file_path=net_path) main_net_focus = read_python_file(net_path) print(f"Net in focus:\n {main_net_focus}") if not invalid_net: try: Net = get_class(net_path) net = Net() except Exception as e: net = curios_net net_path = curios_net_path if isinstance(curios_prompt,str): prompt = curios_prompt elif isinstance(curios_prompt,float) or isinstance(curios_prompt,int): prompt = int_to_prompt[int(curios_prompt)] #breakpoint() try: is_t = is_trainable(net) except Exception: is_t = False if is_t: net_paths.append(net_path) training_prompts.append(prompt) training_nets.append(net) else: invalid_net = True if invalid_net: print(f"The network at {net_path} is not trainable") if training_nets == []: continue if (len(nets_archive.keys()) < cfg.RANDOM_INIT_NETS) and (len(prompt_archive.keys()) < (int(cfg.RANDOM_INIT_NETS*2))): inter_results = [] losses = [] layers_c_net = detect_layers(curios_net) if cfg.DEVICE == "cpu": for net in training_nets: for layer in layers_c_net: try: net = transfer_weights(curios_net,net,layer) print(f"Weights transferred successfully") except Exception: print("Weights can not transfer") for net in training_nets: inter_results.append(forward_pass_on_gpu.remote(net)) elif cfg.DEVICE == "cuda": for net in training_nets: for layer in layers_c_net: try: net = transfer_weights(curios_net,net,layer) print(f"Weights transferred successfully") except Exception: print("Weights can not transfer") for net in training_nets: inter_results.append(forward_pass_on_gpu.remote(net)) else: raise ValueError(f"{cfg.DEVICE} is not a valid device") fitness = [] for i,result in enumerate(inter_results): try: res = ray.get(result) fitness.append([res[0],training_prompts[i],temperature,net_paths[i],res[1]]) if fitness[i][0] <= prev_best_score: prev_best_net_path = net_paths[i] prev_best_prompt = training_prompts[i] prev_best_score = fitness[i][0] temperature += 0.05 else: temperature -= 0.05 if temperature > 1.0: temperature = 1.0 elif temperature < 0.1: temperature = 0.1 fitness[i][2] = temperature except Exception: print("not trainable due to fitness 1") else: if cfg.RAY: inter_results = [] losses = [] layers_c_net = detect_layers(curios_net) if cfg.DEVICE == "cpu": for net in training_nets: for layer in layers_c_net: try: net = transfer_weights(curios_net,net,layer) print(f"Weights transferred successfully") except Exception: print("Weights can not transfer") for net in training_nets: inter_results.append(train_net_on_cpu.remote(net,cfg.NET_TRAINING_EPOCHS)) elif cfg.DEVICE == "both": for i, net in enumerate(training_nets): if i % 2 == 0: inter_results.append(train_net_on_cpu.remote(net,cfg.NET_TRAINING_EPOCHS)) else: inter_results.append(train_net_on_gpu.remote(net,cfg.NET_TRAINING_EPOCHS)) elif cfg.DEVICE == "cuda": for net in training_nets: for layer in layers_c_net: try: net = transfer_weights(curios_net,net,layer) print(f"Weights transferred successfully") except Exception: print("Weights can not transfer") for net in training_nets: inter_results.append(train_net_on_gpu.remote(net,cfg.NET_TRAINING_EPOCHS)) else: raise ValueError(f"{cfg.DEVICE} is not a valid device") fitness = [] for i,result in enumerate(inter_results): try: fitness.append([ray.get(result),training_prompts[i],temperature,net_paths[i],0.0]) if fitness[i][0] <= prev_best_score: prev_best_net_path = net_paths[i] prev_best_prompt = training_prompts[i] prev_best_score = fitness[i][0] temperature += 0.05 else: temperature -= 0.05 if temperature > 1.0: temperature = 1.0 elif temperature < 0.1: temperature = 0.1 fitness[i][2] = temperature except Exception: print("not trainable due to fitness 2") try: infer_results = [] for net in training_nets: infer_results.append(forward_pass_on_gpu.remote(net)) for i,result in enumerate(infer_results): try: res = ray.get(result) fitness[i][4] = res[1] except Exception: print("not trainable due to inference speed 1") except Exception: print("not trainable due to inference speed 2") dummy_inputs = torch.zeros((1, 3, 32, 32)) net_beh_list = [] prompt_beh_list = [] for loss_x,prompt_x,temp_x,net_p,infer_speed in fitness: _nn = read_python_file(net_p) #Avoiding sytax errors Net = get_class(net_p) net = Net() try: flps = FlopCountAnalysis(net, dummy_inputs) flops = flps.total() depth_width_ratio = get_network_width_depth_ratio(net) except Exception: try: flps = FlopCountAnalysis(curios_net, dummy_inputs) flops = flps.total() depth_width_ratio = get_network_width_depth_ratio(curios_net) except Exception: flops = 0 depth_width_ratio = 1 print(f"flops, infer_speed, depth_width_ratio: {flops}, {infer_speed}, {depth_width_ratio}") s_net = to_specie(net,loss_x,np.array([depth_width_ratio, flops]),net_p) net_added = __add_to_archive(s_net, s_net.desc, nets_archive, kdt, type_ind = "network") if net_added: prompt_fit = 1.0 else: prompt_fit = 0.0 s_prompt = to_specie(net,prompt_fit,np.array([prompt_to_int[prompt_x], temp_x]),net_p) prompt_added = __add_to_archive(s_prompt, s_prompt.desc, prompt_archive, kdt, type_ind = "prompt") if not net_added: s_prompt.curiosity = s_prompt.curiosity - 0.5 elif net_added: s_prompt.curiosity = s_prompt.curiosity + 1.0 if gen_i % 1 == 0: cm.__save_archive(nets_archive, gen_i,name="net") cm.__save_archive(prompt_archive, gen_i,name="prompt") # write log if log_file != None: fit_list = np.array([x.fitness for x in nets_archive.values()]) log_file.write("{} {} {} {} {} {} {}\n".format(gen_i, len(nets_archive.keys()), len(prompt_archive.keys()), fit_list.min(), np.mean(fit_list), np.median(fit_list), np.percentile(fit_list, 5), np.percentile(fit_list, 95))) log_file.flush() cm.__save_archive(nets_archive, gen_i,name="net") cm.__save_archive(prompt_archive, gen_i,name="prompt") return nets_archive,prompt_archive if __name__ == "__main__": main()