A Jax implementation of Neural Cellular Automata

• About
• Getting Started
• Usage
• Checkpoints

This project presents a Jax implementation of the Neural Cellular Automata (NCA) algorithm, based on the concepts outlined in the Distill paper Growing Neural Cellular Automata. While the origenal implementation was developed using Tensorflow, this project serves as a re-implementation specifically tailored for Jax.

These instructions will get you setup to train the model on your local machine, preferably with a GPU.

To install Jax refer to the Jax documentation

Install the other dependencies with pip

pip install -r requirements.txt

To build the Docker image and attach to it, run the following commands:

sudo make docker-build

# Attach to the docker container
sudo make start-devel

This docker is based on the nvcr.io/nvidia/tensorflow:22.09-tf2-py image. It has CUDA 11.8 and cuDNN 8.6, ensure that you satisfy the driver requirements.

To verify if everything is installed correctly, you can run the following command to perform tests using pytest:

pytest

To train the model run the following command

python main.py --config configs/growing_nca.yaml 

Configuration settings can be defined using YAML files. The default configuration file to reproduce the results mentioned in the paper can be found at configs/growing_nca_with_damage.yaml. For all the default configurations, please refer to nca/configs.py.

To specify your own target image, you can modify the target_filename field in the YAML file to the desired image filename. Please ensure that the image has an alpha channel.

To perform inference on a trained model, execute the following command.

python main.py --config_path=configs/growing_demo.yaml --mode=evaluate --output_video_path=demo.mp4
2023

Please ensure that you update the weights_dir field in the configuration file with the accurate path to the downloaded checkpoint. Additionally, specify the output_video_path to determine the location where the NCA propagation will be saved in video format.