From 73ac1bcfe05d196f7256ce7cf57de60ca0276553 Mon Sep 17 00:00:00 2001 From: wls2002 Date: Sat, 6 May 2023 11:35:44 +0800 Subject: [PATCH] generally complete, but not work well. Debug --- algorithms/neat/genome/__init__.py | 3 +- algorithms/neat/genome/genome.py | 21 ++++- algorithms/neat/pipeline.py | 132 ++++++++++++++++++++++++++++- algorithms/neat/species.py | 74 ++++++++-------- examples/jax_playground.py | 33 ++------ examples/time_utils.py | 28 ++++++ examples/xor.py | 16 ++-- utils/default_config.json | 10 +-- 8 files changed, 233 insertions(+), 84 deletions(-) create mode 100644 examples/time_utils.py diff --git a/algorithms/neat/genome/__init__.py b/algorithms/neat/genome/__init__.py index 46fbb83..bb905b5 100644 --- a/algorithms/neat/genome/__init__.py +++ b/algorithms/neat/genome/__init__.py @@ -1,4 +1,5 @@ -from .genome import create_initialize_function +from .genome import create_initialize_function, expand, expand_single from .distance import distance from .mutate import create_mutate_function from .forward import create_forward_function +from .crossover import batch_crossover diff --git a/algorithms/neat/genome/genome.py b/algorithms/neat/genome/genome.py index 0173b06..5e7674b 100644 --- a/algorithms/neat/genome/genome.py +++ b/algorithms/neat/genome/genome.py @@ -99,8 +99,8 @@ def initialize_genomes(pop_size: int, def expand(pop_nodes: NDArray, pop_connections: NDArray, new_N: int) -> Tuple[NDArray, NDArray]: """ Expand the genome to accommodate more nodes. - :param pop_nodes: - :param pop_connections: + :param pop_nodes: (pop_size, N, 5) + :param pop_connections: (pop_size, 2, N, N) :param new_N: :return: """ @@ -114,6 +114,23 @@ def expand(pop_nodes: NDArray, pop_connections: NDArray, new_N: int) -> Tuple[ND return new_pop_nodes, new_pop_connections +def expand_single(nodes: NDArray, connections: NDArray, new_N: int) -> Tuple[NDArray, NDArray]: + """ + Expand a single genome to accommodate more nodes. + :param nodes: (N, 5) + :param connections: (2, N, N) + :param new_N: + :return: + """ + old_N = nodes.shape[0] + new_nodes = np.full((new_N, 5), np.nan) + new_nodes[:old_N, :] = nodes + + new_connections = np.full((2, new_N, new_N), np.nan) + new_connections[:, :old_N, :old_N] = connections + + return new_nodes, new_connections + @jit def add_node(new_node_key: int, nodes: Array, connections: Array, bias: float = 0.0, response: float = 1.0, act: int = 0, agg: int = 0) -> Tuple[Array, Array]: diff --git a/algorithms/neat/pipeline.py b/algorithms/neat/pipeline.py index 83ec44f..07bb1bb 100644 --- a/algorithms/neat/pipeline.py +++ b/algorithms/neat/pipeline.py @@ -1,7 +1,13 @@ +from typing import List, Union, Tuple, Callable +import time + import jax +import numpy as np from .species import SpeciesController from .genome import create_initialize_function, create_mutate_function, create_forward_function +from .genome import batch_crossover +from .genome import expand, expand_single class Pipeline: @@ -9,9 +15,13 @@ class Pipeline: Neat algorithm pipeline. """ - def __init__(self, config): + def __init__(self, config, seed=42): + self.randkey = jax.random.PRNGKey(seed) + self.config = config self.N = config.basic.init_maximum_nodes + self.expand_coe = config.basic.expands_coe + self.pop_size = config.neat.population.pop_size self.species_controller = SpeciesController(config) self.initialize_func = create_initialize_function(config) @@ -22,6 +32,11 @@ class Pipeline: self.species_controller.speciate(self.pop_nodes, self.pop_connections, self.generation) + self.new_node_keys_pool: List[int] = [max(self.output_idx) + 1] + + self.generation_timestamp = time.time() + self.best_fitness = float('-inf') + def ask(self, batch: bool): """ Create a forward function for the population. @@ -35,7 +50,120 @@ class Pipeline: def tell(self, fitnesses): self.generation += 1 - print(type(fitnesses), fitnesses) + self.species_controller.update_species_fitnesses(fitnesses) + crossover_pair = self.species_controller.reproduce(self.generation) + self.update_next_generation(crossover_pair) + + self.species_controller.speciate(self.pop_nodes, self.pop_connections, self.generation) + + self.expand() + + def auto_run(self, fitness_func, analysis: Union[Callable, str] = "default"): + for _ in range(self.config.neat.population.generation_limit): + forward_func = self.ask(batch=True) + fitnesses = fitness_func(forward_func) + + if analysis is not None: + if analysis == "default": + self.default_analysis(fitnesses) + else: + assert callable(analysis), f"What the fuck you passed in? A {analysis}?" + analysis(fitnesses) + + self.tell(fitnesses) + print("Generation limit reached!") + + def update_next_generation(self, crossover_pair: List[Union[int, Tuple[int, int]]]) -> None: + """ + create the next generation + :param crossover_pair: created from self.reproduce() + """ + + assert self.pop_nodes.shape[0] == self.pop_size + + k1, k2, self.randkey = jax.random.split(self.randkey, 3) + + # crossover + # prepare elitism mask and crossover pair + elitism_mask = np.full(self.pop_size, False) + for i, pair in enumerate(crossover_pair): + if not isinstance(pair, tuple): # elitism + elitism_mask[i] = True + crossover_pair[i] = (pair, pair) + crossover_pair = np.array(crossover_pair) + + # batch crossover + wpn = self.pop_nodes[crossover_pair[:, 0]] # winner pop nodes + wpc = self.pop_connections[crossover_pair[:, 0]] # winner pop connections + lpn = self.pop_nodes[crossover_pair[:, 1]] # loser pop nodes + lpc = self.pop_connections[crossover_pair[:, 1]] # loser pop connections + crossover_rand_keys = jax.random.split(k1, self.pop_size) + # npn, npc = batch_crossover(crossover_rand_keys, wpn, wpc, lpn, lpc) # new pop nodes, new pop connections + npn, npc = crossover_wrapper(crossover_rand_keys, wpn, wpc, lpn, lpc) + # mutate + new_node_keys = np.array(self.fetch_new_node_keys()) + mutate_rand_keys = jax.random.split(k2, self.pop_size) + m_npn, m_npc = self.mutate_func(mutate_rand_keys, npn, npc, new_node_keys) + m_npn, m_npc = jax.device_get(m_npn), jax.device_get(m_npc) + # elitism don't mutate + # (pop_size, ) to (pop_size, 1, 1) + self.pop_nodes = np.where(elitism_mask[:, None, None], npn, m_npn) + # (pop_size, ) to (pop_size, 1, 1, 1) + self.pop_connections = np.where(elitism_mask[:, None, None, None], npc, m_npc) + + # recycle unused node keys + unused = [] + for i, nodes in enumerate(self.pop_nodes): + node_keys, key = nodes[:, 0], new_node_keys[i] + if not np.isin(key, node_keys): # the new node key is not used + unused.append(key) + self.new_node_keys_pool = unused + self.new_node_keys_pool + + def expand(self): + """ + Expand the population if needed. + :return: + when the maximum node number of the population >= N + the population will expand + """ + pop_node_keys = self.pop_nodes[:, :, 0] + pop_node_sizes = np.sum(~np.isnan(pop_node_keys), axis=1) + max_node_size = np.max(pop_node_sizes) + if max_node_size >= self.N: + print(f"expand to {self.N}!") + self.N = int(self.N * self.expand_coe) + self.pop_nodes, self.pop_connections = expand(self.pop_nodes, self.pop_connections, self.N) + + # don't forget to expand representation genome in species + for s in self.species_controller.species: + s.representative = expand(*s.representative, self.N) + + def fetch_new_node_keys(self): + # if remain unused keys are not enough, create new keys + if len(self.new_node_keys_pool) < self.pop_size: + max_unused_key = max(self.new_node_keys_pool) if self.new_node_keys_pool else -1 + new_keys = list(range(max_unused_key + 1, max_unused_key + 1 + 10 * self.pop_size)) + self.new_node_keys_pool.extend(new_keys) + + # fetch keys from pool + res = self.new_node_keys_pool[:self.pop_size] + self.new_node_keys_pool = self.new_node_keys_pool[self.pop_size:] + return res + + def default_analysis(self, fitnesses): + max_f, min_f, mean_f, std_f = max(fitnesses), min(fitnesses), np.mean(fitnesses), np.std(fitnesses) + species_sizes = [len(s.members) for s in self.species_controller.species.values()] + + new_timestamp = time.time() + cost_time = new_timestamp - self.generation_timestamp + self.generation_timestamp = new_timestamp + + print(f"Generation: {self.generation}", + f"fitness: {max_f}, {min_f}, {mean_f}, {std_f}, Species sizes: {species_sizes}, Cost time: {cost_time}") + + +def crossover_wrapper(*args): + return batch_crossover(*args) diff --git a/algorithms/neat/species.py b/algorithms/neat/species.py index fb46f66..a56544e 100644 --- a/algorithms/neat/species.py +++ b/algorithms/neat/species.py @@ -1,4 +1,4 @@ -from typing import List, Tuple, Dict, Optional +from typing import List, Tuple, Dict, Union from itertools import count import jax @@ -45,7 +45,6 @@ class SpeciesController: self.species_idxer = count(0) self.species: Dict[int, Species] = {} # species_id -> species - self.genome_to_species: Dict[int, int] = {} self.o2m_distance_func = jax.vmap(distance, in_axes=(None, None, 0, 0)) # one to many # self.o2o_distance_func = np_distance # one to one @@ -79,36 +78,37 @@ class SpeciesController: # Partition population into species based on genetic similarity. # First, fast match the population to previous species - rid_list = [new_representatives[sid] for sid in previous_species_list] - res_pop_distance = [ - jax.device_get( - [ + if previous_species_list: # exist previous species + rid_list = [new_representatives[sid] for sid in previous_species_list] + res_pop_distance = [ + jax.device_get( self.o2m_distance_func(pop_nodes[rid], pop_connections[rid], pop_nodes, pop_connections) - for rid in rid_list - ] - ) - ] - pop_res_distance = np.stack(res_pop_distance, axis=0).T - for i in range(pop_res_distance.shape[0]): - if not unspeciated[i]: - continue - min_idx = np.argmin(pop_res_distance[i]) - min_val = pop_res_distance[i, min_idx] - if min_val <= self.compatibility_threshold: - species_id = previous_species_list[min_idx] - new_members[species_id].append(i) - unspeciated[i] = False + ) + for rid in rid_list + ] + + pop_res_distance = np.stack(res_pop_distance, axis=0).T + for i in range(pop_res_distance.shape[0]): + if not unspeciated[i]: + continue + min_idx = np.argmin(pop_res_distance[i]) + min_val = pop_res_distance[i, min_idx] + if min_val <= self.compatibility_threshold: + species_id = previous_species_list[min_idx] + new_members[species_id].append(i) + unspeciated[i] = False # Second, slowly match the lonely population to new-created species. # lonely genome is proved to be not compatible with any previous species, so they only need to be compared with # the new representatives. - new_species_list = [] for i in range(pop_nodes.shape[0]): if not unspeciated[i]: continue unspeciated[i] = False if len(new_representatives) != 0: - rid = [new_representatives[sid] for sid in new_representatives] # the representatives of new species + # the representatives of new species + sid, rid = list(zip(*[(k, v) for k, v in new_representatives.items()])) + distances = [ self.o2o_distance_func(pop_nodes[i], pop_connections[i], pop_nodes[r], pop_connections[r]) for r in rid @@ -117,18 +117,17 @@ class SpeciesController: min_idx = np.argmin(distances) min_val = distances[min_idx] if min_val <= self.compatibility_threshold: - species_id = new_species_list[min_idx] + species_id = sid[min_idx] new_members[species_id].append(i) - continue + continue # create a new species species_id = next(self.species_idxer) - new_species_list.append(species_id) new_representatives[species_id] = i new_members[species_id] = [i] assert np.all(~unspeciated) + # Update species collection based on new speciation. - self.genome_to_species = {} for sid, rid in new_representatives.items(): s = self.species.get(sid) if s is None: @@ -136,12 +135,7 @@ class SpeciesController: self.species[sid] = s members = new_members[sid] - for gid in members: - self.genome_to_species[gid] = sid - s.update((pop_nodes[rid], pop_connections[rid]), members) - for s in self.species.values(): - print(s.members) def update_species_fitnesses(self, fitnesses): """ @@ -189,11 +183,11 @@ class SpeciesController: result.append((sid, s, is_stagnant)) return result - def reproduce(self, generation: int) -> List[Optional[int, Tuple[int, int]]]: + def reproduce(self, generation: int) -> List[Union[int, Tuple[int, int]]]: """ code modified from neat-python! :param generation: - :return: next population indices. + :return: crossover_pair for next generation. # int -> idx in the pop_nodes, pop_connections of elitism # (int, int) -> the father and mother idx to be crossover """ @@ -235,7 +229,7 @@ class SpeciesController: self.species = {} # int -> idx in the pop_nodes, pop_connections of elitism # (int, int) -> the father and mother idx to be crossover - new_population: List[Optional[int, Tuple[int, int]]] = [] + crossover_pair: List[Union[int, Tuple[int, int]]] = [] for spawn, s in zip(spawn_amounts, remaining_species): assert spawn >= self.genome_elitism @@ -248,7 +242,7 @@ class SpeciesController: sorted_members, sorted_fitnesses = sort_element_with_fitnesses(old_members, fitnesses) if self.genome_elitism > 0: for m in sorted_members[:self.genome_elitism]: - new_population.append(m) + crossover_pair.append(m) spawn -= 1 if spawn <= 0: @@ -262,16 +256,16 @@ class SpeciesController: # Randomly choose parents and produce the number of offspring allotted to the species. for _ in range(spawn): - assert len(sorted_members) >= 2 - c1, c2 = np.random.choice(len(sorted_members), size=2, replace=False) + # allow to replace, for the case that the species only has one genome + c1, c2 = np.random.choice(len(sorted_members), size=2, replace=True) idx1, fitness1 = sorted_members[c1], sorted_fitnesses[c1] idx2, fitness2 = sorted_members[c2], sorted_fitnesses[c2] if fitness1 >= fitness2: - new_population.append((idx1, idx2)) + crossover_pair.append((idx1, idx2)) else: - new_population.append((idx2, idx1)) + crossover_pair.append((idx2, idx1)) - return new_population + return crossover_pair def compute_spawn(adjusted_fitness, previous_sizes, pop_size, min_species_size): diff --git a/examples/jax_playground.py b/examples/jax_playground.py index 98c96b9..e8f8374 100644 --- a/examples/jax_playground.py +++ b/examples/jax_playground.py @@ -5,33 +5,10 @@ from jax import random from jax import vmap, jit -def plus1(x): - return x + 1 +seed = jax.random.PRNGKey(42) +seed, *subkeys = random.split(seed, 3) -def minus1(x): - return x - 1 - - -def func(rand_key, x): - r = jax.random.uniform(rand_key, shape=()) - return jax.lax.cond(r > 0.5, plus1, minus1, x) - - -def func2(rand_key): - r = jax.random.uniform(rand_key, ()) - if r < 0.3: - return 1 - elif r < 0.5: - return 2 - else: - return 3 - - - -key = random.PRNGKey(0) -print(func(key, 0)) - -batch_func = vmap(jit(func)) -keys = random.split(key, 100) -print(batch_func(keys, jnp.zeros(100))) \ No newline at end of file +c = random.split(seed, 1) +print(seed, subkeys) +print(c) \ No newline at end of file diff --git a/examples/time_utils.py b/examples/time_utils.py new file mode 100644 index 0000000..2fd6e6f --- /dev/null +++ b/examples/time_utils.py @@ -0,0 +1,28 @@ +import cProfile +from io import StringIO +import pstats + + +def using_cprofile(func, root_abs_path=None, replace_pattern=None, save_path=None): + def inner(*args, **kwargs): + pr = cProfile.Profile() + pr.enable() + ret = func(*args, **kwargs) + pr.disable() + profile_stats = StringIO() + stats = pstats.Stats(pr, stream=profile_stats) + if root_abs_path is not None: + stats.sort_stats('cumulative').print_stats(root_abs_path) + else: + stats.sort_stats('cumulative').print_stats() + output = profile_stats.getvalue() + if replace_pattern is not None: + output = output.replace(replace_pattern, "") + if save_path is None: + print(output) + else: + with open(save_path, "w") as f: + f.write(output) + return ret + + return inner diff --git a/examples/xor.py b/examples/xor.py index c777c28..044d0a6 100644 --- a/examples/xor.py +++ b/examples/xor.py @@ -1,10 +1,12 @@ from typing import Callable, List +from functools import partial import jax import numpy as np from utils import Configer from algorithms.neat import Pipeline +from time_utils import using_cprofile xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]]) xor_outputs = np.array([[0], [1], [1], [0]]) @@ -17,22 +19,24 @@ def evaluate(forward_func: Callable) -> List[float]: """ outs = forward_func(xor_inputs) outs = jax.device_get(outs) - fitnesses = np.mean((outs - xor_outputs) ** 2, axis=(1, 2)) + fitnesses = -np.mean((outs - xor_outputs) ** 2, axis=(1, 2)) + # print(fitnesses) return fitnesses.tolist() # returns a list +# @using_cprofile +@partial(using_cprofile, root_abs_path='/mnt/e/neat-jax/', replace_pattern="/mnt/e/neat-jax/") def main(): config = Configer.load_config() pipeline = Pipeline(config) - forward_func = pipeline.ask(batch=True) - fitnesses = evaluate(forward_func) - pipeline.tell(fitnesses) + pipeline.auto_run(evaluate) - - # for i in range(100): + # for _ in range(100): + # s = time.time() # forward_func = pipeline.ask(batch=True) # fitnesses = evaluate(forward_func) # pipeline.tell(fitnesses) + # print(time.time() - s) if __name__ == '__main__': diff --git a/utils/default_config.json b/utils/default_config.json index e4db63b..a05fbff 100644 --- a/utils/default_config.json +++ b/utils/default_config.json @@ -2,15 +2,15 @@ "basic": { "num_inputs": 2, "num_outputs": 1, - "init_maximum_nodes": 20, - "expands_coe": 1.5 + "init_maximum_nodes": 10, + "expands_coe": 2 }, "neat": { "population": { "fitness_criterion": "max", - "fitness_threshold": 43.9999, + "fitness_threshold": 3, "generation_limit": 100, - "pop_size": 1000, + "pop_size": 20, "reset_on_extinction": "False" }, "gene": { @@ -73,7 +73,7 @@ "node_delete_prob": 0.2 }, "species": { - "compatibility_threshold": 3.5, + "compatibility_threshold": 8, "species_fitness_func": "max", "max_stagnation": 20, "species_elitism": 2,