add visualize module
This commit is contained in:
wls2002
@@ -6,3 +6,5 @@ from .population import update_species, create_next_generation, speciate
|
||||
|
||||
from .genome.activations import act_name2func
|
||||
from .genome.aggregations import agg_name2func
|
||||
|
||||
from .visualize import Genome
|
||||
|
||||
@@ -1,88 +0,0 @@
|
||||
from collections import defaultdict
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def check_array_valid(nodes, cons, input_keys, output_keys):
|
||||
nodes_dict, cons_dict = array2object(nodes, cons, input_keys, output_keys)
|
||||
# assert is_DAG(cons_dict.keys()), "The genome is not a DAG!"
|
||||
|
||||
|
||||
def array2object(nodes, cons, input_keys, output_keys):
|
||||
"""
|
||||
Convert a genome from array to dict.
|
||||
:param nodes: (N, 5)
|
||||
:param cons: (C, 4)
|
||||
:param output_keys:
|
||||
:param input_keys:
|
||||
:return: nodes_dict[key: (bias, response, act, agg)], cons_dict[(i_key, o_key): (weight, enabled)]
|
||||
"""
|
||||
# update nodes_dict
|
||||
nodes_dict = {}
|
||||
for i, node in enumerate(nodes):
|
||||
if np.isnan(node[0]):
|
||||
continue
|
||||
key = int(node[0])
|
||||
assert key not in nodes_dict, f"Duplicate node key: {key}!"
|
||||
|
||||
if key in input_keys:
|
||||
assert np.all(np.isnan(node[1:])), f"Input node {key} must has None bias, response, act, or agg!"
|
||||
nodes_dict[key] = (None,) * 4
|
||||
else:
|
||||
assert np.all(~np.isnan(node[1:])), f"Normal node {key} must has non-None bias, response, act, or agg!"
|
||||
bias = node[1]
|
||||
response = node[2]
|
||||
act = node[3]
|
||||
agg = node[4]
|
||||
nodes_dict[key] = (bias, response, act, agg)
|
||||
|
||||
# check nodes_dict
|
||||
for i in input_keys:
|
||||
assert i in nodes_dict, f"Input node {i} not found in nodes_dict!"
|
||||
|
||||
for o in output_keys:
|
||||
assert o in nodes_dict, f"Output node {o} not found in nodes_dict!"
|
||||
|
||||
# update connections
|
||||
cons_dict = {}
|
||||
for i, con in enumerate(cons):
|
||||
if np.all(np.isnan(con)):
|
||||
pass
|
||||
elif np.all(~np.isnan(con)):
|
||||
i_key = int(con[0])
|
||||
o_key = int(con[1])
|
||||
if (i_key, o_key) in cons_dict:
|
||||
assert False, f"Duplicate connection: {(i_key, o_key)}!"
|
||||
assert i_key in nodes_dict, f"Input node {i_key} not found in nodes_dict!"
|
||||
assert o_key in nodes_dict, f"Output node {o_key} not found in nodes_dict!"
|
||||
weight = con[2]
|
||||
enabled = (con[3] == 1)
|
||||
cons_dict[(i_key, o_key)] = (weight, enabled)
|
||||
else:
|
||||
assert False, f"Connection {i} must has all None or all non-None!"
|
||||
|
||||
return nodes_dict, cons_dict
|
||||
|
||||
|
||||
def is_DAG(edges):
|
||||
all_nodes = set()
|
||||
for a, b in edges:
|
||||
if a == b: # cycle
|
||||
return False
|
||||
all_nodes.union({a, b})
|
||||
|
||||
for node in all_nodes:
|
||||
visited = {n: False for n in all_nodes}
|
||||
def dfs(n):
|
||||
if visited[n]:
|
||||
return False
|
||||
visited[n] = True
|
||||
for a, b in edges:
|
||||
if a == n:
|
||||
if not dfs(b):
|
||||
return False
|
||||
return True
|
||||
|
||||
if not dfs(node):
|
||||
return False
|
||||
return True
|
||||
@@ -1,132 +0,0 @@
|
||||
import time
|
||||
from typing import Union, Callable
|
||||
|
||||
import numpy as np
|
||||
import jax
|
||||
from jax import jit, vmap
|
||||
|
||||
from configs import Configer
|
||||
from algorithms import neat
|
||||
|
||||
|
||||
class Pipeline:
|
||||
"""
|
||||
Neat algorithm pipeline.
|
||||
"""
|
||||
|
||||
def __init__(self, config, seed=42):
|
||||
self.randkey = jax.random.PRNGKey(seed)
|
||||
np.random.seed(seed)
|
||||
|
||||
self.config = config # global config
|
||||
self.jit_config = Configer.create_jit_config(config) # config used in jit-able functions
|
||||
|
||||
self.P = config['pop_size']
|
||||
self.N = config['init_maximum_nodes']
|
||||
self.C = config['init_maximum_connections']
|
||||
self.S = config['init_maximum_species']
|
||||
|
||||
self.generation = 0
|
||||
self.best_genome = None
|
||||
|
||||
self.pop_nodes, self.pop_cons = neat.initialize_genomes(self.N, self.C, self.config)
|
||||
self.species_info = np.full((self.S, 3), np.nan)
|
||||
self.species_info[0, :] = 0, -np.inf, 0
|
||||
self.idx2species = np.zeros(self.P, dtype=np.float32)
|
||||
self.center_nodes = np.full((self.S, self.N, 5), np.nan)
|
||||
self.center_cons = np.full((self.S, self.C, 4), np.nan)
|
||||
self.center_nodes[0, :, :] = self.pop_nodes[0, :, :]
|
||||
self.center_cons[0, :, :] = self.pop_cons[0, :, :]
|
||||
|
||||
self.best_fitness = float('-inf')
|
||||
self.best_genome = None
|
||||
self.generation_timestamp = time.time()
|
||||
|
||||
self.evaluate_time = 0
|
||||
|
||||
self.pop_unflatten_connections = jit(vmap(neat.unflatten_connections))
|
||||
self.pop_topological_sort = jit(vmap(neat.topological_sort))
|
||||
self.forward = neat.create_forward_function(config)
|
||||
|
||||
def ask(self):
|
||||
"""
|
||||
Creates a function that receives a genome and returns a forward function.
|
||||
There are 3 types of config['forward_way']: {'single', 'pop', 'common'}
|
||||
|
||||
single:
|
||||
Create pop_size number of forward functions.
|
||||
Each function receive (batch_size, input_size) and returns (batch_size, output_size)
|
||||
e.g. RL task
|
||||
|
||||
pop:
|
||||
Create a single forward function, which use only once calculation for the population.
|
||||
The function receives (pop_size, batch_size, input_size) and returns (pop_size, batch_size, output_size)
|
||||
|
||||
common:
|
||||
Special case of pop. The population has the same inputs.
|
||||
The function receives (batch_size, input_size) and returns (pop_size, batch_size, output_size)
|
||||
e.g. numerical regression; Hyper-NEAT
|
||||
|
||||
"""
|
||||
u_pop_cons = self.pop_unflatten_connections(self.pop_nodes, self.pop_cons)
|
||||
pop_seqs = self.pop_topological_sort(self.pop_nodes, u_pop_cons)
|
||||
|
||||
# only common mode is supported currently
|
||||
assert self.config['forward_way'] == 'common'
|
||||
return lambda x: self.forward(x, pop_seqs, self.pop_nodes, u_pop_cons)
|
||||
|
||||
def tell(self, fitnesses):
|
||||
self.generation += 1
|
||||
|
||||
k1, k2, self.randkey = jax.random.split(self.randkey, 3)
|
||||
|
||||
self.species_info, self.center_nodes, self.center_cons, winner, loser, elite_mask = \
|
||||
neat.update_species(k1, fitnesses, self.species_info, self.idx2species, self.center_nodes,
|
||||
self.center_cons, self.generation, self.jit_config)
|
||||
|
||||
self.pop_nodes, self.pop_cons = neat.create_next_generation(k2, self.pop_nodes, self.pop_cons, winner, loser,
|
||||
elite_mask, self.generation, self.jit_config)
|
||||
|
||||
self.idx2species, self.center_nodes, self.center_cons, self.species_info = neat.speciate(
|
||||
self.pop_nodes, self.pop_cons, self.species_info, self.center_nodes, self.center_cons, self.generation,
|
||||
self.jit_config)
|
||||
|
||||
def auto_run(self, fitness_func, analysis: Union[Callable, str] = "default"):
|
||||
for _ in range(self.config['generation_limit']):
|
||||
forward_func = self.ask()
|
||||
|
||||
tic = time.time()
|
||||
fitnesses = fitness_func(forward_func)
|
||||
self.evaluate_time += time.time() - tic
|
||||
|
||||
assert np.all(~np.isnan(fitnesses)), "fitnesses should not be nan!"
|
||||
|
||||
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)
|
||||
|
||||
if max(fitnesses) >= self.config['fitness_threshold']:
|
||||
print("Fitness limit reached!")
|
||||
return self.best_genome
|
||||
|
||||
self.tell(fitnesses)
|
||||
print("Generation limit reached!")
|
||||
return self.best_genome
|
||||
|
||||
def default_analysis(self, fitnesses):
|
||||
max_f, min_f, mean_f, std_f = max(fitnesses), min(fitnesses), np.mean(fitnesses), np.std(fitnesses)
|
||||
|
||||
new_timestamp = time.time()
|
||||
cost_time = new_timestamp - self.generation_timestamp
|
||||
self.generation_timestamp = new_timestamp
|
||||
|
||||
max_idx = np.argmax(fitnesses)
|
||||
if fitnesses[max_idx] > self.best_fitness:
|
||||
self.best_fitness = fitnesses[max_idx]
|
||||
self.best_genome = (self.pop_nodes[max_idx], self.pop_cons[max_idx])
|
||||
|
||||
print(f"Generation: {self.generation}",
|
||||
f"fitness: {max_f}, {min_f}, {mean_f}, {std_f}, Cost time: {cost_time}")
|
||||
112
algorithms/neat/visualize.py
Normal file
112
algorithms/neat/visualize.py
Normal file
@@ -0,0 +1,112 @@
|
||||
import jax
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Genome:
|
||||
def __init__(self, nodes, cons, config):
|
||||
self.config = config
|
||||
self.nodes, self.cons = array2object(nodes, cons, config)
|
||||
if config['renumber_nodes']:
|
||||
self.renumber()
|
||||
|
||||
def __repr__(self):
|
||||
return f'Genome(\n' \
|
||||
f'\tinput_keys: {self.config["input_idx"]}, \n' \
|
||||
f'\toutput_keys: {self.config["output_idx"]}, \n' \
|
||||
f'\tnodes: \n\t\t' \
|
||||
f'{self.repr_nodes()} \n' \
|
||||
f'\tconnections: \n\t\t' \
|
||||
f'{self.repr_conns()} \n)'
|
||||
|
||||
def repr_nodes(self):
|
||||
nodes_info = []
|
||||
for key, value in self.nodes.items():
|
||||
bias, response, act, agg = value
|
||||
act_func = self.config['activation_option_names'][int(act)] if act is not None else None
|
||||
agg_func = self.config['aggregation_option_names'][int(agg)] if agg is not None else None
|
||||
s = f"{key}: (bias: {bias}, response: {response}, act: {act_func}, agg: {agg_func})"
|
||||
nodes_info.append(s)
|
||||
return ',\n\t\t'.join(nodes_info)
|
||||
|
||||
def repr_conns(self):
|
||||
conns_info = []
|
||||
for key, value in self.cons.items():
|
||||
weight, enabled = value
|
||||
s = f"{key}: (weight: {weight}, enabled: {enabled})"
|
||||
conns_info.append(s)
|
||||
return ',\n\t\t'.join(conns_info)
|
||||
|
||||
def renumber(self):
|
||||
nodes2new_nodes = {}
|
||||
new_id = len(self.config['input_idx']) + len(self.config['output_idx'])
|
||||
for key in self.nodes.keys():
|
||||
if key in self.config['input_idx'] or key in self.config['output_idx']:
|
||||
nodes2new_nodes[key] = key
|
||||
else:
|
||||
nodes2new_nodes[key] = new_id
|
||||
new_id += 1
|
||||
|
||||
new_nodes, new_cons = {}, {}
|
||||
for key, value in self.nodes.items():
|
||||
new_nodes[nodes2new_nodes[key]] = value
|
||||
for key, value in self.cons.items():
|
||||
i_key, o_key = key
|
||||
new_cons[(nodes2new_nodes[i_key], nodes2new_nodes[o_key])] = value
|
||||
self.nodes = new_nodes
|
||||
self.cons = new_cons
|
||||
|
||||
|
||||
def array2object(nodes, cons, config):
|
||||
"""
|
||||
Convert a genome from array to dict.
|
||||
:param nodes: (N, 5)
|
||||
:param cons: (C, 4)
|
||||
:return: nodes_dict[key: (bias, response, act, agg)], cons_dict[(i_key, o_key): (weight, enabled)]
|
||||
"""
|
||||
nodes, cons = jax.device_get((nodes, cons))
|
||||
# update nodes_dict
|
||||
nodes_dict = {}
|
||||
for i, node in enumerate(nodes):
|
||||
if np.isnan(node[0]):
|
||||
continue
|
||||
key = int(node[0])
|
||||
assert key not in nodes_dict, f"Duplicate node key: {key}!"
|
||||
|
||||
if key in config['input_idx']:
|
||||
assert np.all(np.isnan(node[1:])), f"Input node {key} must has None bias, response, act, or agg!"
|
||||
nodes_dict[key] = (None,) * 4
|
||||
else:
|
||||
assert np.all(
|
||||
~np.isnan(node[1:])), f"Normal node {key} must has non-None bias, response, act, or agg!"
|
||||
bias = node[1]
|
||||
response = node[2]
|
||||
act = node[3]
|
||||
agg = node[4]
|
||||
nodes_dict[key] = (bias, response, act, agg)
|
||||
|
||||
# check nodes_dict
|
||||
for i in config['input_idx']:
|
||||
assert i in nodes_dict, f"Input node {i} not found in nodes_dict!"
|
||||
|
||||
for o in config['output_idx']:
|
||||
assert o in nodes_dict, f"Output node {o} not found in nodes_dict!"
|
||||
|
||||
# update connections
|
||||
cons_dict = {}
|
||||
for i, con in enumerate(cons):
|
||||
if np.all(np.isnan(con)):
|
||||
pass
|
||||
elif np.all(~np.isnan(con)):
|
||||
i_key = int(con[0])
|
||||
o_key = int(con[1])
|
||||
if (i_key, o_key) in cons_dict:
|
||||
assert False, f"Duplicate connection: {(i_key, o_key)}!"
|
||||
assert i_key in nodes_dict, f"Input node {i_key} not found in nodes_dict!"
|
||||
assert o_key in nodes_dict, f"Output node {o_key} not found in nodes_dict!"
|
||||
weight = con[2]
|
||||
enabled = (con[3] == 1)
|
||||
cons_dict[(i_key, o_key)] = (weight, enabled)
|
||||
else:
|
||||
assert False, f"Connection {i} must has all None or all non-None!"
|
||||
|
||||
return nodes_dict, cons_dict
|
||||
Reference in New Issue
Block a user