Variational Autoencoder ¶

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 cax.nn.vae import VAE, vae_loss

import time import jax import jax.numpy as jnp import mediapy import optax import torchvision from flax import nnx from cax.nn.vae import VAE, vae_loss

Configuration¶

In [4]:

seed = 0

spatial_dims = (28, 28)
features = (1, 32, 32)
latent_size = 4

batch_size = 32
learning_rate = 1e-2

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

seed = 0 spatial_dims = (28, 28) features = (1, 32, 32) latent_size = 4 batch_size = 32 learning_rate = 1e-2 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([y.resize(spatial_dims) for y, _ in ds_train])[..., None] / 255
x_test = jnp.array([y.resize(spatial_dims) for y, _ 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([y.resize(spatial_dims) for y, _ in ds_train])[..., None] / 255 x_test = jnp.array([y.resize(spatial_dims) for y, _ in ds_test])[..., None] / 255 # Visualize mediapy.show_images(x_train[:8].repeat(4, axis=-3).repeat(4, axis=-2))

Instantiate system¶

In [6]:

vae = VAE(
    spatial_dims=spatial_dims,
    features=features,
    latent_size=latent_size,
    rngs=rngs,
)

vae = VAE( spatial_dims=spatial_dims, features=features, latent_size=latent_size, rngs=rngs, )

In [7]:

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

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

Number of params: 2499689

Train¶

Optimizer¶

In [8]:

lr_sched = optax.linear_schedule(
    init_value=learning_rate, end_value=0.01 * learning_rate, transition_steps=8_192
)

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

optimizer = nnx.Optimizer(vae, optimizer, wrt=nnx.Param)

lr_sched = optax.linear_schedule( init_value=learning_rate, end_value=0.01 * learning_rate, transition_steps=8_192 ) optimizer = optax.chain( optax.clip_by_global_norm(1.0), optax.adam(learning_rate=lr_sched), ) optimizer = nnx.Optimizer(vae, optimizer, wrt=nnx.Param)

Loss¶

In [9]:

@nnx.jit
def loss_fn(vae, image):
    """Loss function."""
    image_recon, mean, logvar = vae(image)
    return vae_loss(image_recon, image, mean, logvar)

@nnx.jit def loss_fn(vae, image): """Loss function.""" image_recon, mean, logvar = vae(image) return vae_loss(image_recon, image, mean, logvar)

Train step¶

In [10]:

@nnx.jit
def train_step(vae, optimizer, key):
    """Train step."""
    image_idx = jax.random.choice(key, x_train.shape[0], shape=(batch_size,))
    image = x_train[image_idx]

    loss, grad = nnx.value_and_grad(loss_fn)(vae, image)
    optimizer.update(vae, grad)

    return loss

@nnx.jit def train_step(vae, optimizer, key): """Train step.""" image_idx = jax.random.choice(key, x_train.shape[0], shape=(batch_size,)) image = x_train[image_idx] loss, grad = nnx.value_and_grad(loss_fn)(vae, image) optimizer.update(vae, grad) return loss

Main loop¶

In [11]:

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(vae, 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(vae, 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")

Step    0/8192 |    5.9s | Loss 1.759e+04

Step  128/8192 |    6.3s | Loss 8.535e+03

Step  256/8192 |    6.7s | Loss 5.065e+03

Step  384/8192 |    7.1s | Loss 4.826e+03

Step  512/8192 |    7.4s | Loss 4.660e+03

Step  640/8192 |    7.8s | Loss 4.604e+03

Step  768/8192 |    8.2s | Loss 4.546e+03

Step  896/8192 |    8.5s | Loss 4.512e+03

Step 1024/8192 |    8.9s | Loss 4.519e+03

Step 1152/8192 |    9.3s | Loss 4.479e+03

Step 1280/8192 |    9.7s | Loss 4.441e+03

Step 1408/8192 |   10.0s | Loss 4.459e+03

Step 1536/8192 |   10.4s | Loss 4.419e+03

Step 1664/8192 |   10.8s | Loss 4.391e+03

Step 1792/8192 |   11.1s | Loss 4.385e+03

Step 1920/8192 |   11.5s | Loss 4.396e+03

Step 2048/8192 |   11.8s | Loss 4.354e+03

Step 2176/8192 |   12.2s | Loss 4.367e+03

Step 2304/8192 |   12.6s | Loss 4.364e+03

Step 2432/8192 |   12.9s | Loss 4.354e+03

Step 2560/8192 |   13.3s | Loss 4.325e+03

Step 2688/8192 |   13.7s | Loss 4.313e+03

Step 2816/8192 |   14.1s | Loss 4.296e+03

Step 2944/8192 |   14.4s | Loss 4.321e+03

Step 3072/8192 |   14.8s | Loss 4.289e+03

Step 3200/8192 |   15.2s | Loss 4.292e+03

Step 3328/8192 |   15.5s | Loss 4.291e+03

Step 3456/8192 |   15.9s | Loss 4.257e+03

Step 3584/8192 |   16.3s | Loss 4.254e+03

Step 3712/8192 |   16.6s | Loss 4.248e+03

Step 3840/8192 |   17.0s | Loss 4.249e+03

Step 3968/8192 |   17.3s | Loss 4.277e+03

Step 4096/8192 |   17.7s | Loss 4.244e+03

Step 4224/8192 |   18.1s | Loss 4.240e+03

Step 4352/8192 |   18.4s | Loss 4.251e+03

Step 4480/8192 |   18.8s | Loss 4.205e+03

Step 4608/8192 |   19.2s | Loss 4.198e+03

Step 4736/8192 |   19.5s | Loss 4.237e+03

Step 4864/8192 |   19.9s | Loss 4.203e+03

Step 4992/8192 |   20.3s | Loss 4.180e+03

Step 5120/8192 |   20.6s | Loss 4.245e+03

Step 5248/8192 |   21.0s | Loss 4.197e+03

Step 5376/8192 |   21.4s | Loss 4.201e+03

Step 5504/8192 |   21.7s | Loss 4.198e+03

Step 5632/8192 |   22.1s | Loss 4.166e+03

Step 5760/8192 |   22.5s | Loss 4.148e+03

Step 5888/8192 |   22.8s | Loss 4.192e+03

Step 6016/8192 |   23.2s | Loss 4.175e+03

Step 6144/8192 |   23.6s | Loss 4.164e+03

Step 6272/8192 |   24.0s | Loss 4.174e+03

Step 6400/8192 |   24.3s | Loss 4.149e+03

Step 6528/8192 |   24.7s | Loss 4.144e+03

Step 6656/8192 |   25.1s | Loss 4.151e+03

Step 6784/8192 |   25.4s | Loss 4.132e+03

Step 6912/8192 |   25.8s | Loss 4.132e+03

Step 7040/8192 |   26.2s | Loss 4.148e+03

Step 7168/8192 |   26.5s | Loss 4.174e+03

Step 7296/8192 |   26.9s | Loss 4.135e+03

Step 7424/8192 |   27.3s | Loss 4.106e+03

Step 7552/8192 |   27.6s | Loss 4.097e+03

Step 7680/8192 |   28.0s | Loss 4.075e+03

Step 7808/8192 |   28.4s | Loss 4.128e+03

Step 7936/8192 |   28.7s | Loss 4.114e+03

Step 8064/8192 |   29.1s | Loss 4.134e+03

Step 8191/8192 |   29.5s | Loss 4.114e+03
✨ Trained for 8192 steps in 29s

Visualize¶

In [12]:

num_examples = 8

key, subkey = jax.random.split(key)
z = jax.random.normal(subkey, shape=(num_examples, latent_size))
x = vae.generate(z)

mediapy.show_images(x.repeat(4, axis=-3).repeat(4, axis=-2))

num_examples = 8 key, subkey = jax.random.split(key) z = jax.random.normal(subkey, shape=(num_examples, latent_size)) x = vae.generate(z) mediapy.show_images(x.repeat(4, axis=-3).repeat(4, axis=-2))

In [13]:

num_examples = 8

key, subkey = jax.random.split(key)
x_idx = jax.random.choice(subkey, x_test.shape[0], shape=(num_examples,))
x = x_test[x_idx]

state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None})
x_recon, _, _ = nnx.split_rngs(splits=num_examples)(
    nnx.vmap(
        lambda vae, x: vae(x),
        in_axes=(state_axes, 0),
    )
)(vae, x)

mediapy.show_images(x.repeat(4, axis=-3).repeat(4, axis=-2))
mediapy.show_images(jax.nn.sigmoid(x_recon).repeat(4, axis=-3).repeat(4, axis=-2))

num_examples = 8 key, subkey = jax.random.split(key) x_idx = jax.random.choice(subkey, x_test.shape[0], shape=(num_examples,)) x = x_test[x_idx] state_axes = nnx.StateAxes({nnx.RngState: 0, ...: None}) x_recon, _, _ = nnx.split_rngs(splits=num_examples)( nnx.vmap( lambda vae, x: vae(x), in_axes=(state_axes, 0), ) )(vae, x) mediapy.show_images(x.repeat(4, axis=-3).repeat(4, axis=-2)) mediapy.show_images(jax.nn.sigmoid(x_recon).repeat(4, axis=-3).repeat(4, axis=-2))