tony/tensorneat-mend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perfect! fix bug about jax auto recompile

Browse Source 
add task xor-3d
This commit is contained in:
wls2002
2023-07-02 22:15:26 +08:00
parent e711146f41
commit c4d34e877b
11 changed files with 234 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
pipeline.py
View File

@@ -5,8 +5,8 @@ import numpy as np
import jax
from jax import jit, vmap
from configs import Configer
from algorithms import neat
from configs.configer import Configer
class Pipeline:
@@ -14,58 +14,40 @@ class Pipeline:
Neat algorithm pipeline.
"""
def __init__(self, config, seed=42):
self.randkey = jax.random.PRNGKey(seed)
np.random.seed(seed)
def __init__(self, config):
self.config = config # global config
self.jit_config = Configer.create_jit_config(config) # config used in jit-able functions
self.jit_config = Configer.create_jit_config(config)
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, 4), np.nan)
self.species_info[0, :] = 0, -np.inf, 0, self.P
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.neat_states = neat.initialize(config)
self.best_fitness = float('-inf')
self.generation_timestamp = time.time()
self.evaluate_time = 0
self.randkey, self.pop_nodes, self.pop_cons, self.species_info, self.idx2species, self.center_nodes, \
self.center_cons, self.generation, self.next_node_key, self.next_species_key = neat.initialize(config)
self.forward = neat.create_forward_function(config)
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)
# fitness_lower = np.zeros(self.P, dtype=np.float32)
# randkey_lower = np.zeros(2, dtype=np.uint32)
# pop_nodes_lower = np.zeros((self.P, self.N, 5), dtype=np.float32)
# pop_cons_lower = np.zeros((self.P, self.C, 4), dtype=np.float32)
# species_info_lower = np.zeros((self.S, 4), dtype=np.float32)
# idx2species_lower = np.zeros(self.P, dtype=np.float32)
# center_nodes_lower = np.zeros((self.S, self.N, 5), dtype=np.float32)
# center_cons_lower = np.zeros((self.S, self.C, 4), dtype=np.float32)
#
# self.tell_func = jit(neat.tell).lower(fitness_lower,
# randkey_lower,
# pop_nodes_lower,
# pop_cons_lower,
# species_info_lower,
# idx2species_lower,
# center_nodes_lower,
# center_cons_lower,
# 0,
# self.jit_config).compile()
# self.tell_func = neat.tell.lower(np.zeros(config['pop_size'], dtype=np.float32),
# self.randkey,
# self.pop_nodes,
# self.pop_cons,
# self.species_info,
# self.idx2species,
# self.center_nodes,
# self.center_cons,
# self.generation,
# self.next_node_key,
# self.next_species_key,
# self.jit_config).compile()
def ask(self):
"""
@@ -97,9 +79,19 @@ class Pipeline:
def tell(self, fitness):
self.randkey, self.pop_nodes, self.pop_cons, self.species_info, self.idx2species, self.center_nodes, \
self.center_cons, self.generation = neat.tell(fitness, self.randkey, self.pop_nodes, self.pop_cons,
self.species_info, self.idx2species, self.center_nodes,
self.center_cons, self.generation, self.jit_config)
self.center_cons, self.generation, self.next_node_key, self.next_species_key = neat.tell(fitness,
self.randkey,
self.pop_nodes,
self.pop_cons,
self.species_info,
self.idx2species,
self.center_nodes,
self.center_cons,
self.generation,
self.next_node_key,
self.next_species_key,
self.jit_config)
def auto_run(self, fitness_func, analysis: Union[Callable, str] = "default"):
for _ in range(self.config['generation_limit']):
@@ -109,7 +101,7 @@ class Pipeline:
fitnesses = fitness_func(forward_func)
self.evaluate_time += time.time() - tic
assert np.all(~np.isnan(fitnesses)), "fitnesses should not be nan!"
# assert np.all(~np.isnan(fitnesses)), "fitnesses should not be nan!"
if analysis is not None:
if analysis == "default":
@@ -138,7 +130,8 @@ class Pipeline:
self.best_fitness = fitnesses[max_idx]
self.best_genome = (self.pop_nodes[max_idx], self.pop_cons[max_idx])
species_sizes = [int(i) for i in self.species_info[:, 3] if i > 0]
member_count = jax.device_get(self.species_info[:, 3])
species_sizes = [int(i) for i in member_count if i > 0]
print(f"Generation: {self.generation}",
f"species: {len(species_sizes)}, {species_sizes}",