From 5bd6e5c357de942f5631aca315d08d1f7159f59b Mon Sep 17 00:00:00 2001 From: wls2002 Date: Thu, 30 May 2024 19:44:52 +0800 Subject: [PATCH] add "update_by_batch" in gene; add flatten_conns as an inverse function for unflatten_conns; add "test_flatten.ipynb" as test for them. --- tensorneat/algorithm/neat/gene/base.py | 3 + tensorneat/algorithm/neat/gene/conn/base.py | 7 + tensorneat/algorithm/neat/gene/node/base.py | 9 + .../gene/node/default_without_response.py | 8 +- .../algorithm/neat/gene/node/normalized.py | 204 +++++++++++++++++ tensorneat/algorithm/neat/genome/base.py | 6 + tensorneat/algorithm/neat/genome/default.py | 3 + tensorneat/test/test_flatten.ipynb | 207 ++++++++++++++++++ tensorneat/utils/tools.py | 45 +++- 9 files changed, 481 insertions(+), 11 deletions(-) create mode 100644 tensorneat/algorithm/neat/gene/node/normalized.py create mode 100644 tensorneat/test/test_flatten.ipynb diff --git a/tensorneat/algorithm/neat/gene/base.py b/tensorneat/algorithm/neat/gene/base.py index bef2182..99a57da 100644 --- a/tensorneat/algorithm/neat/gene/base.py +++ b/tensorneat/algorithm/neat/gene/base.py @@ -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) diff --git a/tensorneat/algorithm/neat/gene/conn/base.py b/tensorneat/algorithm/neat/gene/conn/base.py index 9819129..17e6cb8 100644 --- a/tensorneat/algorithm/neat/gene/conn/base.py +++ b/tensorneat/algorithm/neat/gene/conn/base.py @@ -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, + ) diff --git a/tensorneat/algorithm/neat/gene/node/base.py b/tensorneat/algorithm/neat/gene/node/base.py index 55361f4..d694a44 100644 --- a/tensorneat/algorithm/neat/gene/node/base.py +++ b/tensorneat/algorithm/neat/gene/node/base.py @@ -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, + ) diff --git a/tensorneat/algorithm/neat/gene/node/default_without_response.py b/tensorneat/algorithm/neat/gene/node/default_without_response.py index 32c9419..9f4fd5b 100644 --- a/tensorneat/algorithm/neat/gene/node/default_without_response.py +++ b/tensorneat/algorithm/neat/gene/node/default_without_response.py @@ -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): diff --git a/tensorneat/algorithm/neat/gene/node/normalized.py b/tensorneat/algorithm/neat/gene/node/normalized.py new file mode 100644 index 0000000..94f2558 --- /dev/null +++ b/tensorneat/algorithm/neat/gene/node/normalized.py @@ -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 diff --git a/tensorneat/algorithm/neat/genome/base.py b/tensorneat/algorithm/neat/genome/base.py index 7f7bf39..c78c7f1 100644 --- a/tensorneat/algorithm/neat/genome/base.py +++ b/tensorneat/algorithm/neat/genome/base.py @@ -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 diff --git a/tensorneat/algorithm/neat/genome/default.py b/tensorneat/algorithm/neat/genome/default.py index a2b6a50..b0ff770 100644 --- a/tensorneat/algorithm/neat/genome/default.py +++ b/tensorneat/algorithm/neat/genome/default.py @@ -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 diff --git a/tensorneat/test/test_flatten.ipynb b/tensorneat/test/test_flatten.ipynb new file mode 100644 index 0000000..de1d77e --- /dev/null +++ b/tensorneat/test/test_flatten.ipynb @@ -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 +} diff --git a/tensorneat/utils/tools.py b/tensorneat/utils/tools.py index 5f9bc11..f6f1de1 100644 --- a/tensorneat/utils/tools.py +++ b/tensorneat/utils/tools.py @@ -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