Views
No views yet
1git clone git@github.com:shyamsn97/jax-nca.git
2cd jax-nca
3python setup.py installpip install jax-ncajax.lax.scan and jax.jit (https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)1def multi_step(params, nca, current_state, num_steps):
2 # params: parameters for NCA
3 # nca: Flax Module describing NCA
4 # current_state: Current NCA state
5 # num_steps: number of steps to run
6
7 for i in range(num_steps):
8 current_state = nca.apply(params, current_state)
9 return current_statejax.lax.scan1def multi_step(params, nca, current_state, num_steps):
2 # params: parameters for NCA
3 # nca: Flax Module describing NCA
4 # current_state: Current NCA state
5 # num_steps: number of steps to run
6
7 def forward(carry, inp):
8 carry = nca.apply({"params": params}, carry)
9 return carry, carry
10
11 final_state, nca_states = jax.lax.scan(forward, current_state, None, length=num_steps)
12 return final_statecell_fire_rate = 1.0 works at the moment 1class NCA(nn.Module):
2 num_hidden_channels: int
3 num_target_channels: int = 3
4 alpha_living_threshold: float = 0.1
5 cell_fire_rate: float = 1.0
6 trainable_perception: bool = False
7 alpha: float = 1.0
8
9 """
10 num_hidden_channels: Number of hidden channels for each cell to use
11 num_target_channels: Number of target channels to be used
12 alpha_living_threshold: threshold to determine whether a cell lives or dies
13 cell_fire_rate: probability that a cell receives an update per step
14 trainable_perception: if true, instead of using sobel filters use a trainable conv net
15 alpha: scalar value to be multiplied to updates
16 """
17 ...
18
19from jax_nca.nca import NCA
20
21# usage
22nca = NCA(
23 num_hidden_channels = 16,
24 num_target_channels = 3,
25 trainable_perception = False,
26 cell_fire_rate = 1.0,
27 alpha_living_threshold = 0.1
28)
29
30nca_seed = nca.create_seed(
31 nca.num_hidden_channels, nca.num_target_channels, shape=(64,64), batch_size=1
32)
33rng = jax.random.PRNGKey(0)
34params = = nca.init(rng, nca_seed, rng)["params"]
35update = nca.apply({"params":params}, nca_seed, jax.random.PRNGKey(10))
36
37# multi step
38
39final_state, nca_states = nca.multi_step(poarams, nca_seed, jax.random.PRNGKey(10), num_steps=32)1from jax_nca.dataset import ImageDataset
2from jax_nca.trainer import EmojiTrainer
3
4
5dataset = ImageDataset(emoji='🦎', img_size=64)
6
7
8nca = NCA(
9 num_hidden_channels = 16,
10 num_target_channels = 3,
11 trainable_perception = False,
12 cell_fire_rate = 1.0,
13 alpha_living_threshold = 0.1
14)
15
16trainer = EmojiTrainer(dataset, nca, n_damage=0)
17
18trainer.train(100000, batch_size=8, seed=10, lr=2e-4, min_steps=64, max_steps=96)
19
20# to access train state:
21
22state = trainer.state
23
24# save
25nca.save(state.params, "saved_params")
26
27# load params
28loaded_params = nca.load("saved_params")
29