diff --git a/.gitignore b/.gitignore index 28be23c..e80a817 100644 --- a/.gitignore +++ b/.gitignore @@ -118,3 +118,4 @@ cython_debug/ tutorials/.ipynb_checkpoints/* docs/_build/* +examples/func_fit/evolving_state.pkl \ No newline at end of file diff --git a/examples/func_fit/xor_restore_evolving.py b/examples/func_fit/xor_restore_evolving.py new file mode 100644 index 0000000..74075e1 --- /dev/null +++ b/examples/func_fit/xor_restore_evolving.py @@ -0,0 +1,52 @@ +import jax +import numpy as np +from tensorneat.common import State +from tensorneat.pipeline import Pipeline +from tensorneat import algorithm, genome, problem +from tensorneat.common import ACT + +# neccessary settings +algorithm = algorithm.NEAT( + pop_size=1000, + species_size=20, + survival_threshold=0.01, + genome=genome.DefaultGenome( + num_inputs=3, + num_outputs=1, + max_nodes=7, + output_transform=ACT.sigmoid, + ), +) +problem = problem.XOR3d() + +pipeline = Pipeline( + algorithm, + problem, + generation_limit=200, # actually useless when we don't using auto_run() + fitness_target=-1e-6, # actually useless when we don't using auto_run() + seed=42, +) + +# load the previous evolving state +state = State.load("./evolving_state.pkl") +print("load the evolving state from ./evolving_state.pkl") + + +# compile step to speed up +compiled_step = jax.jit(pipeline.step).lower(state).compile() + +current_generation = 0 +# run 50 generations +for i in range(50): + state, previous_pop, fitnesses = compiled_step(state) + fitnesses = jax.device_get(fitnesses) # move fitness from gpu to cpu for printing + print(f"Generation {current_generation}, best fitness: {max(fitnesses)}") + current_generation += 1 + +# obtain the best individual +best_idx = np.argmax(fitnesses) +best_nodes, best_conns = previous_pop[0][best_idx], previous_pop[1][best_idx] +# make it inference +transformed = algorithm.genome.transform(state, best_nodes, best_conns) +xor3d_outputs = jax.vmap(algorithm.genome.forward, in_axes=(None, None, 0))(state, transformed, problem.inputs) +print(f"{xor3d_outputs=}") \ No newline at end of file diff --git a/examples/func_fit/xor_save_the_evolving_state.py b/examples/func_fit/xor_save_the_evolving_state.py new file mode 100644 index 0000000..9159eb7 --- /dev/null +++ b/examples/func_fit/xor_save_the_evolving_state.py @@ -0,0 +1,51 @@ +import jax +import numpy as np +from tensorneat.pipeline import Pipeline +from tensorneat import algorithm, genome, problem +from tensorneat.common import ACT + +# neccessary settings +algorithm = algorithm.NEAT( + pop_size=1000, + species_size=20, + survival_threshold=0.01, + genome=genome.DefaultGenome( + num_inputs=3, + num_outputs=1, + max_nodes=7, + output_transform=ACT.sigmoid, + ), +) +problem = problem.XOR3d() + +pipeline = Pipeline( + algorithm, + problem, + generation_limit=200, # actually useless when we don't using auto_run() + fitness_target=-1e-6, # actually useless when we don't using auto_run() + seed=42, +) +state = pipeline.setup() + +# compile step to speed up +compiled_step = jax.jit(pipeline.step).lower(state).compile() + +current_generation = 0 +# run 50 generations +for i in range(50): + state, previous_pop, fitnesses = compiled_step(state) + fitnesses = jax.device_get(fitnesses) # move fitness from gpu to cpu for printing + print(f"Generation {current_generation}, best fitness: {max(fitnesses)}") + current_generation += 1 + +# obtain the best individual +best_idx = np.argmax(fitnesses) +best_nodes, best_conns = previous_pop[0][best_idx], previous_pop[1][best_idx] +# make it inference +transformed = algorithm.genome.transform(state, best_nodes, best_conns) +xor3d_outputs = jax.vmap(algorithm.genome.forward, in_axes=(None, None, 0))(state, transformed, problem.inputs) +print(f"{xor3d_outputs=}") + +# save the evolving state +state.save("./evolving_state.pkl") +print("save the evolving state to ./evolving_state.pkl") \ No newline at end of file diff --git a/src/tensorneat/pipeline.py b/src/tensorneat/pipeline.py index 0ddf3ed..95b1ce3 100644 --- a/src/tensorneat/pipeline.py +++ b/src/tensorneat/pipeline.py @@ -83,6 +83,13 @@ class Pipeline(StatefulBaseClass): return state def step(self, state): + """ + returns: + state, previous_pop, fitnesses + state: updated state + previous_pop: previous population + fitnesses: fitnesses of previous population + """ randkey_, randkey = jax.random.split(state.randkey)