Package: jax-nca
Neural Cellular Automata (Based onhttps://distill.pub/2020/growing-ca/) implemented in Jax (Flax)Installationfrom source:git clone git@github.com:shyamsn97/jax-nca.git
cd jax-nca
python setup.py installfrom PYPIpip install jax-ncaHow do NCAs work?For more information, view the awesome articlehttps://distill.pub/2020/growing-ca/-- Mordvintsev, et al., "Growing Neural Cellular Automata", Distill, 2020Image below describes a single update step:https://github.com/distillpub/post--growing-ca/blob/master/public/figures/model.svgWhy Jax?Note: This project served as a nice introduction to jax, so its performance can probably be improvedNCAs are autoregressive models like RNNs, where new states are calculated from previous ones. With jax, we can make these operations a lot more performant withjax.lax.scanandjax.jit(https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html)Instead of writing the nca growth process as:defmulti_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 runforiinrange(num_steps):current_state=nca.apply(params,current_state)returncurrent_stateWe can write this withjax.lax.scandefmulti_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 rundefforward(carry,inp):carry=nca.apply({"params":params},carry)returncarry,carryfinal_state,nca_states=jax.lax.scan(forward,current_state,None,length=num_steps)returnfinal_stateThe actual multi_step implementation can be found here:https://github.com/shyamsn97/jax-nca/blob/main/jax_nca/nca.py#L103UsageSeenotebooks/Gecko.ipynbfor a full exampleCurrently there's a bug with the stochastic update, so onlycell_fire_rate = 1.0works at the momentCreating and using NCA:classNCA(nn.Module):num_hidden_channels:intnum_target_channels:int=3alpha_living_threshold:float=0.1cell_fire_rate:float=1.0trainable_perception:bool=Falsealpha:float=1.0"""num_hidden_channels: Number of hidden channels for each cell to usenum_target_channels: Number of target channels to be usedalpha_living_threshold: threshold to determine whether a cell lives or diescell_fire_rate: probability that a cell receives an update per steptrainable_perception: if true, instead of using sobel filters use a trainable conv netalpha: scalar value to be multiplied to updates"""...fromjax_nca.ncaimportNCA# usagenca=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 stepfinal_state,nca_states=nca.multi_step(poarams,nca_seed,jax.random.PRNGKey(10),num_steps=32)To train the NCA:fromjax_nca.datasetimportImageDatasetfromjax_nca.trainerimportEmojiTrainerdataset=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# savenca.save(state.params,"saved_params")# load paramsloaded_params=nca.load("saved_params")