https://github.com/shyamsn97/jax-nca
Neural Cellular Automata implemented with Jax
https://github.com/shyamsn97/jax-nca
Last synced: 8 months ago
JSON representation
Neural Cellular Automata implemented with Jax
- Host: GitHub
- URL: https://github.com/shyamsn97/jax-nca
- Owner: shyamsn97
- License: mit
- Created: 2022-04-21T19:55:25.000Z (over 3 years ago)
- Default Branch: main
- Last Pushed: 2022-04-21T22:58:10.000Z (over 3 years ago)
- Last Synced: 2024-10-11T09:27:09.929Z (about 1 year ago)
- Language: Jupyter Notebook
- Size: 551 KB
- Stars: 6
- Watchers: 3
- Forks: 2
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
- awesome-neural-cellular-automata - JAX-NCA - NCA Implementation in JAX (Flax Linen) (Implementations / Growing Neural Cellular Automata)
README
# Neural Cellular Automata (Based on https://distill.pub/2020/growing-ca/) implemented in Jax (Flax)
---
## Installation
from source:
```
git clone git@github.com:shyamsn97/jax-nca.git
cd jax-nca
python setup.py install
```
from PYPI
```
pip install jax-nca
```
---
## How do NCAs work?
For more information, view the awesome article https://distill.pub/2020/growing-ca/ -- Mordvintsev, et al., "Growing Neural Cellular Automata", Distill, 2020
Image below describes a single update step: https://github.com/distillpub/post--growing-ca/blob/master/public/figures/model.svg
---
## Why Jax?
Note: This project served as a nice introduction to jax, so its performance can probably be improved
NCAs are autoregressive models like RNNs, where new states are calculated from previous ones. With jax, we can make these operations a lot more performant with `jax.lax.scan` and `jax.jit` (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)
Instead of writing the nca growth process as:
```python
def multi_step(params, nca, current_state, num_steps):
# params: parameters for NCA
# nca: Flax Module describing NCA
# current_state: Current NCA state
# num_steps: number of steps to run
for i in range(num_steps):
current_state = nca.apply(params, current_state)
return current_state
```
We can write this with `jax.lax.scan`
```python
def multi_step(params, nca, current_state, num_steps):
# params: parameters for NCA
# nca: Flax Module describing NCA
# current_state: Current NCA state
# num_steps: number of steps to run
def forward(carry, inp):
carry = nca.apply({"params": params}, carry)
return carry, carry
final_state, nca_states = jax.lax.scan(forward, current_state, None, length=num_steps)
return final_state
```
The actual multi_step implementation can be found here: https://github.com/shyamsn97/jax-nca/blob/main/jax_nca/nca.py#L103
---
## Usage
See [notebooks/Gecko.ipynb](notebooks/Gecko.ipynb) for a full example
Currently there's a bug with the stochastic update, so only `cell_fire_rate = 1.0` works at the moment
Creating and using NCA:
```python
class NCA(nn.Module):
num_hidden_channels: int
num_target_channels: int = 3
alpha_living_threshold: float = 0.1
cell_fire_rate: float = 1.0
trainable_perception: bool = False
alpha: float = 1.0
"""
num_hidden_channels: Number of hidden channels for each cell to use
num_target_channels: Number of target channels to be used
alpha_living_threshold: threshold to determine whether a cell lives or dies
cell_fire_rate: probability that a cell receives an update per step
trainable_perception: if true, instead of using sobel filters use a trainable conv net
alpha: scalar value to be multiplied to updates
"""
...
from jax_nca.nca import NCA
# usage
nca = NCA(
num_hidden_channels = 16,
num_target_channels = 3,
trainable_perception = False,
cell_fire_rate = 1.0,
alpha_living_threshold = 0.1
)
nca_seed = nca.create_seed(
nca.num_hidden_channels, nca.num_target_channels, shape=(64,64), batch_size=1
)
rng = jax.random.PRNGKey(0)
params = = nca.init(rng, nca_seed, rng)["params"]
update = nca.apply({"params":params}, nca_seed, jax.random.PRNGKey(10))
# multi step
final_state, nca_states = nca.multi_step(poarams, nca_seed, jax.random.PRNGKey(10), num_steps=32)
```
To train the NCA:
```python
from jax_nca.dataset import ImageDataset
from jax_nca.trainer import EmojiTrainer
dataset = ImageDataset(emoji='🦎', img_size=64)
nca = NCA(
num_hidden_channels = 16,
num_target_channels = 3,
trainable_perception = False,
cell_fire_rate = 1.0,
alpha_living_threshold = 0.1
)
trainer = EmojiTrainer(dataset, nca, n_damage=0)
trainer.train(100000, batch_size=8, seed=10, lr=2e-4, min_steps=64, max_steps=96)
# to access train state:
state = trainer.state
# save
nca.save(state.params, "saved_params")
# load params
loaded_params = nca.load("saved_params")
```