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

add "update_by_batch" in genome;

Browse Source 
add "normalized" gene, which can do normalization before activation func.
add related test.
This commit is contained in:
wls2002
2024-05-30 23:12:11 +08:00
parent 5bd6e5c357
commit 3ea9986bd4
6 changed files with 403 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
tensorneat/algorithm/neat/gene/node/normalized.py
View File

@@ -29,12 +29,12 @@ class NormalizedNode(BaseNodeGene):
aggregation_default: callable = Agg.sum,
aggregation_options: Tuple = (Agg.sum,),
aggregation_replace_rate: float = 0.1,
alpha_init_mean: float = 0.0,
alpha_init_mean: float = 1.0,
alpha_init_std: float = 1.0,
alpha_mutate_power: float = 0.5,
alpha_mutate_rate: float = 0.7,
alpha_replace_rate: float = 0.1,
beta_init_mean: float = 1.0,
beta_init_mean: float = 0.0,
beta_init_std: float = 1.0,
beta_mutate_power: float = 0.5,
beta_mutate_rate: float = 0.7,
@@ -92,7 +92,7 @@ class NormalizedNode(BaseNodeGene):
alpha = jax.random.normal(k5, ()) * self.alpha_init_std + self.alpha_init_mean
beta = jax.random.normal(k6, ()) * self.beta_init_std + self.beta_init_mean
return jnp.array([bias, act, agg, 0, 1, alpha, beta])
return jnp.array([bias, act, agg, mean, std, alpha, beta])
def mutate(self, state, randkey, node):
k1, k2, k3, k4, k5, k6 = jax.random.split(state.randkey, num=6)
@@ -178,13 +178,13 @@ class NormalizedNode(BaseNodeGene):
batch_z = bias + batch_z
# calculate mean
valid_values_count = jnp.sum(~jnp.isnan(batch_inputs))
valid_values_sum = jnp.sum(jnp.where(jnp.isnan(batch_inputs), 0, batch_inputs))
valid_values_count = jnp.sum(~jnp.isnan(batch_z))
valid_values_sum = jnp.sum(jnp.where(jnp.isnan(batch_z), 0, batch_z))
mean = valid_values_sum / valid_values_count
# calculate std
std = jnp.sqrt(
jnp.sum(jnp.where(jnp.isnan(batch_inputs), 0, (batch_inputs - mean) ** 2))
jnp.sum(jnp.where(jnp.isnan(batch_z), 0, (batch_z - mean) ** 2))
/ valid_values_count
)

5
tensorneat/algorithm/neat/genome/base.py
View File

@@ -39,6 +39,9 @@ class BaseGenome:
def transform(self, state, nodes, conns):
raise NotImplementedError
def restore(self, state, transformed):
raise NotImplementedError
def forward(self, state, inputs, transformed):
raise NotImplementedError
@@ -121,7 +124,7 @@ class BaseGenome:
return nodes, conns
def update_by_batch(self, state, batch_input, nodes, conns):
def update_by_batch(self, state, batch_input, transformed):
"""
Update the genome by a batch of data.
"""

77
tensorneat/algorithm/neat/genome/default.py
View File

@@ -1,7 +1,7 @@
from typing import Callable
import jax, jax.numpy as jnp
from utils import unflatten_conns, topological_sort, I_INF
from utils import unflatten_conns, flatten_conns, topological_sort, I_INF
from . import BaseGenome
from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
@@ -53,17 +53,21 @@ class DefaultGenome(BaseGenome):
return seqs, nodes, u_conns
def forward(self, state, inputs, transformed):
cal_seqs, nodes, conns = transformed
def restore(self, state, transformed):
seqs, nodes, u_conns = transformed
conns = flatten_conns(nodes, u_conns, C=self.max_conns)
return nodes, conns
N = nodes.shape[0]
ini_vals = jnp.full((N,), jnp.nan)
def forward(self, state, inputs, transformed):
cal_seqs, nodes, u_conns = transformed
ini_vals = jnp.full((self.max_nodes,), jnp.nan)
ini_vals = ini_vals.at[self.input_idx].set(inputs)
nodes_attrs = nodes[:, 1:]
def cond_fun(carry):
values, idx = carry
return (idx < N) & (cal_seqs[idx] != I_INF)
return (idx < self.max_nodes) & (cal_seqs[idx] != I_INF)
def body_func(carry):
values, idx = carry
@@ -71,7 +75,7 @@ class DefaultGenome(BaseGenome):
def hit():
ins = jax.vmap(self.conn_gene.forward, in_axes=(None, 1, 0))(
state, conns[:, :, i], values
state, u_conns[:, :, i], values
)
z = self.node_gene.forward(
state,
@@ -80,6 +84,7 @@ class DefaultGenome(BaseGenome):
is_output_node=jnp.isin(i, self.output_idx),
)
new_values = values.at[i].set(z)
return new_values
# the val of input nodes is obtained by the task, not by calculation
@@ -94,5 +99,59 @@ class DefaultGenome(BaseGenome):
else:
return self.output_transform(vals[self.output_idx])
def update_by_batch(self, state, batch_input, nodes, conns):
pass
def update_by_batch(self, state, batch_input, transformed):
cal_seqs, nodes, u_conns = transformed
batch_size = batch_input.shape[0]
batch_ini_vals = jnp.full((batch_size, self.max_nodes), jnp.nan)
batch_ini_vals = batch_ini_vals.at[:, self.input_idx].set(batch_input)
nodes_attrs = nodes[:, 1:]
def cond_fun(carry):
batch_values, nodes_attrs_, u_conns_, idx = carry
return (idx < self.max_nodes) & (cal_seqs[idx] != I_INF)
def body_func(carry):
batch_values, nodes_attrs_, u_conns_, idx = carry
i = cal_seqs[idx]
def hit():
batch_ins, new_conn_attrs = jax.vmap(
self.conn_gene.update_by_batch, in_axes=(None, 1, 1), out_axes=(1, 1)
)(state, u_conns_[:, :, i], batch_values)
batch_z, new_node_attrs = self.node_gene.update_by_batch(
state,
nodes_attrs[i],
batch_ins,
is_output_node=jnp.isin(i, self.output_idx),
)
new_batch_values = batch_values.at[:, i].set(batch_z)
return (
new_batch_values,
nodes_attrs_.at[i].set(new_node_attrs),
u_conns_.at[:, :, i].set(new_conn_attrs),
)
(batch_values, nodes_attrs_, u_conns_) = jax.lax.cond(
jnp.isin(i, self.input_idx),
lambda: (batch_values, nodes_attrs_, u_conns_),
hit,
)
# the val of input nodes is obtained by the task, not by calculation
return batch_values, nodes_attrs_, u_conns_, idx + 1
batch_vals, nodes_attrs, u_conns, _ = jax.lax.while_loop(
cond_fun, body_func, (batch_ini_vals, nodes_attrs, u_conns, 0)
)
nodes = nodes.at[:, 1:].set(nodes_attrs)
new_transformed = (cal_seqs, nodes, u_conns)
if self.output_transform is None:
return batch_vals[:, self.output_idx], new_transformed
else:
return (
jax.vmap(self.output_transform)(batch_vals[:, self.output_idx]),
new_transformed,
)

12
tensorneat/algorithm/neat/genome/recurrent.py
View File

@@ -1,7 +1,7 @@
from typing import Callable
import jax, jax.numpy as jnp
from utils import unflatten_conns
from utils import unflatten_conns, flatten_conns
from . import BaseGenome
from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
@@ -54,11 +54,15 @@ class RecurrentGenome(BaseGenome):
return nodes, u_conns
def restore(self, state, transformed):
nodes, u_conns = transformed
conns = flatten_conns(nodes, u_conns, C=self.max_conns)
return nodes, conns
def forward(self, state, inputs, transformed):
nodes, conns = transformed
N = nodes.shape[0]
vals = jnp.full((N,), jnp.nan)
vals = jnp.full((self.max_nodes,), jnp.nan)
nodes_attrs = nodes[:, 1:] # remove index
def body_func(_, values):
@@ -73,7 +77,7 @@ class RecurrentGenome(BaseGenome):
)(state, conns, values)
# calculate nodes
is_output_nodes = jnp.isin(jnp.arange(N), self.output_idx)
is_output_nodes = jnp.isin(jnp.arange(self.max_nodes), self.output_idx)
values = jax.vmap(self.node_gene.forward, in_axes=(None, 0, 0, 0))(
state, nodes_attrs, node_ins.T, is_output_nodes
)