Diffusing Neural Cellular Automata ¶

Installation¶

You will need Python 3.12 or later, and a working JAX installation. For example, you can install JAX with:

In [1]:

%pip install -U "jax[cuda]"

%pip install -U "jax[cuda]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Then, install CAX from PyPi:

In [2]:

%pip install -U "cax[examples]"

%pip install -U "cax[examples]"

/home/faldor_google_com/dev/cax/.venv/bin/python3: No module named pip

Note: you may need to restart the kernel to use updated packages.

Import¶

In [3]:

import time

import jax
import jax.numpy as jnp
import mediapy
import optax
from flax import nnx
from jax import Array

from cax.core import ComplexSystem
from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel
from cax.core.update import NCAUpdate
from cax.utils import clip_and_uint8, get_emoji_array, rgba_to_rgb

import time import jax import jax.numpy as jnp import mediapy import optax from flax import nnx from jax import Array from cax.core import ComplexSystem from cax.core.perceive import ConvPerceive, grad_kernel, identity_kernel from cax.core.update import NCAUpdate from cax.utils import clip_and_uint8, get_emoji_array, rgba_to_rgb

Configuration¶

In [4]:

seed = 0

channel_size = 64
num_kernels = 3
hidden_size = 256
cell_dropout_rate = 0.5

num_steps = 64
batch_size = 8
learning_rate = 1e-3

emoji = "🦎"
size = 40
pad_width = 16

key = jax.random.key(seed)
rngs = nnx.Rngs(seed)

seed = 0 channel_size = 64 num_kernels = 3 hidden_size = 256 cell_dropout_rate = 0.5 num_steps = 64 batch_size = 8 learning_rate = 1e-3 emoji = "🦎" size = 40 pad_width = 16 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

In [5]:

y = get_emoji_array(emoji, size, pad_width)

mediapy.show_image(y)

y = get_emoji_array(emoji, size, pad_width) mediapy.show_image(y)

Instantiate system¶

In [6]:

class DiffusingNCA(ComplexSystem):
    """Diffusing Neural Cellular Automata class."""

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize Diffusing NCA.

        Args:
                rngs: rng key.

        """
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            feature_group_count=channel_size,
            rngs=rngs,
        )
        self.update = NCAUpdate(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            hidden_layer_sizes=(hidden_size,),
            cell_dropout_rate=cell_dropout_rate,
            zeros_init=True,
            rngs=rngs,
        )

        # Initialize kernel with sobel filters
        kernel = jnp.concatenate(
            [identity_kernel(num_dims=2), grad_kernel(num_dims=2)], axis=-1
        )
        kernel = jnp.expand_dims(
            jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2
        )
        self.perceive.conv.kernel[...] = kernel

    def _step(self, state: Array, input: Array | None = None) -> Array:
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        return next_state

    @nnx.jit
    def render(self, state):
        """Render state to RGB."""
        rgba = state[..., -4:]
        rgb = rgba_to_rgb(rgba)

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgb)

    @nnx.jit
    def render_rgba(self, state):
        """Render state to RGBA."""
        rgba = state[..., -4:]

        # Clip values to valid range and convert to uint8
        return clip_and_uint8(rgba)

class DiffusingNCA(ComplexSystem): """Diffusing Neural Cellular Automata class.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize Diffusing NCA. Args: rngs: rng key. """ self.perceive = ConvPerceive( channel_size=channel_size, perception_size=num_kernels * channel_size, feature_group_count=channel_size, rngs=rngs, ) self.update = NCAUpdate( channel_size=channel_size, perception_size=num_kernels * channel_size, hidden_layer_sizes=(hidden_size,), cell_dropout_rate=cell_dropout_rate, zeros_init=True, rngs=rngs, ) # Initialize kernel with sobel filters kernel = jnp.concatenate( [identity_kernel(num_dims=2), grad_kernel(num_dims=2)], axis=-1 ) kernel = jnp.expand_dims( jnp.concatenate([kernel] * channel_size, axis=-1), axis=-2 ) self.perceive.conv.kernel[...] = kernel def _step(self, state: Array, input: Array | None = None) -> Array: perception = self.perceive(state) next_state = self.update(state, perception, input) return next_state @nnx.jit def render(self, state): """Render state to RGB.""" rgba = state[..., -4:] rgb = rgba_to_rgb(rgba) # Clip values to valid range and convert to uint8 return clip_and_uint8(rgb) @nnx.jit def render_rgba(self, state): """Render state to RGBA.""" rgba = state[..., -4:] # Clip values to valid range and convert to uint8 return clip_and_uint8(rgba)

In [7]:

cs = DiffusingNCA(rngs=rngs)

cs = DiffusingNCA(rngs=rngs)

In [8]:

params = nnx.state(cs, nnx.Param)
print("Number of params:", sum(x.size for x in jax.tree.leaves(params)))

params = nnx.state(cs, nnx.Param) print("Number of params:", sum(x.size for x in jax.tree.leaves(params)))

Number of params: 67584

Sample initial state¶

In [9]:

def add_noise(image, alpha, key):
    """Add noise to the image with a given alpha value."""
    noise = jax.random.normal(key, image.shape)
    noisy_image = (1 - alpha) * image + alpha * noise
    return jnp.clip(noisy_image, 0.0, 1.0)


def sample_state(key):
    """Sample a state with added noise."""
    state = jnp.zeros(y.shape[:2] + (channel_size,))

    alpha_key, noise_key = jax.random.split(key)

    # Add noise
    alpha = jax.random.uniform(alpha_key)
    noisy_y = add_noise(y, alpha, noise_key)

    return state.at[..., -4:].set(noisy_y)

def add_noise(image, alpha, key): """Add noise to the image with a given alpha value.""" noise = jax.random.normal(key, image.shape) noisy_image = (1 - alpha) * image + alpha * noise return jnp.clip(noisy_image, 0.0, 1.0) def sample_state(key): """Sample a state with added noise.""" state = jnp.zeros(y.shape[:2] + (channel_size,)) alpha_key, noise_key = jax.random.split(key) # Add noise alpha = jax.random.uniform(alpha_key) noisy_y = add_noise(y, alpha, noise_key) return state.at[..., -4:].set(noisy_y)

Train¶

Optimizer¶

In [10]:

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=2_000
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=lr_sched),
)

update_params = nnx.All(nnx.Param, nnx.PathContains("update"))
optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.1 * learning_rate, transition_steps=2_000 ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=lr_sched), ) update_params = nnx.All(nnx.Param, nnx.PathContains("update")) optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

Loss¶

In [11]:

def mse(state):
    """Mean Squared Error."""
    return jnp.mean(jnp.square(state[..., -4:] - y))

def mse(state): """Mean Squared Error.""" return jnp.mean(jnp.square(state[..., -4:] - y))

In [12]:

@nnx.jit
def loss_fn(cs, state):
    """Loss function."""
    state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
    state = nnx.split_rngs(splits=batch_size)(
        nnx.vmap(
            lambda cs, state: cs(state, num_steps=num_steps),
            in_axes=(state_axes, 0),
        )
    )(cs, state)

    loss = mse(state)
    return loss

@nnx.jit def loss_fn(cs, state): """Loss function.""" state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state = nnx.split_rngs(splits=batch_size)( nnx.vmap( lambda cs, state: cs(state, num_steps=num_steps), in_axes=(state_axes, 0), ) )(cs, state) loss = mse(state) return loss

Train step¶

In [13]:

@nnx.jit
def train_step(cs, optimizer, key):
    """Train step."""
    keys = jax.random.split(key, batch_size)
    current_state = jax.vmap(sample_state)(keys)

    loss, grad = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, update_params))(
        cs, current_state
    )
    optimizer.update(cs, grad)

    return loss

@nnx.jit def train_step(cs, optimizer, key): """Train step.""" keys = jax.random.split(key, batch_size) current_state = jax.vmap(sample_state)(keys) loss, grad = nnx.value_and_grad(loss_fn, argnums=nnx.DiffState(0, update_params))( cs, current_state ) optimizer.update(cs, grad) return loss

Main loop¶

In [14]:

num_train_steps = 8_192
print_interval = 128

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    loss = train_step(cs, optimizer, subkey)
    losses.append(loss)

    if i % print_interval == 0 or i == num_train_steps - 1:
        avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:])
        elapsed = time.perf_counter() - start
        print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}")

print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

num_train_steps = 8_192 print_interval = 128 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) loss = train_step(cs, optimizer, subkey) losses.append(loss) if i % print_interval == 0 or i == num_train_steps - 1: avg_loss = sum(losses[-print_interval:]) / len(losses[-print_interval:]) elapsed = time.perf_counter() - start print(f"Step {i:>4}/{num_train_steps} | {elapsed:6.1f}s | Loss {avg_loss:.3e}") print(f"✨ Trained for {num_train_steps} steps in {time.perf_counter() - start:.0f}s")

W0904 21:36:26.914117  452785 cuda_timer.cc:88] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.

Step    0/8192 |    4.0s | Loss 1.152e-01

Step  128/8192 |   13.2s | Loss 2.095e-02

Step  256/8192 |   22.3s | Loss 1.312e-02

Step  384/8192 |   31.4s | Loss 1.341e-02

Step  512/8192 |   40.5s | Loss 1.093e-02

Step  640/8192 |   49.6s | Loss 1.210e-02

Step  768/8192 |   58.7s | Loss 1.016e-02

Step  896/8192 |   67.8s | Loss 9.561e-03

Step 1024/8192 |   76.9s | Loss 9.112e-03

Step 1152/8192 |   86.0s | Loss 7.720e-03

Step 1280/8192 |   95.1s | Loss 7.250e-03

Step 1408/8192 |  104.2s | Loss 6.294e-03

Step 1536/8192 |  113.3s | Loss 6.304e-03

Step 1664/8192 |  122.4s | Loss 5.879e-03

Step 1792/8192 |  131.5s | Loss 5.394e-03

Step 1920/8192 |  140.6s | Loss 5.704e-03

Step 2048/8192 |  149.7s | Loss 5.165e-03

Step 2176/8192 |  158.8s | Loss 5.310e-03

Step 2304/8192 |  167.9s | Loss 5.292e-03

Step 2432/8192 |  177.0s | Loss 5.289e-03

Step 2560/8192 |  186.2s | Loss 5.393e-03

Step 2688/8192 |  195.3s | Loss 5.224e-03

Step 2816/8192 |  204.4s | Loss 5.017e-03

Step 2944/8192 |  213.5s | Loss 5.056e-03

Step 3072/8192 |  222.6s | Loss 4.923e-03

Step 3200/8192 |  231.7s | Loss 4.805e-03

Step 3328/8192 |  240.8s | Loss 4.889e-03

Step 3456/8192 |  249.9s | Loss 5.022e-03

Step 3584/8192 |  259.0s | Loss 4.935e-03

Step 3712/8192 |  268.1s | Loss 4.987e-03

Step 3840/8192 |  277.2s | Loss 5.029e-03

Step 3968/8192 |  286.3s | Loss 4.970e-03

Step 4096/8192 |  295.4s | Loss 4.704e-03

Step 4224/8192 |  304.5s | Loss 4.679e-03

Step 4352/8192 |  313.6s | Loss 4.810e-03

Step 4480/8192 |  322.7s | Loss 4.538e-03

Step 4608/8192 |  331.8s | Loss 4.699e-03

Step 4736/8192 |  340.9s | Loss 4.375e-03

Step 4864/8192 |  350.0s | Loss 4.374e-03

Step 4992/8192 |  359.1s | Loss 4.287e-03

Step 5120/8192 |  368.2s | Loss 4.355e-03

Step 5248/8192 |  377.3s | Loss 4.429e-03

Step 5376/8192 |  386.4s | Loss 4.216e-03

Step 5504/8192 |  395.5s | Loss 4.145e-03

Step 5632/8192 |  404.6s | Loss 4.055e-03

Step 5760/8192 |  413.7s | Loss 4.212e-03

Step 5888/8192 |  422.8s | Loss 4.092e-03

Step 6016/8192 |  431.9s | Loss 3.849e-03

Step 6144/8192 |  441.0s | Loss 3.782e-03

Step 6272/8192 |  450.1s | Loss 3.712e-03

Step 6400/8192 |  459.2s | Loss 3.859e-03

Step 6528/8192 |  468.3s | Loss 3.701e-03

Step 6656/8192 |  477.4s | Loss 3.686e-03

Step 6784/8192 |  486.5s | Loss 3.840e-03

Step 6912/8192 |  495.6s | Loss 3.595e-03

Step 7040/8192 |  504.7s | Loss 3.476e-03

Step 7168/8192 |  513.7s | Loss 3.610e-03

Step 7296/8192 |  522.8s | Loss 3.453e-03

Step 7424/8192 |  531.9s | Loss 3.764e-03

Step 7552/8192 |  541.0s | Loss 3.303e-03

Step 7680/8192 |  550.1s | Loss 3.468e-03

Step 7808/8192 |  559.2s | Loss 3.232e-03

Step 7936/8192 |  568.2s | Loss 3.385e-03

Step 8064/8192 |  577.3s | Loss 3.531e-03

Step 8191/8192 |  586.3s | Loss 3.540e-03
✨ Trained for 8192 steps in 586s

Run¶

In [15]:

num_examples = 8

key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, num_examples)
state_init = jax.vmap(sample_state)(keys)

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
state_final, states = nnx.split_rngs(splits=num_examples)(
    nnx.vmap(
        lambda cs, state: cs(state, num_steps=num_steps, return_states=True),
        in_axes=(state_axes, 0),
    )
)(cs, state_init)

num_examples = 8 key, subkey = jax.random.split(key) keys = jax.random.split(subkey, num_examples) state_init = jax.vmap(sample_state)(keys) state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) state_final, states = nnx.split_rngs(splits=num_examples)( nnx.vmap( lambda cs, state: cs(state, num_steps=num_steps, return_states=True), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [16]:

frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, states)

frames_final = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, state_final)

frames_final_rgba = nnx.vmap(
    lambda cs, state: cs.render_rgba(state),
    in_axes=(None, 0),
)(cs, state_final)

mediapy.show_images(frames_final.repeat(2, axis=-3).repeat(2, axis=-2))
mediapy.show_videos(frames.repeat(2, axis=-3).repeat(2, axis=-2))

frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) frames_final = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, state_final) frames_final_rgba = nnx.vmap( lambda cs, state: cs.render_rgba(state), in_axes=(None, 0), )(cs, state_final) mediapy.show_images(frames_final.repeat(2, axis=-3).repeat(2, axis=-2)) mediapy.show_videos(frames.repeat(2, axis=-3).repeat(2, axis=-2))