Self-autoencoding MNIST Digits ¶

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
import torchvision
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.nca_update import NCAUpdate
from cax.nn.pool import Pool
from cax.utils import clip_and_uint8

import time import jax import jax.numpy as jnp import mediapy import optax import torchvision 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.nca_update import NCAUpdate from cax.nn.pool import Pool from cax.utils import clip_and_uint8

Configuration¶

In [4]:

seed = 0

channel_size = 16
spatial_dims = (28, 28, 42)
num_kernels = 4
hidden_size = 256
cell_dropout_rate = 0.5

num_steps = 96
pool_size = 1_024
batch_size = 8
learning_rate = 1e-3

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

seed = 0 channel_size = 16 spatial_dims = (28, 28, 42) num_kernels = 4 hidden_size = 256 cell_dropout_rate = 0.5 num_steps = 96 pool_size = 1_024 batch_size = 8 learning_rate = 1e-3 key = jax.random.key(seed) rngs = nnx.Rngs(seed)

Dataset¶

In [5]:

# Load MNIST dataset
ds_train = torchvision.datasets.MNIST(root="./data", train=True, download=True)
ds_test = torchvision.datasets.MNIST(root="./data", train=False, download=True)

# Convert to jax.Array
x_train = jnp.array([x.resize(spatial_dims[:2]) for x, _ in ds_train])[..., None] / 255
x_test = jnp.array([x.resize(spatial_dims[:2]) for x, _ in ds_test])[..., None] / 255

# Visualize
mediapy.show_images(x_train[:8].repeat(4, axis=-3).repeat(4, axis=-2))

# Load MNIST dataset ds_train = torchvision.datasets.MNIST(root="./data", train=True, download=True) ds_test = torchvision.datasets.MNIST(root="./data", train=False, download=True) # Convert to jax.Array x_train = jnp.array([x.resize(spatial_dims[:2]) for x, _ in ds_train])[..., None] / 255 x_test = jnp.array([x.resize(spatial_dims[:2]) for x, _ in ds_test])[..., None] / 255 # Visualize mediapy.show_images(x_train[:8].repeat(4, axis=-3).repeat(4, axis=-2))

Instantiate system¶

In [6]:

class SelfAutoencodingNCA(ComplexSystem):
    """Self-Autoencoding Neural Cellular Automata class."""

    def __init__(self, *, rngs: nnx.Rngs):
        """Initialize Self-Autoencoding NCA."""
        self.perceive = ConvPerceive(
            channel_size=channel_size,
            perception_size=num_kernels * channel_size,
            kernel_size=(3, 3, 3),
            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,
            kernel_size=(3, 3, 3),
            zeros_init=True,
            rngs=rngs,
        )

        # Initialize kernel with sobel filters
        kernel = jnp.concatenate(
            [identity_kernel(num_dims=3), grad_kernel(num_dims=3)], 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:
        """Perform a single step."""
        # Extract x
        x = state[..., 0, -1:]

        # Step
        perception = self.perceive(state)
        next_state = self.update(state, perception, input)

        # Mask
        mid = tuple(size // 2 for size in spatial_dims)
        center = next_state[..., *mid, :]
        next_state = next_state.at[..., mid[-1], :].set(0.0)  # Mask
        next_state = next_state.at[..., *mid, :].set(center)  # Except center cell

        # Override
        next_state = next_state.at[..., 0, -1:].set(x)

        return next_state

    @nnx.jit
    def render(self, state):
        """Render state to RGB."""
        gray = state[..., -1, -1:]
        rgb = jnp.repeat(gray, 3, axis=-1)

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

class SelfAutoencodingNCA(ComplexSystem): """Self-Autoencoding Neural Cellular Automata class.""" def __init__(self, *, rngs: nnx.Rngs): """Initialize Self-Autoencoding NCA.""" self.perceive = ConvPerceive( channel_size=channel_size, perception_size=num_kernels * channel_size, kernel_size=(3, 3, 3), 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, kernel_size=(3, 3, 3), zeros_init=True, rngs=rngs, ) # Initialize kernel with sobel filters kernel = jnp.concatenate( [identity_kernel(num_dims=3), grad_kernel(num_dims=3)], 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: """Perform a single step.""" # Extract x x = state[..., 0, -1:] # Step perception = self.perceive(state) next_state = self.update(state, perception, input) # Mask mid = tuple(size // 2 for size in spatial_dims) center = next_state[..., *mid, :] next_state = next_state.at[..., mid[-1], :].set(0.0) # Mask next_state = next_state.at[..., *mid, :].set(center) # Except center cell # Override next_state = next_state.at[..., 0, -1:].set(x) return next_state @nnx.jit def render(self, state): """Render state to RGB.""" gray = state[..., -1, -1:] rgb = jnp.repeat(gray, 3, axis=-1) # Clip values to valid range and convert to uint8 return clip_and_uint8(rgb)

In [7]:

cs = SelfAutoencodingNCA(rngs=rngs)

cs = SelfAutoencodingNCA(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: 22480

Sample initial state¶

In [9]:

def sample_state(key):
    """Sample a state with a random image."""
    # Init state
    state = jnp.zeros(spatial_dims + (channel_size,))

    # Sample random image
    x_idx = jax.random.choice(key, x_train.shape[0])
    x = x_train[x_idx]

    # Set image in state
    state = state.at[..., 0, -1:].set(x)
    return state, x_idx

def sample_state(key): """Sample a state with a random image.""" # Init state state = jnp.zeros(spatial_dims + (channel_size,)) # Sample random image x_idx = jax.random.choice(key, x_train.shape[0]) x = x_train[x_idx] # Set image in state state = state.at[..., 0, -1:].set(x) return state, x_idx

Train¶

Pool¶

In [10]:

key, subkey = jax.random.split(key)

keys = jax.random.split(subkey, pool_size)
state, x_idx = jax.vmap(sample_state)(keys)

pool = Pool.create({"state": state, "x_idx": x_idx})

key, subkey = jax.random.split(key) keys = jax.random.split(subkey, pool_size) state, x_idx = jax.vmap(sample_state)(keys) pool = Pool.create({"state": state, "x_idx": x_idx})

Optimizer¶

In [11]:

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.5 * 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)
optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.5 * 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) optimizer = nnx.Optimizer(cs, optimizer, wrt=update_params)

Loss¶

In [12]:

def mse(state, x):
    """Mean Squared Error."""
    return jnp.mean(jnp.square(state[..., :, -1:] - x[..., None, :]))

def mse(state, x): """Mean Squared Error.""" return jnp.mean(jnp.square(state[..., :, -1:] - x[..., None, :]))

In [13]:

@nnx.jit
def loss_fn(cs, state, x):
    """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, return_states=True),
            in_axes=(state_axes, 0),
        )
    )(cs, state)

    # Sample a random step
    idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps)
    state = state[jnp.arange(batch_size), idx]

    loss = mse(state, x)
    return loss, state

@nnx.jit def loss_fn(cs, state, x): """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, return_states=True), in_axes=(state_axes, 0), ) )(cs, state) # Sample a random step idx = jax.random.randint(key, (batch_size,), num_steps // 2, num_steps) state = state[jnp.arange(batch_size), idx] loss = mse(state, x) return loss, state

Train step¶

In [14]:

@nnx.jit
def train_step(cs, optimizer, pool, key):
    """Train step."""
    sample_key, sample_state_key = jax.random.split(key)

    # Sample from pool
    pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False)
    current_state = batch["state"]
    current_x_idx = batch["x_idx"]
    current_x = x_train[current_x_idx]

    # Sort by descending loss
    sort_idx = jnp.argsort(jax.vmap(mse)(current_state, current_x), descending=True)
    pool_idx = pool_idx[sort_idx]
    current_state = current_state[sort_idx]
    current_x_idx = current_x_idx[sort_idx]

    # Sample a new state to replace the worst
    new_state, new_x_idx = sample_state(sample_state_key)
    current_state = current_state.at[0].set(new_state)
    current_x_idx = current_x_idx.at[0].set(new_x_idx)
    current_x = x_train[current_x_idx]

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

    pool = pool.update(pool_idx, {"state": current_state, "x_idx": current_x_idx})
    return loss, pool

@nnx.jit def train_step(cs, optimizer, pool, key): """Train step.""" sample_key, sample_state_key = jax.random.split(key) # Sample from pool pool_idx, batch = pool.sample(sample_key, batch_size=batch_size, replace=False) current_state = batch["state"] current_x_idx = batch["x_idx"] current_x = x_train[current_x_idx] # Sort by descending loss sort_idx = jnp.argsort(jax.vmap(mse)(current_state, current_x), descending=True) pool_idx = pool_idx[sort_idx] current_state = current_state[sort_idx] current_x_idx = current_x_idx[sort_idx] # Sample a new state to replace the worst new_state, new_x_idx = sample_state(sample_state_key) current_state = current_state.at[0].set(new_state) current_x_idx = current_x_idx.at[0].set(new_x_idx) current_x = x_train[current_x_idx] (loss, current_state), grad = nnx.value_and_grad( loss_fn, has_aux=True, argnums=nnx.DiffState(0, update_params) )(cs, current_state, current_x) optimizer.update(cs, grad) pool = pool.update(pool_idx, {"state": current_state, "x_idx": current_x_idx}) return loss, pool

Main loop¶

In [15]:

num_train_steps = 2 * 8_192
print_interval = 256

losses = []
start = time.perf_counter()
for i in range(num_train_steps):
    key, subkey = jax.random.split(key)
    loss, pool = train_step(cs, optimizer, pool, 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:>5}/{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 = 2 * 8_192 print_interval = 256 losses = [] start = time.perf_counter() for i in range(num_train_steps): key, subkey = jax.random.split(key) loss, pool = train_step(cs, optimizer, pool, 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:>5}/{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:47:59.574688  495457 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.
W0904 21:47:59.748837  495457 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/16384 |    7.4s | Loss 1.028e-01

Step   256/16384 |  112.3s | Loss 5.497e-02

Step   512/16384 |  216.9s | Loss 3.901e-02

Step   768/16384 |  321.6s | Loss 3.382e-02

Step  1024/16384 |  426.3s | Loss 3.142e-02

Step  1280/16384 |  530.9s | Loss 2.872e-02

Step  1536/16384 |  635.6s | Loss 2.614e-02

Step  1792/16384 |  740.3s | Loss 2.411e-02

Step  2048/16384 |  844.9s | Loss 2.240e-02

Step  2304/16384 |  949.6s | Loss 2.160e-02

Step  2560/16384 | 1054.3s | Loss 2.077e-02

Step  2816/16384 | 1159.3s | Loss 1.943e-02

Step  3072/16384 | 1264.2s | Loss 1.872e-02

Step  3328/16384 | 1368.8s | Loss 1.767e-02

Step  3584/16384 | 1473.5s | Loss 1.691e-02

Step  3840/16384 | 1578.2s | Loss 1.640e-02

Step  4096/16384 | 1682.8s | Loss 1.611e-02

Step  4352/16384 | 1787.4s | Loss 1.504e-02

Step  4608/16384 | 1892.1s | Loss 1.446e-02

Step  4864/16384 | 1996.7s | Loss 1.432e-02

Step  5120/16384 | 2101.3s | Loss 1.423e-02

Step  5376/16384 | 2205.9s | Loss 1.339e-02

Step  5632/16384 | 2310.6s | Loss 1.347e-02

Step  5888/16384 | 2415.6s | Loss 1.285e-02

Step  6144/16384 | 2520.5s | Loss 1.249e-02

Step  6400/16384 | 2625.5s | Loss 1.216e-02

Step  6656/16384 | 2730.5s | Loss 1.177e-02

Step  6912/16384 | 2835.4s | Loss 1.198e-02

Step  7168/16384 | 2940.4s | Loss 1.156e-02

Step  7424/16384 | 3045.4s | Loss 1.138e-02

Step  7680/16384 | 3150.3s | Loss 1.124e-02

Step  7936/16384 | 3255.3s | Loss 1.128e-02

Step  8192/16384 | 3360.3s | Loss 1.156e-02

Step  8448/16384 | 3465.2s | Loss 1.109e-02

Step  8704/16384 | 3570.2s | Loss 1.084e-02

Step  8960/16384 | 3675.2s | Loss 1.115e-02

Step  9216/16384 | 3780.1s | Loss 1.073e-02

Step  9472/16384 | 3885.1s | Loss 1.094e-02

Step  9728/16384 | 3990.1s | Loss 1.071e-02

Step  9984/16384 | 4095.0s | Loss 1.014e-02

Step 10240/16384 | 4200.0s | Loss 1.015e-02

Step 10496/16384 | 4305.0s | Loss 1.060e-02

Step 10752/16384 | 4409.9s | Loss 1.025e-02

Step 11008/16384 | 4514.9s | Loss 9.959e-03

Step 11264/16384 | 4619.9s | Loss 1.054e-02

Step 11520/16384 | 4724.8s | Loss 1.008e-02

Step 11776/16384 | 4829.8s | Loss 9.844e-03

Step 12032/16384 | 4934.8s | Loss 9.414e-03

Step 12288/16384 | 5039.7s | Loss 1.034e-02

Step 12544/16384 | 5144.7s | Loss 1.002e-02

Step 12800/16384 | 5249.6s | Loss 9.586e-03

Step 13056/16384 | 5354.6s | Loss 9.531e-03

Step 13312/16384 | 5459.6s | Loss 9.318e-03

Step 13568/16384 | 5564.5s | Loss 9.310e-03

Step 13824/16384 | 5669.5s | Loss 9.165e-03

Step 14080/16384 | 5774.5s | Loss 9.008e-03

Step 14336/16384 | 5879.4s | Loss 9.297e-03

Step 14592/16384 | 5984.4s | Loss 9.159e-03

Step 14848/16384 | 6089.4s | Loss 9.139e-03

Step 15104/16384 | 6194.3s | Loss 9.283e-03

Step 15360/16384 | 6299.3s | Loss 9.131e-03

Step 15616/16384 | 6404.2s | Loss 8.919e-03

Step 15872/16384 | 6509.2s | Loss 8.894e-03

Step 16128/16384 | 6614.1s | Loss 9.235e-03

Step 16383/16384 | 6718.7s | Loss 9.431e-03
✨ Trained for 16384 steps in 6719s

Run¶

In [16]:

num_examples = 8

key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, num_examples)
state_init, x_idx = 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=2 * 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, x_idx = 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=2 * num_steps, return_states=True), in_axes=(state_axes, 0), ) )(cs, state_init)

Visualize¶

In [17]:

states = jnp.concatenate([state_init[:, None], states], axis=1)
frames = nnx.vmap(
    lambda cs, state: cs.render(state),
    in_axes=(None, 0),
)(cs, states)

mediapy.show_images(x_train[x_idx].repeat(4, axis=-3).repeat(4, axis=-2))
mediapy.show_videos(frames.repeat(4, axis=-3).repeat(4, axis=-2))

states = jnp.concatenate([state_init[:, None], states], axis=1) frames = nnx.vmap( lambda cs, state: cs.render(state), in_axes=(None, 0), )(cs, states) mediapy.show_images(x_train[x_idx].repeat(4, axis=-3).repeat(4, axis=-2)) mediapy.show_videos(frames.repeat(4, axis=-3).repeat(4, axis=-2))