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

63
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_INT
from utils import unflatten_conns, topological_sort, I_INF
from . import BaseGenome
from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
@@ -10,18 +10,21 @@ from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
class DefaultGenome(BaseGenome):
"""Default genome class, with the same behavior as the NEAT-Python"""
network_type = 'feedforward'
network_type = "feedforward"
def __init__(self,
num_inputs: int,
num_outputs: int,
max_nodes=5,
max_conns=4,
node_gene: BaseNodeGene = DefaultNodeGene(),
conn_gene: BaseConnGene = DefaultConnGene(),
output_transform: Callable = None
):
super().__init__(num_inputs, num_outputs, max_nodes, max_conns, node_gene, conn_gene)
def __init__(
self,
num_inputs: int,
num_outputs: int,
max_nodes=5,
max_conns=4,
node_gene: BaseNodeGene = DefaultNodeGene(),
conn_gene: BaseConnGene = DefaultConnGene(),
output_transform: Callable = None,
):
super().__init__(
num_inputs, num_outputs, max_nodes, max_conns, node_gene, conn_gene
)
if output_transform is not None:
try:
@@ -38,7 +41,7 @@ class DefaultGenome(BaseGenome):
u_conns = jnp.where(conn_enable, u_conns[1:, :], jnp.nan)
seqs = topological_sort(nodes, conn_enable)
return state, seqs, nodes, u_conns
return seqs, nodes, u_conns
def forward(self, state, inputs, transformed):
cal_seqs, nodes, conns = transformed
@@ -49,32 +52,34 @@ class DefaultGenome(BaseGenome):
nodes_attrs = nodes[:, 1:]
def cond_fun(carry):
state_, values, idx = carry
return (idx < N) & (cal_seqs[idx] != I_INT)
values, idx = carry
return (idx < N) & (cal_seqs[idx] != I_INF)
def body_func(carry):
state_, values, idx = carry
values, idx = carry
i = cal_seqs[idx]
def hit():
s, ins = jax.vmap(self.conn_gene.forward,
in_axes=(None, 1, 0), out_axes=(None, 0))(state_, conns[:, :, i], values)
s, z = self.node_gene.forward(s, nodes_attrs[i], ins, is_output_node=jnp.isin(i, self.output_idx))
ins = jax.vmap(self.conn_gene.forward, in_axes=(None, 1, 0))(
state, conns[:, :, i], values
)
z = self.node_gene.forward(
state,
nodes_attrs[i],
ins,
is_output_node=jnp.isin(i, self.output_idx),
)
new_values = values.at[i].set(z)
return s, new_values
return new_values
# the val of input nodes is obtained by the task, not by calculation
state_, values = jax.lax.cond(
jnp.isin(i, self.input_idx),
lambda: (state_, values),
hit
)
values = jax.lax.cond(jnp.isin(i, self.input_idx), lambda: values, hit)
return state_, values, idx + 1
return values, idx + 1
state, vals, _ = jax.lax.while_loop(cond_fun, body_func, (state, ini_vals, 0))
vals, _ = jax.lax.while_loop(cond_fun, body_func, (ini_vals, 0))
if self.output_transform is None:
return state, vals[self.output_idx]
return vals[self.output_idx]
else:
return state, self.output_transform(vals[self.output_idx])
return self.output_transform(vals[self.output_idx])