JAX AI Stack

The UNETR decoder relies on two specific types of upsampling modules to bridge the gap between transformer sequence representations and spatial feature maps:

1. UnetrPrUpBlock (Projection Upsampling): Used to create skip connections from the intermediate feature maps provided by the ViT encoder. It uses nnx.ConvTranspose to project and upsample the transformer's sequence representations into a spatial format.
2. UnetrUpBlock: The standard upsampling module used in the decoder path. It performs a transposed convolution on the current feature map and then concatenates the result with a skip connection from the encoder before passing it through a UnetResBlock.

UnetrUpBlock Signature: __call__(self, x: jax.Array, skip: jax.Array) -> jax.Array

• x: The upsampled feature map from the previous decoder layer.
• skip: The feature map from the encoder skip connection.

To save the best models during training (e.g., based on validation Mean IoU), use orbax.checkpoint.

When saving an nnx.Module, you must extract its state using nnx.state(model). Note that for compatibility with certain layers like Dropout, you may need to map jax._src.prng.PRNGKeyArray to a serializable format using jax.random.key_data(x) before passing the state to mngr.save().

import orbax.checkpoint as ocp

# Setup manager
path = ocp.test_utils.erase_and_create_empty("/tmp/output-oxford-model/")
options = ocp.CheckpointManagerOptions(max_to_keep=2)
mngr = ocp.CheckpointManager(path, options=options)

def save_model(epoch):
    state = nnx.state(model)
    # Convert PRNGKeyArray to serializable format
    def get_key_data(x):
        if isinstance(x, jax._src.prng.PRNGKeyArray):
            if isinstance(x.dtype, jax._src.prng.KeyTy):
                return jax.random.key_data(x)
        return x
    
    serializable_state = jax.tree.map(get_key_data, state)
    mngr.save(epoch, args=ocp.args.StandardSave(serializable_state))
    mngr.wait_until_finished()

To prepare the Food 101 dataset for training a Vision Transformer, use datasets.load_dataset to download the data. You can reduce training time by selecting a subset of classes. For efficient data loading, use grain alongside TorchVision for image transformations.

Key steps include:

1. Loading the dataset from HuggingFace.
2. Creating a label mapping to normalize class indices.
3. Applying image transformations (resizing, normalization, and augmentation) using torchvision.transforms.v2.
4. Using grain.DataLoader with a grain.IndexSampler for efficient, parallelized data loading.

from datasets import load_dataset
import grain.python as grain
from torchvision.transforms import v2 as T

# 1. Load subset
train_dataset = load_dataset("food101", split="train[:15000]")

# 2. Define transforms
tv_train_transforms = T.Compose([
    T.RandomResizedCrop((224, 224), scale=(0.7, 1.0)),
    T.RandomHorizontalFlip(),
    T.Lambda(lambda x: np.asarray(x.convert("RGB"))),
    T.Lambda(normalize),
])

# 3. Setup Grain Loader
train_sampler = grain.IndexSampler(
    len(train_dataset),
    shuffle=True,
    seed=12,
    shard_options=grain.NoSharding(),
    num_epochs=1,
)

train_loader = grain.DataLoader(
    data_source=train_dataset,
    sampler=train_sampler,
    worker_count=4,
    operations=[grain.Batch(32, drop_remainder=True)],
)

Grain is a data-loading library that can be used with JAX. To use it, install it via pip, define a custom Dataset class (e.g., wrapping torchvision.datasets.MNIST), and initialize a pygrain.DataLoader using a pygrain.SequentialSampler for sequential batch sampling.

!pip install grain

import numpy as np
import grain.python as pygrain
from torchvision.datasets import MNIST

class Dataset:
    def __init__(self, data_dir, train=True):
        self.data_dir = data_dir
        self.train = train
        self.load_data()

    def load_data(self):
        self.dataset = MNIST(self.data_dir, download=True, train=self.train)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        img, label = self.dataset[index]
        return np.ravel(np.array(img, dtype=np.float32)), label

# Initialize
mnist_dataset = Dataset(data_dir)

# Setup DataLoader
sampler = pygrain.SequentialSampler(
    num_records=len(mnist_dataset),
    shard_options=pygrain.NoSharding())

def pygrain_training_generator():
    return pygrain.DataLoader(
        data_source=mnist_dataset,
        sampler=sampler,
        operations=[pygrain.Batch(batch_size)],
    )