{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "# Training a neural network with DALI and JAX\n", "\n", "This simple example shows how to train a neural network implemented in JAX with DALI pipelines. It builds on MNIST training example from JAX codebase that can be found [here](https://github.com/google/jax/blob/jax-v0.4.13/examples/mnist_classifier_fromscratch.py).\n", "\n", "We will use MNIST in Caffe2 format from [DALI_extra](https://github.com/NVIDIA/DALI_extra)." ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "training_data_path = os.path.join(\n", " os.environ[\"DALI_EXTRA_PATH\"], \"db/MNIST/training/\"\n", ")\n", "validation_data_path = os.path.join(\n", " os.environ[\"DALI_EXTRA_PATH\"], \"db/MNIST/testing/\"\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "First step is to create a definition function that will later be used to create instances of DALI iterators. It defines all steps of the preprocessing. \n", "\n", "In this simple example we have `fn.readers.caffe2` for reading data in Caffe2 format, `fn.decoders.image` for image decoding, `fn.crop_mirror_normalize` used to normalize the images and `fn.reshape` to adjust the shape of the output tensors. We also move the labels from the CPU to the GPU memory with `labels.gpu()`. Our model expects labels to be in one-hot encoding, so we use `fn.one_hot` to convert them.\n", "\n", "This example focuses on how to use DALI to train a model defined in JAX. For more information on DALI and JAX integration look into [Getting started with JAX and DALI](jax-getting_started.ipynb) and [pipeline documentation](../../../pipeline.rst)" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "from nvidia.dali.plugin.jax import data_iterator\n", "import nvidia.dali.fn as fn\n", "import nvidia.dali.types as types\n", "\n", "batch_size = 200\n", "image_size = 28\n", "num_classes = 10\n", "\n", "\n", "@data_iterator(output_map=[\"images\", \"labels\"], reader_name=\"caffe2_reader\")\n", "def mnist_iterator(data_path, random_shuffle):\n", " jpegs, labels = fn.readers.caffe2(\n", " path=data_path, random_shuffle=random_shuffle, name=\"caffe2_reader\"\n", " )\n", " images = fn.decoders.image(jpegs, device=\"mixed\", output_type=types.GRAY)\n", " images = fn.crop_mirror_normalize(\n", " images, dtype=types.FLOAT, std=[255.0], output_layout=\"CHW\"\n", " )\n", " images = fn.reshape(images, shape=[image_size * image_size])\n", "\n", " labels = labels.gpu()\n", "\n", " if random_shuffle:\n", " labels = fn.one_hot(labels, num_classes=num_classes)\n", "\n", " return images, labels" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Next, we use the function to create DALI iterators for training and validation." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Creating iterators\n", "\n", "\n" ] } ], "source": [ "print(\"Creating iterators\")\n", "\n", "training_iterator = mnist_iterator(\n", " data_path=training_data_path, random_shuffle=True, batch_size=batch_size\n", ")\n", "\n", "validation_iterator = mnist_iterator(\n", " data_path=validation_data_path, random_shuffle=False, batch_size=batch_size\n", ")\n", "\n", "print(training_iterator)\n", "print(validation_iterator)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "With the setup above, DALI iterators are ready for the training. \n", "\n", "Finally, we import training utilities implemented in JAX. `init_model` will create the model instance and initialize its parameters. In this simple example it is a MLP model with two hidden layers. `update` performs one iteration of the training. `accuracy` is a helper function to run validation after each epoch on the test set and get current accuracy of the model." ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "from model import init_model, update, accuracy" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "`jax.jit` traces, compiles, and caches functions lazily on first invocation for a given input signature. During this process, XLA may capture CUDA graphs, which forbids some CUDA calls that DALI's background thread uses internally. Since subsequent calls to the JAX function with inputs of the same shape and dtype don't trigger compilation again, we can work around this by warming up with dummy inputs before starting any DALI workload:" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "import jax.numpy as jnp\n", "\n", "model = init_model()\n", "dummy_images = jnp.empty(\n", " (batch_size, image_size * image_size), dtype=jnp.float32\n", ")\n", "dummy_labels = jnp.empty((batch_size, num_classes), dtype=jnp.float32)\n", "_ = update(model, {\"images\": dummy_images, \"labels\": dummy_labels})" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", "\n", " Warning
\n", " \n", " If you skip this step, CUDA graph capture will happen on the first call to `update` and may overlap with DALI's execution, causing CUDA errors in JAX.\n", " \n", " Alternatively, you can disable XLA command buffers entirely by setting `XLA_FLAGS=\"--xla_gpu_enable_command_buffer=\"`, at the cost of some performance.\n", " \n", "
" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "At this point, everything is ready to run the training." ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Starting training\n", "Epoch 0 sec\n", "Test set accuracy 0.674500048160553\n", "Epoch 1 sec\n", "Test set accuracy 0.7854000329971313\n", "Epoch 2 sec\n", "Test set accuracy 0.8252000212669373\n", "Epoch 3 sec\n", "Test set accuracy 0.847100019454956\n", "Epoch 4 sec\n", "Test set accuracy 0.8618000149726868\n" ] } ], "source": [ "print(\"Starting training\")\n", "\n", "num_epochs = 5\n", "\n", "for epoch in range(num_epochs):\n", " for batch in training_iterator:\n", " model = update(model, batch)\n", "\n", " test_acc = accuracy(model, validation_iterator)\n", " print(f\"Epoch {epoch} sec\")\n", " print(f\"Test set accuracy {test_acc}\")" ] } ], "metadata": { "celltoolbar": "Raw Cell Format", "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.20" } }, "nbformat": 4, "nbformat_minor": 4 }