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

add "update_by_batch" in gene;

Browse Source 
add flatten_conns as an inverse function for unflatten_conns;
add "test_flatten.ipynb" as test for them.
This commit is contained in:
wls2002
2024-05-30 19:44:52 +08:00
parent cd92f411dc
commit 5bd6e5c357
9 changed files with 481 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tensorneat/algorithm/neat/gene/base.py
View File

@@ -32,6 +32,9 @@ class BaseGene:
def forward(self, state, attrs, inputs):
raise NotImplementedError
def update_by_batch(self, state, attrs, batch_inputs):
raise NotImplementedError
@property
def length(self):
return len(self.fixed_attrs) + len(self.custom_attrs)

7
tensorneat/algorithm/neat/gene/conn/base.py
View File

@@ -29,3 +29,10 @@ class BaseConnGene(BaseGene):
def forward(self, state, attrs, inputs):
raise NotImplementedError
def update_by_batch(self, state, attrs, batch_inputs):
# default: do not update attrs, but to calculate batch_res
return (
jax.vmap(self.forward, in_axes=(None, None, 0))(state, attrs, batch_inputs),
attrs,
)

9
tensorneat/algorithm/neat/gene/node/base.py
View File

@@ -18,3 +18,12 @@ class BaseNodeGene(BaseGene):
def forward(self, state, attrs, inputs, is_output_node=False):
raise NotImplementedError
def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False):
# default: do not update attrs, but to calculate batch_res
return (
jax.vmap(self.forward, in_axes=(None, None, 0, None))(
state, attrs, batch_inputs, is_output_node
),
attrs,
)

8
tensorneat/algorithm/neat/gene/node/default_without_response.py
View File

@@ -76,11 +76,11 @@ class NodeGeneWithoutResponse(BaseNodeGene):
)
act = mutate_int(
k3, node[3], self.activation_indices, self.activation_replace_rate
k3, node[2], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int(
k4, node[4], self.aggregation_indices, self.aggregation_replace_rate
k4, node[3], self.aggregation_indices, self.aggregation_replace_rate
)
return jnp.array([index, bias, act, agg])
@@ -88,8 +88,8 @@ class NodeGeneWithoutResponse(BaseNodeGene):
def distance(self, state, node1, node2):
return (
jnp.abs(node1[1] - node2[1]) # bias
+ (node1[3] != node2[3]) # activation
+ (node1[4] != node2[4]) # aggregation
+ (node1[2] != node2[2]) # activation
+ (node1[3] != node2[3]) # aggregation
)
def forward(self, state, attrs, inputs, is_output_node=False):

204
tensorneat/algorithm/neat/gene/node/normalized.py Normal file
View File

@@ -0,0 +1,204 @@
from typing import Tuple
import jax, jax.numpy as jnp
from utils import Act, Agg, act, agg, mutate_int, mutate_float
from . import BaseNodeGene
class NormalizedNode(BaseNodeGene):
"""
Node with normalization before activation.
"""
# alpha and beta is used for normalization, just like BatchNorm
# norm: (data - mean) / (std + eps) * alpha + beta
custom_attrs = ["bias", "aggregation", "activation", "mean", "std", "alpha", "beta"]
eps = 1e-6
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,
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,
alpha_init_mean: float = 0.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_std: float = 1.0,
beta_mutate_power: float = 0.5,
beta_mutate_rate: float = 0.7,
beta_replace_rate: float = 0.1,
):
super().__init__()
self.bias_init_mean = bias_init_mean
self.bias_init_std = bias_init_std
self.bias_mutate_power = bias_mutate_power
self.bias_mutate_rate = bias_mutate_rate
self.bias_replace_rate = bias_replace_rate
self.activation_default = activation_options.index(activation_default)
self.activation_options = activation_options
self.activation_indices = jnp.arange(len(activation_options))
self.activation_replace_rate = activation_replace_rate
self.aggregation_default = aggregation_options.index(aggregation_default)
self.aggregation_options = aggregation_options
self.aggregation_indices = jnp.arange(len(aggregation_options))
self.aggregation_replace_rate = aggregation_replace_rate
self.alpha_init_mean = alpha_init_mean
self.alpha_init_std = alpha_init_std
self.alpha_mutate_power = alpha_mutate_power
self.alpha_mutate_rate = alpha_mutate_rate
self.alpha_replace_rate = alpha_replace_rate
self.beta_init_mean = beta_init_mean
self.beta_init_std = beta_init_std
self.beta_mutate_power = beta_mutate_power
self.beta_mutate_rate = beta_mutate_rate
self.beta_replace_rate = beta_replace_rate
def new_custom_attrs(self, state):
return jnp.array(
[
self.bias_init_mean,
self.activation_default,
self.aggregation_default,
0, # mean
1, # std
self.alpha_init_mean, # alpha
self.beta_init_mean, # beta
]
)
def new_random_attrs(self, state, randkey):
k1, k2, k3, k4, k5, k6 = jax.random.split(randkey, num=6)
bias = jax.random.normal(k1, ()) * self.bias_init_std + self.bias_init_mean
act = jax.random.randint(k3, (), 0, len(self.activation_options))
agg = jax.random.randint(k4, (), 0, len(self.aggregation_options))
mean = 0
std = 1
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])
def mutate(self, state, randkey, node):
k1, k2, k3, k4, k5, k6 = jax.random.split(state.randkey, num=6)
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,
)
act = mutate_int(
k3, node[2], self.activation_indices, self.activation_replace_rate
)
agg = mutate_int(
k4, node[3], self.aggregation_indices, self.aggregation_replace_rate
)
mean = node[4]
std = node[5]
alpha = mutate_float(
k5,
node[6],
self.alpha_init_mean,
self.alpha_init_std,
self.alpha_mutate_power,
self.alpha_mutate_rate,
self.alpha_replace_rate,
)
beta = mutate_float(
k6,
node[7],
self.beta_init_mean,
self.beta_init_std,
self.beta_mutate_power,
self.beta_mutate_rate,
self.beta_replace_rate,
)
return jnp.array([index, bias, act, agg, mean, std, alpha, beta])
def distance(self, state, node1, node2):
return (
jnp.abs(node1[1] - node2[1]) # bias
+ (node1[2] != node2[2]) # activation
+ (node1[3] != node2[3]) # aggregation
+ (node1[6] - node2[6]) # alpha
+ (node1[7] - node2[7]) # beta
)
def forward(self, state, attrs, inputs, is_output_node=False):
"""
post_act = (agg(inputs) + bias - mean) / std * alpha + beta
"""
bias, act_idx, agg_idx, mean, std, alpha, beta = attrs
z = agg(agg_idx, inputs, self.aggregation_options)
z = bias + z
z = (z - mean) / (std + self.eps) * alpha + beta # normalization
# 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)
)
return z
def update_by_batch(self, state, attrs, batch_inputs, is_output_node=False):
bias, act_idx, agg_idx, mean, std, alpha, beta = attrs
batch_z = jax.vmap(agg, in_axes=(None, 0, None))(
agg_idx, batch_inputs, self.aggregation_options
)
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))
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))
/ valid_values_count
)
batch_z = (batch_z - mean) / (std + self.eps) * alpha + beta # normalization
batch_z = jax.lax.cond(
is_output_node,
lambda: batch_z,
lambda: jax.vmap(act, in_axes=(None, 0, None))(
act_idx, batch_z, self.activation_options
),
)
# update mean and std to the attrs
attrs = attrs.at[3].set(mean)
attrs = attrs.at[4].set(std)
return batch_z, attrs

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

@@ -120,3 +120,9 @@ class BaseGenome:
conns = conns.at[: len(conn_keys), 3:].set(random_conn_attrs)
return nodes, conns
def update_by_batch(self, state, batch_input, nodes, conns):
"""
Update the genome by a batch of data.
"""
raise NotImplementedError

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

@@ -93,3 +93,6 @@ class DefaultGenome(BaseGenome):
return vals[self.output_idx]
else:
return self.output_transform(vals[self.output_idx])
def update_by_batch(self, state, batch_input, nodes, conns):
pass

207
tensorneat/test/test_flatten.ipynb Normal file
View File

@@ -0,0 +1,207 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "initial_id",
"metadata": {
"collapsed": true,
"ExecuteTime": {
"end_time": "2024-05-30T11:40:55.584592400Z",
"start_time": "2024-05-30T11:40:53.016051600Z"
}
},
"outputs": [],
"source": [
"from algorithm.neat.genome import DefaultGenome\n",
"from utils.tools import flatten_conns, unflatten_conns\n",
"import jax, jax.numpy as jnp\n",
"from jax import vmap"
]
},
{
"cell_type": "code",
"execution_count": 2,
"outputs": [
{
"data": {
"text/plain": "((10, 5), (10, 4))"
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"genome = DefaultGenome(num_inputs=3, num_outputs=2, max_nodes=10, max_conns=10)\n",
"state = genome.setup()\n",
"key = jax.random.PRNGKey(0)\n",
"nodes, conns = genome.initialize(state, key)\n",
"nodes.shape, conns.shape"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-05-30T11:40:59.021858400Z",
"start_time": "2024-05-30T11:40:55.592593400Z"
}
},
"id": "89fb5cd0e77a028d"
},
{
"cell_type": "code",
"execution_count": 3,
"outputs": [
{
"data": {
"text/plain": "(2, 10, 10)"
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"unflatten = unflatten_conns(nodes, conns)\n",
"unflatten.shape"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-05-30T11:40:59.472701700Z",
"start_time": "2024-05-30T11:40:59.021858400Z"
}
},
"id": "aaa88227bbf29936"
},
{
"cell_type": "code",
"execution_count": 4,
"outputs": [
{
"data": {
"text/plain": "(Array([[ 0. , 5. , 1. , -0.41923347],\n [ 1. , 5. , 1. , -3.1815007 ],\n [ 2. , 5. , 1. , 0.5184341 ],\n [ 5. , 3. , 1. , -1.9784615 ],\n [ 5. , 4. , 1. , 0.7132204 ],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan]], dtype=float32, weak_type=True),\n Array([[ 0. , 5. , 1. , -0.41923347],\n [ 1. , 5. , 1. , -3.1815007 ],\n [ 2. , 5. , 1. , 0.5184341 ],\n [ 5. , 3. , 1. , -1.9784615 ],\n [ 5. , 4. , 1. , 0.7132204 ],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan],\n [ nan, nan, nan, nan]], dtype=float32))"
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# single flatten\n",
"flatten = flatten_conns(nodes, unflatten, C=10)\n",
"conns, flatten"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-05-30T11:41:00.308954100Z",
"start_time": "2024-05-30T11:40:59.469541500Z"
}
},
"id": "f2c65de38fdcff8f"
},
{
"cell_type": "code",
"execution_count": 8,
"outputs": [
{
"data": {
"text/plain": "((3, 10, 5), (3, 10, 4))"
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"# batch_flatten\n",
"key = jax.random.PRNGKey(0)\n",
"keys = jax.random.split(key, 3)\n",
"pop_nodes, pop_conns = jax.vmap(genome.initialize, in_axes=(None, 0))(state, keys)\n",
"pop_nodes.shape, pop_conns.shape"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-05-30T11:43:09.287012800Z",
"start_time": "2024-05-30T11:43:09.230179800Z"
}
},
"id": "fe89b178b721d656"
},
{
"cell_type": "code",
"execution_count": 9,
"outputs": [
{
"data": {
"text/plain": "(3, 2, 10, 10)"
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"pop_unflatten = jax.vmap(unflatten_conns)(pop_nodes, pop_conns)\n",
"pop_unflatten.shape"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-05-30T11:43:10.004429100Z",
"start_time": "2024-05-30T11:43:09.404949800Z"
}
},
"id": "14bbb257e5ddeab"
},
{
"cell_type": "code",
"execution_count": 10,
"outputs": [
{
"data": {
"text/plain": "(3, 10, 4)"
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"flatten = jax.vmap(flatten_conns, in_axes=(0, 0, None))(pop_nodes, pop_unflatten, 10)\n",
"flatten.shape"
],
"metadata": {
"collapsed": false,
"ExecuteTime": {
"end_time": "2024-05-30T11:43:39.983690700Z",
"start_time": "2024-05-30T11:43:39.208549Z"
}
},
"id": "8e5cdf6140c81dc0"
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 2
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

45
tensorneat/utils/tools.py
View File

@@ -13,24 +13,55 @@ def unflatten_conns(nodes, conns):
connection length, N means the number of nodes, C means the number of connections
returns the un_flattened connections with shape (CL-2, N, N)
"""
N = nodes.shape[0]
CL = conns.shape[1]
N = nodes.shape[0] # max_nodes
CL = conns.shape[1] # connection length = (fix_attrs + custom_attrs)
node_keys = nodes[:, 0]
i_keys, o_keys = conns[:, 0], conns[:, 1]
def key_to_indices(key, keys):
return fetch_first(key == keys)
i_idxs = vmap(key_to_indices, in_axes=(0, None))(i_keys, node_keys)
o_idxs = vmap(key_to_indices, in_axes=(0, None))(o_keys, node_keys)
res = jnp.full((CL - 2, N, N), jnp.nan)
unflatten = jnp.full((CL - 2, N, N), jnp.nan)
# Is interesting that jax use clip when attach data in array
# however, it will do nothing set values in an array
# put all attributes include enable in res
res = res.at[:, i_idxs, o_idxs].set(conns[:, 2:].T)
unflatten = unflatten.at[:, i_idxs, o_idxs].set(conns[:, 2:].T)
assert unflatten.shape == (CL - 2, N, N)
return res
return unflatten
def key_to_indices(key, keys):
return fetch_first(key == keys)
def flatten_conns(nodes, unflatten, C):
"""
the inverse function of unflatten_conns
transform the unflatten conn (CL-2, N, N) to (C, CL)
"""
N = nodes.shape[0]
CL = unflatten.shape[0] + 2
node_keys = nodes[:, 0]
def extract_conn(i, j):
return jnp.where(
jnp.isnan(unflatten[0, i, j]),
jnp.nan,
jnp.concatenate([jnp.array([node_keys[i], node_keys[j]]), unflatten[:, i, j]]),
)
x, y = jnp.meshgrid(jnp.arange(N), jnp.arange(N), indexing="ij")
conns = vmap(extract_conn)(x.flatten(), y.flatten())
assert conns.shape == (N * N, CL)
# put nan to the tail of the conns
sorted_idx = jnp.argsort(conns[:, 0])
sorted_conn = conns[sorted_idx]
# truncate the conns to the number of connections
conns = sorted_conn[:C]
assert conns.shape == (C, CL)
return conns
@jit