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 .base import BaseSubstrate
from tensorneat.genome.utils import set_conn_attrs 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): def __init__(self, num_inputs, num_outputs, coors, nodes, conns):
self.inputs = num_inputs self.inputs = num_inputs
self.outputs = num_outputs self.outputs = num_outputs
self.coors = jnp.array(coors) self.coors = np.array(coors)
self.nodes = jnp.array(nodes) self.nodes = np.array(nodes)
self.conns = jnp.array(conns) self.conns = np.array(conns)
def make_nodes(self, query_res): def make_nodes(self, query_res):
return self.nodes 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 from jax.tree_util import register_pytree_node_class
@@ -39,6 +41,15 @@ class State:
def __contains__(self, item): def __contains__(self, item):
return item in self.state_dict 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): def tree_flatten(self):
children = list(self.state_dict.values()) children = list(self.state_dict.values())
aux_data = list(self.state_dict.keys()) 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()): def setup(self, state=State()):
return 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): def show_config(self, registered_objects=None):
if registered_objects is None: # root call if registered_objects is None: # root call
registered_objects = [] registered_objects = []
@@ -47,27 +23,53 @@ class StatefulBaseClass:
config[str(key)] = str(value) config[str(key)] = str(value)
return config return config
@classmethod # TODO: Bug need be fixed
def load(cls, path: str, with_state: bool = False, warning: bool = True): # def save(self, state: Optional[State] = None, path: Optional[str] = None):
with open(path, "rb") as f: # if path is None:
obj = pickle.load(f) # time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
if with_state: # path = f"./{self.__class__.__name__} {time}.pkl"
if "aux_for_state" not in obj.__dict__: # if state is not None:
if warning: # self.__dict__["aux_for_state"] = state
warnings.warn( # with open(path, "wb") as f:
"This object does not have state to load, return empty state", # pickle.dump(self, f)
category=UserWarning,
) # def __getstate__(self):
return obj, State() # # only pickle the picklable attributes
state = obj.__dict__["aux_for_state"] # state = self.__dict__.copy()
del obj.__dict__["aux_for_state"] # non_picklable_keys = []
return obj, state # for key, value in state.items():
else: # try:
if "aux_for_state" in obj.__dict__: # pickle.dumps(value)
if warning: # except Exception as e:
warnings.warn( # print(f"Cannot pickle key {key}: {e}")
"This object has state to load, ignore it", # non_picklable_keys.append(key)
category=UserWarning,
) # for key in non_picklable_keys:
del obj.__dict__["aux_for_state"] # state.pop(key)
return obj
# 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()): def setup(self, state=State()):
state = self.node_gene.setup(state) state = self.node_gene.setup(state)
state = self.conn_gene.setup(state) state = self.conn_gene.setup(state)
state = self.mutation.setup(state, self) state = self.mutation.setup(state)
state = self.crossover.setup(state, self) state = self.crossover.setup(state)
state = self.distance.setup(state, self) state = self.distance.setup(state)
return state return state
def transform(self, state, nodes, conns): def transform(self, state, nodes, conns):
@@ -114,13 +114,13 @@ class BaseGenome(StatefulBaseClass):
raise NotImplementedError raise NotImplementedError
def execute_mutation(self, state, randkey, nodes, conns, new_node_key): 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): 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): 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): def initialize(self, state, randkey):
k1, k2 = jax.random.split(randkey) # k1 for nodes, k2 for conns 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 get_func_name
) )
from . import BaseNode from .base import BaseNode
class BiasNode(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): class BaseCrossover(StatefulBaseClass):
def setup(self, state=State(), genome = None): def __call__(self, state, genome, randkey, nodes1, nodes2, conns1, conns2):
assert genome is not None, "genome should not be None"
self.genome = genome
return state
def __call__(self, state, randkey, nodes1, nodes2, conns1, conns2):
raise NotImplementedError raise NotImplementedError

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

@@ -11,14 +11,14 @@ from ...utils import (
class DefaultCrossover(BaseCrossover): 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 use genome1 and genome2 to generate a new genome
notice that genome1 should have higher fitness than genome2 (genome1 is winner!) notice that genome1 should have higher fitness than genome2 (genome1 is winner!)
""" """
randkey1, randkey2 = jax.random.split(randkey, 2) randkey1, randkey2 = jax.random.split(randkey, 2)
randkeys1 = jax.random.split(randkey1, self.genome.max_nodes) randkeys1 = jax.random.split(randkey1, genome.max_nodes)
randkeys2 = jax.random.split(randkey2, self.genome.max_conns) randkeys2 = jax.random.split(randkey2, genome.max_conns)
# crossover nodes # crossover nodes
keys1, keys2 = nodes1[:, 0], nodes2[:, 0] keys1, keys2 = nodes1[:, 0], nodes2[:, 0]
@@ -33,7 +33,7 @@ class DefaultCrossover(BaseCrossover):
new_node_attrs = jnp.where( new_node_attrs = jnp.where(
jnp.isnan(node_attrs1) | jnp.isnan(node_attrs2), # one of them is nan 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) 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 state, randkeys1, node_attrs1, node_attrs2
), # homologous or both nan ), # homologous or both nan
) )
@@ -49,7 +49,7 @@ class DefaultCrossover(BaseCrossover):
new_conn_attrs = jnp.where( new_conn_attrs = jnp.where(
jnp.isnan(conns_attrs1) | jnp.isnan(conns_attrs2), jnp.isnan(conns_attrs1) | jnp.isnan(conns_attrs2),
conns_attrs1, # not homologous genes or both nan, use the value of conns1(winner) 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 state, randkeys2, conns_attrs1, conns_attrs2
), # homologous or both nan ), # 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): class BaseDistance(StatefulBaseClass):
def setup(self, state=State(), genome = None): def __call__(self, state, genome, nodes1, nodes2, conns1, conns2):
assert genome is not None, "genome should not be None"
self.genome = genome
return state
def __call__(self, state, nodes1, nodes2, conns1, conns2):
""" """
The distance between two genomes 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_disjoint = compatibility_disjoint
self.compatibility_weight = compatibility_weight 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 The distance between two genomes
""" """
d = self.node_distance(state, nodes1, nodes2) + self.conn_distance( d = self.node_distance(state, genome, nodes1, nodes2) + self.conn_distance(
state, conns1, conns2 state, genome, conns1, conns2
) )
return d 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 The distance of the nodes part for two genomes
""" """
@@ -50,7 +50,7 @@ class DefaultDistance(BaseDistance):
# calculate the distance of homologous nodes # calculate the distance of homologous nodes
fr_attrs = vmap(extract_node_attrs)(fr) fr_attrs = vmap(extract_node_attrs)(fr)
sr_attrs = vmap(extract_node_attrs)(sr) 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 state, fr_attrs, sr_attrs
) # homologous node distance ) # homologous node distance
hnd = jnp.where(jnp.isnan(hnd), 0, hnd) hnd = jnp.where(jnp.isnan(hnd), 0, hnd)
@@ -65,7 +65,7 @@ class DefaultDistance(BaseDistance):
return val 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 The distance of the conns part for two genomes
""" """
@@ -89,7 +89,7 @@ class DefaultDistance(BaseDistance):
fr_attrs = vmap(extract_conn_attrs)(fr) fr_attrs = vmap(extract_conn_attrs)(fr)
sr_attrs = vmap(extract_conn_attrs)(sr) 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 state, fr_attrs, sr_attrs
) # homologous connection distance ) # homologous connection distance
hcd = jnp.where(jnp.isnan(hcd), 0, hcd) 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): class BaseMutation(StatefulBaseClass):
def setup(self, state=State(), genome = None): def __call__(self, state, genome, randkey, nodes, conns, new_node_key):
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):
raise NotImplementedError 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_add = node_add
self.node_delete = node_delete 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) k1, k2 = jax.random.split(randkey)
nodes, conns = self.mutate_structure( 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 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_): def mutate_add_node(key_, nodes_, conns_):
""" """
add a node while do not influence the output of the network 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 # add a new node with identity attrs
new_nodes = add_node( 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 # add two new connections
@@ -71,7 +71,7 @@ class DefaultMutation(BaseMutation):
new_conns, new_conns,
i_key, i_key,
new_node_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 # second is with the origin attrs
new_conns = add_conn( new_conns = add_conn(
@@ -97,8 +97,8 @@ class DefaultMutation(BaseMutation):
key, idx = self.choose_node_key( key, idx = self.choose_node_key(
key_, key_,
nodes_, nodes_,
self.genome.input_idx, genome.input_idx,
self.genome.output_idx, genome.output_idx,
allow_input_keys=False, allow_input_keys=False,
allow_output_keys=False, allow_output_keys=False,
) )
@@ -136,8 +136,8 @@ class DefaultMutation(BaseMutation):
i_key, from_idx = self.choose_node_key( i_key, from_idx = self.choose_node_key(
k1_, k1_,
nodes_, nodes_,
self.genome.input_idx, genome.input_idx,
self.genome.output_idx, genome.output_idx,
allow_input_keys=True, allow_input_keys=True,
allow_output_keys=True, allow_output_keys=True,
) )
@@ -146,8 +146,8 @@ class DefaultMutation(BaseMutation):
o_key, to_idx = self.choose_node_key( o_key, to_idx = self.choose_node_key(
k2_, k2_,
nodes_, nodes_,
self.genome.input_idx, genome.input_idx,
self.genome.output_idx, genome.output_idx,
allow_input_keys=False, allow_input_keys=False,
allow_output_keys=True, allow_output_keys=True,
) )
@@ -161,10 +161,10 @@ class DefaultMutation(BaseMutation):
def successful(): def successful():
# add a connection with zero attrs # add a connection with zero attrs
return nodes_, add_conn( 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_) u_conns = unflatten_conns(nodes_, conns_)
conns_exist = u_conns != I_INF conns_exist = u_conns != I_INF
is_cycle = check_cycles(nodes_, conns_exist, from_idx, to_idx) is_cycle = check_cycles(nodes_, conns_exist, from_idx, to_idx)
@@ -175,7 +175,7 @@ class DefaultMutation(BaseMutation):
successful, successful,
) )
elif self.genome.network_type == "recurrent": elif genome.network_type == "recurrent":
return jax.lax.cond( return jax.lax.cond(
is_already_exist | (remain_conn_space < 1), is_already_exist | (remain_conn_space < 1),
nothing, nothing,
@@ -183,7 +183,7 @@ class DefaultMutation(BaseMutation):
) )
else: 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_): def mutate_delete_conn(key_, nodes_, conns_):
# randomly choose a connection # randomly choose a connection
@@ -223,19 +223,19 @@ class DefaultMutation(BaseMutation):
return nodes, conns 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) k1, k2 = jax.random.split(randkey)
nodes_randkeys = jax.random.split(k1, num=self.genome.max_nodes) nodes_randkeys = jax.random.split(k1, num=genome.max_nodes)
conns_randkeys = jax.random.split(k2, num=self.genome.max_conns) conns_randkeys = jax.random.split(k2, num=genome.max_conns)
node_attrs = vmap(extract_node_attrs)(nodes) 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 state, nodes_randkeys, node_attrs
) )
new_nodes = vmap(set_node_attrs)(nodes, new_node_attrs) new_nodes = vmap(set_node_attrs)(nodes, new_node_attrs)
conn_attrs = vmap(extract_conn_attrs)(conns) 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 state, conns_randkeys, conn_attrs
) )
new_conns = vmap(set_conn_attrs)(conns, new_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): def show(self, state, best, *args, **kwargs):
transformed = self.algorithm.transform(state, best) transformed = self.algorithm.transform(state, best)
self.problem.show( return self.problem.show(
state, state.randkey, self.algorithm.forward, transformed, *args, **kwargs state, state.randkey, self.algorithm.forward, transformed, *args, **kwargs
) )