134 lines
4.5 KiB
Python
134 lines
4.5 KiB
Python
import warnings
|
|
import os
|
|
import time
|
|
import numpy as np
|
|
|
|
import jax
|
|
from jax.experimental import io_callback
|
|
from evox import Monitor
|
|
from evox import State as EvoXState
|
|
|
|
from tensorneat.algorithm import BaseAlgorithm as TensorNEATAlgorithm
|
|
from tensorneat.common import State as TensorNEATState
|
|
|
|
|
|
class TensorNEATMonitor(Monitor):
|
|
|
|
def __init__(
|
|
self,
|
|
neat_algorithm: TensorNEATAlgorithm,
|
|
save_dir: str = None,
|
|
is_save: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.neat_algorithm = neat_algorithm
|
|
|
|
self.generation_timestamp = time.time()
|
|
self.alg_state: TensorNEATState = None
|
|
self.fitness = None
|
|
self.best_fitness = -np.inf
|
|
self.best_genome = None
|
|
|
|
self.is_save = is_save
|
|
|
|
if is_save:
|
|
if save_dir is None:
|
|
now = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
|
|
self.save_dir = f"./{self.__class__.__name__} {now}"
|
|
else:
|
|
self.save_dir = save_dir
|
|
print(f"save to {self.save_dir}")
|
|
if not os.path.exists(self.save_dir):
|
|
os.makedirs(self.save_dir)
|
|
self.genome_dir = os.path.join(self.save_dir, "genomes")
|
|
if not os.path.exists(self.genome_dir):
|
|
os.makedirs(self.genome_dir)
|
|
|
|
def hooks(self):
|
|
return ["pre_tell"]
|
|
|
|
def pre_tell(self, state: EvoXState, cand_sol, transformed_cand_sol, fitness, transformed_fitness):
|
|
io_callback(
|
|
self.store_info,
|
|
None,
|
|
state,
|
|
transformed_fitness,
|
|
)
|
|
|
|
def store_info(self, state: EvoXState, fitness):
|
|
self.alg_state: TensorNEATState = state.query_state("algorithm").alg_state
|
|
self.fitness = jax.device_get(fitness)
|
|
|
|
def show(self):
|
|
pop = self.neat_algorithm.ask(self.alg_state)
|
|
valid_fitnesses = self.fitness[~np.isinf(self.fitness)]
|
|
|
|
max_f, min_f, mean_f, std_f = (
|
|
max(valid_fitnesses),
|
|
min(valid_fitnesses),
|
|
np.mean(valid_fitnesses),
|
|
np.std(valid_fitnesses),
|
|
)
|
|
|
|
new_timestamp = time.time()
|
|
|
|
cost_time = new_timestamp - self.generation_timestamp
|
|
self.generation_timestamp = new_timestamp
|
|
|
|
max_idx = np.argmax(self.fitness)
|
|
if self.fitness[max_idx] > self.best_fitness:
|
|
self.best_fitness = self.fitness[max_idx]
|
|
self.best_genome = pop[0][max_idx], pop[1][max_idx]
|
|
|
|
if self.is_save:
|
|
best_genome = jax.device_get((pop[0][max_idx], pop[1][max_idx]))
|
|
with open(
|
|
os.path.join(
|
|
self.genome_dir,
|
|
f"{int(self.neat_algorithm.generation(self.alg_state))}.npz",
|
|
),
|
|
"wb",
|
|
) as f:
|
|
np.savez(
|
|
f,
|
|
nodes=best_genome[0],
|
|
conns=best_genome[1],
|
|
fitness=self.best_fitness,
|
|
)
|
|
|
|
# save best if save path is not None
|
|
member_count = jax.device_get(self.neat_algorithm.member_count(self.alg_state))
|
|
species_sizes = [int(i) for i in member_count if i > 0]
|
|
|
|
pop = jax.device_get(pop)
|
|
pop_nodes, pop_conns = pop # (P, N, NL), (P, C, CL)
|
|
nodes_cnt = (~np.isnan(pop_nodes[:, :, 0])).sum(axis=1) # (P,)
|
|
conns_cnt = (~np.isnan(pop_conns[:, :, 0])).sum(axis=1) # (P,)
|
|
|
|
max_node_cnt, min_node_cnt, mean_node_cnt = (
|
|
max(nodes_cnt),
|
|
min(nodes_cnt),
|
|
np.mean(nodes_cnt),
|
|
)
|
|
|
|
max_conn_cnt, min_conn_cnt, mean_conn_cnt = (
|
|
max(conns_cnt),
|
|
min(conns_cnt),
|
|
np.mean(conns_cnt),
|
|
)
|
|
|
|
print(
|
|
f"Generation: {self.neat_algorithm.generation(self.alg_state)}, Cost time: {cost_time * 1000:.2f}ms\n",
|
|
f"\tnode counts: max: {max_node_cnt}, min: {min_node_cnt}, mean: {mean_node_cnt:.2f}\n",
|
|
f"\tconn counts: max: {max_conn_cnt}, min: {min_conn_cnt}, mean: {mean_conn_cnt:.2f}\n",
|
|
f"\tspecies: {len(species_sizes)}, {species_sizes}\n",
|
|
f"\tfitness: valid cnt: {len(valid_fitnesses)}, max: {max_f:.4f}, min: {min_f:.4f}, mean: {mean_f:.4f}, std: {std_f:.4f}\n",
|
|
)
|
|
|
|
# append log
|
|
if self.is_save:
|
|
with open(os.path.join(self.save_dir, "log.txt"), "a") as f:
|
|
f.write(
|
|
f"{self.neat_algorithm.generation(self.alg_state)},{max_f},{min_f},{mean_f},{std_f},{cost_time}\n"
|
|
)
|