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

update some files for save

Browse Source
This commit is contained in:
root
2024-07-15 14:25:51 +08:00
parent f032564a43
commit 6edf083d4f
12 changed files with 110 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/tensorneat/algorithm/hyperneat/substrate/default.py
View File

@@ -1,4 +1,5 @@
from jax import vmap, numpy as jnp
from jax import vmap
import numpy as np
from .base import BaseSubstrate
from tensorneat.genome.utils import set_conn_attrs
@@ -8,9 +9,9 @@ class DefaultSubstrate(BaseSubstrate):
def __init__(self, num_inputs, num_outputs, coors, nodes, conns):
self.inputs = num_inputs
self.outputs = num_outputs
self.coors = jnp.array(coors)
self.nodes = jnp.array(nodes)
self.conns = jnp.array(conns)
self.coors = np.array(coors)
self.nodes = np.array(nodes)
self.conns = np.array(conns)
def make_nodes(self, query_res):
return self.nodes

11
src/tensorneat/common/state.py
View File

@@ -1,3 +1,5 @@
import pickle
from jax.tree_util import register_pytree_node_class
@@ -39,6 +41,15 @@ class State:
def __contains__(self, item):
return item in self.state_dict
def save(self, file_name):
with open(file_name, "wb") as f:
pickle.dump(self, f)
@classmethod
def load(cls, file_name):
with open(file_name, "rb") as f:
return pickle.load(f)
def tree_flatten(self):
children = list(self.state_dict.values())
aux_data = list(self.state_dict.keys())

98
src/tensorneat/common/stateful_class.py
View File

@@ -9,30 +9,6 @@ class StatefulBaseClass:
def setup(self, state=State()):
return state
def save(self, state: Optional[State] = None, path: Optional[str] = None):
if path is None:
time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
path = f"./{self.__class__.__name__} {time}.pkl"
if state is not None:
self.__dict__["aux_for_state"] = state
with open(path, "wb") as f:
pickle.dump(self, f)
def __getstate__(self):
# only pickle the picklable attributes
state = self.__dict__.copy()
non_picklable_keys = []
for key, value in state.items():
try:
pickle.dumps(value)
except Exception:
non_picklable_keys.append(key)
for key in non_picklable_keys:
state.pop(key)
return state
def show_config(self, registered_objects=None):
if registered_objects is None: # root call
registered_objects = []
@@ -47,27 +23,53 @@ class StatefulBaseClass:
config[str(key)] = str(value)
return config
@classmethod
def load(cls, path: str, with_state: bool = False, warning: bool = True):
with open(path, "rb") as f:
obj = pickle.load(f)
if with_state:
if "aux_for_state" not in obj.__dict__:
if warning:
warnings.warn(
"This object does not have state to load, return empty state",
category=UserWarning,
)
return obj, State()
state = obj.__dict__["aux_for_state"]
del obj.__dict__["aux_for_state"]
return obj, state
else:
if "aux_for_state" in obj.__dict__:
if warning:
warnings.warn(
"This object has state to load, ignore it",
category=UserWarning,
)
del obj.__dict__["aux_for_state"]
return obj
# TODO: Bug need be fixed
# def save(self, state: Optional[State] = None, path: Optional[str] = None):
# if path is None:
# time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
# path = f"./{self.__class__.__name__} {time}.pkl"
# if state is not None:
# self.__dict__["aux_for_state"] = state
# with open(path, "wb") as f:
# pickle.dump(self, f)
# def __getstate__(self):
# # only pickle the picklable attributes
# state = self.__dict__.copy()
# non_picklable_keys = []
# for key, value in state.items():
# try:
# pickle.dumps(value)
# except Exception as e:
# print(f"Cannot pickle key {key}: {e}")
# non_picklable_keys.append(key)
# for key in non_picklable_keys:
# state.pop(key)
# return state
# @classmethod
# def load(cls, path: str, with_state: bool = False, warning: bool = True):
# with open(path, "rb") as f:
# obj = pickle.load(f)
# if with_state:
# if "aux_for_state" not in obj.__dict__:
# if warning:
# warnings.warn(
# "This object does not have state to load, return empty state",
# category=UserWarning,
# )
# return obj, State()
# state = obj.__dict__["aux_for_state"]
# del obj.__dict__["aux_for_state"]
# return obj, state
# else:
# if "aux_for_state" in obj.__dict__:
# if warning:
# warnings.warn(
# "This object has state to load, ignore it",
# category=UserWarning,
# )
# del obj.__dict__["aux_for_state"]
# return obj

12
src/tensorneat/genome/base.py
View File

@@ -96,9 +96,9 @@ class BaseGenome(StatefulBaseClass):
def setup(self, state=State()):
state = self.node_gene.setup(state)
state = self.conn_gene.setup(state)
state = self.mutation.setup(state, self)
state = self.crossover.setup(state, self)
state = self.distance.setup(state, self)
state = self.mutation.setup(state)
state = self.crossover.setup(state)
state = self.distance.setup(state)
return state
def transform(self, state, nodes, conns):
@@ -114,13 +114,13 @@ class BaseGenome(StatefulBaseClass):
raise NotImplementedError
def execute_mutation(self, state, randkey, nodes, conns, new_node_key):
return self.mutation(state, randkey, nodes, conns, new_node_key)
return self.mutation(state, self, randkey, nodes, conns, new_node_key)
def execute_crossover(self, state, randkey, nodes1, conns1, nodes2, conns2):
return self.crossover(state, randkey, nodes1, conns1, nodes2, conns2)
return self.crossover(state, self, randkey, nodes1, conns1, nodes2, conns2)
def execute_distance(self, state, nodes1, conns1, nodes2, conns2):
return self.distance(state, nodes1, conns1, nodes2, conns2)
return self.distance(state, self, nodes1, conns1, nodes2, conns2)
def initialize(self, state, randkey):
k1, k2 = jax.random.split(randkey) # k1 for nodes, k2 for conns

2
src/tensorneat/genome/gene/node/bias.py
View File

@@ -13,7 +13,7 @@ from tensorneat.common import (
get_func_name
)
from . import BaseNode
from .base import BaseNode
class BiasNode(BaseNode):

7
src/tensorneat/genome/operations/crossover/base.py
View File

@@ -3,10 +3,5 @@ from tensorneat.common import StatefulBaseClass, State
class BaseCrossover(StatefulBaseClass):
def setup(self, state=State(), genome = None):
assert genome is not None, "genome should not be None"
self.genome = genome
return state
def __call__(self, state, randkey, nodes1, nodes2, conns1, conns2):
def __call__(self, state, genome, randkey, nodes1, nodes2, conns1, conns2):
raise NotImplementedError

10
src/tensorneat/genome/operations/crossover/default.py
View File

@@ -11,14 +11,14 @@ from ...utils import (
class DefaultCrossover(BaseCrossover):
def __call__(self, state, randkey, nodes1, conns1, nodes2, conns2):
def __call__(self, state, genome, randkey, nodes1, conns1, nodes2, conns2):
"""
use genome1 and genome2 to generate a new genome
notice that genome1 should have higher fitness than genome2 (genome1 is winner!)
"""
randkey1, randkey2 = jax.random.split(randkey, 2)
randkeys1 = jax.random.split(randkey1, self.genome.max_nodes)
randkeys2 = jax.random.split(randkey2, self.genome.max_conns)
randkeys1 = jax.random.split(randkey1, genome.max_nodes)
randkeys2 = jax.random.split(randkey2, genome.max_conns)
# crossover nodes
keys1, keys2 = nodes1[:, 0], nodes2[:, 0]
@@ -33,7 +33,7 @@ class DefaultCrossover(BaseCrossover):
new_node_attrs = jnp.where(
jnp.isnan(node_attrs1) | jnp.isnan(node_attrs2), # one of them is nan
node_attrs1, # not homologous genes or both nan, use the value of nodes1(winner)
vmap(self.genome.node_gene.crossover, in_axes=(None, 0, 0, 0))(
vmap(genome.node_gene.crossover, in_axes=(None, 0, 0, 0))(
state, randkeys1, node_attrs1, node_attrs2
), # homologous or both nan
)
@@ -49,7 +49,7 @@ class DefaultCrossover(BaseCrossover):
new_conn_attrs = jnp.where(
jnp.isnan(conns_attrs1) | jnp.isnan(conns_attrs2),
conns_attrs1, # not homologous genes or both nan, use the value of conns1(winner)
vmap(self.genome.conn_gene.crossover, in_axes=(None, 0, 0, 0))(
vmap(genome.conn_gene.crossover, in_axes=(None, 0, 0, 0))(
state, randkeys2, conns_attrs1, conns_attrs2
), # homologous or both nan
)

7
src/tensorneat/genome/operations/distance/base.py
View File

@@ -3,12 +3,7 @@ from tensorneat.common import StatefulBaseClass, State
class BaseDistance(StatefulBaseClass):
def setup(self, state=State(), genome = None):
assert genome is not None, "genome should not be None"
self.genome = genome
return state
def __call__(self, state, nodes1, nodes2, conns1, conns2):
def __call__(self, state, genome, nodes1, nodes2, conns1, conns2):
"""
The distance between two genomes
"""

14
src/tensorneat/genome/operations/distance/default.py
View File

@@ -13,16 +13,16 @@ class DefaultDistance(BaseDistance):
self.compatibility_disjoint = compatibility_disjoint
self.compatibility_weight = compatibility_weight
def __call__(self, state, nodes1, conns1, nodes2, conns2):
def __call__(self, state, genome, nodes1, conns1, nodes2, conns2):
"""
The distance between two genomes
"""
d = self.node_distance(state, nodes1, nodes2) + self.conn_distance(
state, conns1, conns2
d = self.node_distance(state, genome, nodes1, nodes2) + self.conn_distance(
state, genome, conns1, conns2
)
return d
def node_distance(self, state, nodes1, nodes2):
def node_distance(self, state, genome, nodes1, nodes2):
"""
The distance of the nodes part for two genomes
"""
@@ -50,7 +50,7 @@ class DefaultDistance(BaseDistance):
# calculate the distance of homologous nodes
fr_attrs = vmap(extract_node_attrs)(fr)
sr_attrs = vmap(extract_node_attrs)(sr)
hnd = vmap(self.genome.node_gene.distance, in_axes=(None, 0, 0))(
hnd = vmap(genome.node_gene.distance, in_axes=(None, 0, 0))(
state, fr_attrs, sr_attrs
) # homologous node distance
hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
@@ -65,7 +65,7 @@ class DefaultDistance(BaseDistance):
return val
def conn_distance(self, state, conns1, conns2):
def conn_distance(self, state, genome, conns1, conns2):
"""
The distance of the conns part for two genomes
"""
@@ -89,7 +89,7 @@ class DefaultDistance(BaseDistance):
fr_attrs = vmap(extract_conn_attrs)(fr)
sr_attrs = vmap(extract_conn_attrs)(sr)
hcd = vmap(self.genome.conn_gene.distance, in_axes=(None, 0, 0))(
hcd = vmap(genome.conn_gene.distance, in_axes=(None, 0, 0))(
state, fr_attrs, sr_attrs
) # homologous connection distance
hcd = jnp.where(jnp.isnan(hcd), 0, hcd)

7
src/tensorneat/genome/operations/mutation/base.py
View File

@@ -3,10 +3,5 @@ from tensorneat.common import StatefulBaseClass, State
class BaseMutation(StatefulBaseClass):
def setup(self, state=State(), genome = None):
assert genome is not None, "genome should not be None"
self.genome = genome
return state
def __call__(self, state, randkey, nodes, conns, new_node_key):
def __call__(self, state, genome, randkey, nodes, conns, new_node_key):
raise NotImplementedError

42
src/tensorneat/genome/operations/mutation/default.py
View File

@@ -33,17 +33,17 @@ class DefaultMutation(BaseMutation):
self.node_add = node_add
self.node_delete = node_delete
def __call__(self, state, randkey, nodes, conns, new_node_key):
def __call__(self, state, genome, randkey, nodes, conns, new_node_key):
k1, k2 = jax.random.split(randkey)
nodes, conns = self.mutate_structure(
state, k1, nodes, conns, new_node_key
state, genome, k1, nodes, conns, new_node_key
)
nodes, conns = self.mutate_values(state, k2, nodes, conns)
nodes, conns = self.mutate_values(state, genome, k2, nodes, conns)
return nodes, conns
def mutate_structure(self, state, randkey, nodes, conns, new_node_key):
def mutate_structure(self, state, genome, randkey, nodes, conns, new_node_key):
def mutate_add_node(key_, nodes_, conns_):
"""
add a node while do not influence the output of the network
@@ -62,7 +62,7 @@ class DefaultMutation(BaseMutation):
# add a new node with identity attrs
new_nodes = add_node(
nodes_, new_node_key, self.genome.node_gene.new_identity_attrs(state)
nodes_, new_node_key, genome.node_gene.new_identity_attrs(state)
)
# add two new connections
@@ -71,7 +71,7 @@ class DefaultMutation(BaseMutation):
new_conns,
i_key,
new_node_key,
self.genome.conn_gene.new_identity_attrs(state),
genome.conn_gene.new_identity_attrs(state),
)
# second is with the origin attrs
new_conns = add_conn(
@@ -97,8 +97,8 @@ class DefaultMutation(BaseMutation):
key, idx = self.choose_node_key(
key_,
nodes_,
self.genome.input_idx,
self.genome.output_idx,
genome.input_idx,
genome.output_idx,
allow_input_keys=False,
allow_output_keys=False,
)
@@ -136,8 +136,8 @@ class DefaultMutation(BaseMutation):
i_key, from_idx = self.choose_node_key(
k1_,
nodes_,
self.genome.input_idx,
self.genome.output_idx,
genome.input_idx,
genome.output_idx,
allow_input_keys=True,
allow_output_keys=True,
)
@@ -146,8 +146,8 @@ class DefaultMutation(BaseMutation):
o_key, to_idx = self.choose_node_key(
k2_,
nodes_,
self.genome.input_idx,
self.genome.output_idx,
genome.input_idx,
genome.output_idx,
allow_input_keys=False,
allow_output_keys=True,
)
@@ -161,10 +161,10 @@ class DefaultMutation(BaseMutation):
def successful():
# add a connection with zero attrs
return nodes_, add_conn(
conns_, i_key, o_key, self.genome.conn_gene.new_zero_attrs(state)
conns_, i_key, o_key, genome.conn_gene.new_zero_attrs(state)
)
if self.genome.network_type == "feedforward":
if genome.network_type == "feedforward":
u_conns = unflatten_conns(nodes_, conns_)
conns_exist = u_conns != I_INF
is_cycle = check_cycles(nodes_, conns_exist, from_idx, to_idx)
@@ -175,7 +175,7 @@ class DefaultMutation(BaseMutation):
successful,
)
elif self.genome.network_type == "recurrent":
elif genome.network_type == "recurrent":
return jax.lax.cond(
is_already_exist | (remain_conn_space < 1),
nothing,
@@ -183,7 +183,7 @@ class DefaultMutation(BaseMutation):
)
else:
raise ValueError(f"Invalid network type: {self.genome.network_type}")
raise ValueError(f"Invalid network type: {genome.network_type}")
def mutate_delete_conn(key_, nodes_, conns_):
# randomly choose a connection
@@ -223,19 +223,19 @@ class DefaultMutation(BaseMutation):
return nodes, conns
def mutate_values(self, state, randkey, nodes, conns):
def mutate_values(self, state, genome, randkey, nodes, conns):
k1, k2 = jax.random.split(randkey)
nodes_randkeys = jax.random.split(k1, num=self.genome.max_nodes)
conns_randkeys = jax.random.split(k2, num=self.genome.max_conns)
nodes_randkeys = jax.random.split(k1, num=genome.max_nodes)
conns_randkeys = jax.random.split(k2, num=genome.max_conns)
node_attrs = vmap(extract_node_attrs)(nodes)
new_node_attrs = vmap(self.genome.node_gene.mutate, in_axes=(None, 0, 0))(
new_node_attrs = vmap(genome.node_gene.mutate, in_axes=(None, 0, 0))(
state, nodes_randkeys, node_attrs
)
new_nodes = vmap(set_node_attrs)(nodes, new_node_attrs)
conn_attrs = vmap(extract_conn_attrs)(conns)
new_conn_attrs = vmap(self.genome.conn_gene.mutate, in_axes=(None, 0, 0))(
new_conn_attrs = vmap(genome.conn_gene.mutate, in_axes=(None, 0, 0))(
state, conns_randkeys, conn_attrs
)
new_conns = vmap(set_conn_attrs)(conns, new_conn_attrs)

2
src/tensorneat/pipeline.py
View File

@@ -184,6 +184,6 @@ class Pipeline(StatefulBaseClass):
def show(self, state, best, *args, **kwargs):
transformed = self.algorithm.transform(state, best)
self.problem.show(
return self.problem.show(
state, state.randkey, self.algorithm.forward, transformed, *args, **kwargs
)