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

use black format all files;

Browse Source 
remove "return state" for functions which will be executed in vmap;
recover randkey as args in mutation methods
This commit is contained in:
wls2002
2024-05-26 15:46:04 +08:00
parent 79d53ea7af
commit cf69b916af
38 changed files with 932 additions and 582 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
tensorneat/algorithm/neat/gene/node/default.py
View File

@@ -9,29 +9,26 @@ from . import BaseNodeGene
class DefaultNodeGene(BaseNodeGene):
"Default node gene, with the same behavior as in NEAT-python."
custom_attrs = ['bias', 'response', 'aggregation', 'activation']
custom_attrs = ["bias", "response", "aggregation", "activation"]
def __init__(
self,
bias_init_mean: float = 0.0,
bias_init_std: float = 1.0,
bias_mutate_power: float = 0.5,
bias_mutate_rate: float = 0.7,
bias_replace_rate: float = 0.1,
response_init_mean: float = 1.0,
response_init_std: float = 0.0,
response_mutate_power: float = 0.5,
response_mutate_rate: float = 0.7,
response_replace_rate: float = 0.1,
activation_default: callable = Act.sigmoid,
activation_options: Tuple = (Act.sigmoid,),
activation_replace_rate: float = 0.1,
aggregation_default: callable = Agg.sum,
aggregation_options: Tuple = (Agg.sum,),
aggregation_replace_rate: float = 0.1,
self,
bias_init_mean: float = 0.0,
bias_init_std: float = 1.0,
bias_mutate_power: float = 0.5,
bias_mutate_rate: float = 0.7,
bias_replace_rate: float = 0.1,
response_init_mean: float = 1.0,
response_init_std: float = 0.0,
response_mutate_power: float = 0.5,
response_mutate_rate: float = 0.7,
response_replace_rate: float = 0.1,
activation_default: callable = Act.sigmoid,
activation_options: Tuple = (Act.sigmoid,),
activation_replace_rate: float = 0.1,
aggregation_default: callable = Agg.sum,
aggregation_options: Tuple = (Agg.sum,),
aggregation_replace_rate: float = 0.1,
):
super().__init__()
self.bias_init_mean = bias_init_mean
@@ -56,33 +53,66 @@ class DefaultNodeGene(BaseNodeGene):
self.aggregation_indices = jnp.arange(len(aggregation_options))
self.aggregation_replace_rate = aggregation_replace_rate
def new_attrs(self, state):
return state, jnp.array(
[self.bias_init_mean, self.response_init_mean, self.activation_default, self.aggregation_default]
def new_custom_attrs(self, state):
return jnp.array(
[
self.bias_init_mean,
self.response_init_mean,
self.activation_default,
self.aggregation_default,
]
)
def mutate(self, state, node):
k1, k2, k3, k4, randkey = jax.random.split(state.randkey, num=5)
def new_random_attrs(self, state, randkey):
k1, k2, k3, k4 = jax.random.split(randkey, num=4)
bias = jax.random.normal(k1, ()) * self.bias_init_std + self.bias_init_mean
res = (
jax.random.normal(k2, ()) * self.response_init_std + self.response_init_mean
)
act = jax.random.randint(k3, (), 0, len(self.activation_options))
agg = jax.random.randint(k4, (), 0, len(self.aggregation_options))
return jnp.array([bias, res, act, agg])
def mutate(self, state, randkey, node):
k1, k2, k3, k4 = jax.random.split(state.randkey, num=4)
index = node[0]
bias = mutate_float(k1, node[1], self.bias_init_mean, self.bias_init_std,
self.bias_mutate_power, self.bias_mutate_rate, self.bias_replace_rate)
bias = mutate_float(
k1,
node[1],
self.bias_init_mean,
self.bias_init_std,
self.bias_mutate_power,
self.bias_mutate_rate,
self.bias_replace_rate,
)
res = mutate_float(k2, node[2], self.response_init_mean, self.response_init_std,
self.response_mutate_power, self.response_mutate_rate, self.response_replace_rate)
res = mutate_float(
k2,
node[2],
self.response_init_mean,
self.response_init_std,
self.response_mutate_power,
self.response_mutate_rate,
self.response_replace_rate,
)
act = mutate_int(k3, node[3], self.activation_indices, self.activation_replace_rate)
act = mutate_int(
k3, node[3], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int(k4, node[4], self.aggregation_indices, self.aggregation_replace_rate)
agg = mutate_int(
k4, node[4], self.aggregation_indices, self.aggregation_replace_rate
)
return state.update(randkey=randkey), jnp.array([index, bias, res, act, agg])
return jnp.array([index, bias, res, act, agg])
def distance(self, state, node1, node2):
return state, (
jnp.abs(node1[1] - node2[1]) +
jnp.abs(node1[2] - node2[2]) +
(node1[3] != node2[3]) +
(node1[4] != node2[4])
return (
jnp.abs(node1[1] - node2[1])
+ jnp.abs(node1[2] - node2[2])
+ (node1[3] != node2[3])
+ (node1[4] != node2[4])
)
def forward(self, state, attrs, inputs, is_output_node=False):
@@ -93,9 +123,7 @@ class DefaultNodeGene(BaseNodeGene):
# the last output node should not be activated
z = jax.lax.cond(
is_output_node,
lambda: z,
lambda: act(act_idx, z, self.activation_options)
is_output_node, lambda: z, lambda: act(act_idx, z, self.activation_options)
)
return state, z
return z