File size: 4,116 Bytes
f595bfc 434b57f f595bfc 434b57f f595bfc 434b57f f595bfc 434b57f f595bfc 434b57f f595bfc 434b57f f595bfc 434b57f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
---
tags:
- image-generation
---
# Neural Cellular Automata (Based on https://distill.pub/2020/growing-ca/) implemented in Jax (Flax)
## Installation
from source
```bash
git clone git@github.com:shyamsn97/jax-nca.git
cd jax-nca
python setup.py install
```
from PYPI
```bash
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?
<b> Note: This project served as a nice introduction to jax, so its performance can probably be improved </b>
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
<b> Currently there's a bug with the stochastic update, so only `cell_fire_rate = 1.0` works at the moment </b>
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")
``` |