[flax:examples:mnist] Update mnist example to use NNX.

PiperOrigin-RevId: 815862166

Author: danielsuo

docs_nnx/mnist_tutorial.ipynb

@@ -26,15 +26,15 @@ Let’s get started!
 
 If `flax` is not installed in your Python environment, use `pip` to install the package from PyPI (below, just uncomment the code in the cell if you are working from Google Colab/Jupyter Notebook):
 
-```{code-cell}
+```{code-cell} ipython3
 # !pip install flax
 ```
 
 ## 2. Load the MNIST dataset
 
 First, you need to load the MNIST dataset and then prepare the training and testing sets via Tensorflow Datasets (TFDS). You normalize image values, shuffle the data and divide it into batches, and prefetch samples to enhance performance.
 
-```{code-cell}
+```{code-cell} ipython3
 import tensorflow_datasets as tfds  # TFDS to download MNIST.
 import tensorflow as tf  # TensorFlow / `tf.data` operations.
 
@@ -72,7 +72,7 @@ test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1)
 
 Create a CNN for classification with Flax NNX by subclassing `nnx.Module`:
 
-```{code-cell}
+```{code-cell} ipython3
 from flax import nnx  # The Flax NNX API.
 from functools import partial
 
@@ -82,19 +82,19 @@ class CNN(nnx.Module):
   def __init__(self, *, rngs: nnx.Rngs):
     self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)
     self.batch_norm1 = nnx.BatchNorm(32, rngs=rngs)
-    self.dropout1 = nnx.Dropout(rate=0.025, rngs=rngs)
+    self.dropout1 = nnx.Dropout(rate=0.025)
     self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)
     self.batch_norm2 = nnx.BatchNorm(64, rngs=rngs)
     self.avg_pool = partial(nnx.avg_pool, window_shape=(2, 2), strides=(2, 2))
     self.linear1 = nnx.Linear(3136, 256, rngs=rngs)
-    self.dropout2 = nnx.Dropout(rate=0.025, rngs=rngs)
+    self.dropout2 = nnx.Dropout(rate=0.025)
     self.linear2 = nnx.Linear(256, 10, rngs=rngs)
 
-  def __call__(self, x):
-    x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x)))))
+  def __call__(self, x, rngs: nnx.Rngs):
+    x = self.avg_pool(nnx.relu(self.batch_norm1(self.dropout1(self.conv1(x), rngs=rngs))))
     x = self.avg_pool(nnx.relu(self.batch_norm2(self.conv2(x))))
     x = x.reshape(x.shape[0], -1)  # flatten
-    x = nnx.relu(self.dropout2(self.linear1(x)))
+    x = nnx.relu(self.dropout2(self.linear1(x), rngs=rngs))
     x = self.linear2(x)
     return x
 
@@ -108,18 +108,18 @@ nnx.display(model)
 
 Let's put the CNN model to the test!  Here, you’ll perform a forward pass with arbitrary data and print the results.
 
-```{code-cell}
+```{code-cell} ipython3
 import jax.numpy as jnp  # JAX NumPy
 
-y = model(jnp.ones((1, 28, 28, 1)))
+y = model(jnp.ones((1, 28, 28, 1)), nnx.Rngs(0))
 y
 ```
 
 ## 4. Create the optimizer and define some metrics
 
 In Flax NNX, you need to create an `nnx.Optimizer` object to manage the model's parameters and apply gradients during training. `nnx.Optimizer` receives the model's reference, so that it can update its parameters, and an [Optax](https://optax.readthedocs.io/) optimizer to define the update rules. Additionally, you will define an `nnx.MultiMetric` object to keep track of the `Accuracy` and the `Average` loss.
 
-```{code-cell}
+```{code-cell} ipython3
 import optax
 
 learning_rate = 0.005
@@ -144,25 +144,25 @@ In addition to the `loss`, during training and testing you will also get the `lo
 
 During training - the `train_step` - you will use `nnx.value_and_grad` to compute the gradients and update the model's parameters using the `optimizer` you have already defined. And during both training and testing (the `eval_step`), the `loss` and `logits` will be used to calculate the metrics.
 
-```{code-cell}
-def loss_fn(model: CNN, batch):
-  logits = model(batch['image'])
+```{code-cell} ipython3
+def loss_fn(model: CNN, batch, rngs):
+  logits = model(batch['image'], rngs)
   loss = optax.softmax_cross_entropy_with_integer_labels(
     logits=logits, labels=batch['label']
   ).mean()
   return loss, logits
 
 @nnx.jit
-def train_step(model: CNN, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric, batch):
+def train_step(model: CNN, optimizer: nnx.Optimizer, metrics: nnx.MultiMetric, batch, rngs):
   """Train for a single step."""
   grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)
-  (loss, logits), grads = grad_fn(model, batch)
+  (loss, logits), grads = grad_fn(model, batch, rngs)
   metrics.update(loss=loss, logits=logits, labels=batch['label'])  # In-place updates.
-  optimizer.update(grads)  # In-place updates.
+  optimizer.update(model, grads)  # In-place updates.
 
 @nnx.jit
 def eval_step(model: CNN, metrics: nnx.MultiMetric, batch):
-  loss, logits = loss_fn(model, batch)
+  loss, logits = loss_fn(model, batch, None)
   metrics.update(loss=loss, logits=logits, labels=batch['label'])  # In-place updates.
 ```
 
@@ -177,7 +177,7 @@ Now, you can train the CNN model using batches of data for 10 epochs, evaluate t
 on the test set after each epoch, and log the training and testing metrics (the loss and
 the accuracy) during the process. Typically this leads to the model achieving around 99% accuracy.
 
-```{code-cell}
+```{code-cell} ipython3
 from IPython.display import clear_output
 import matplotlib.pyplot as plt
 
@@ -188,13 +188,15 @@ metrics_history = {
   'test_accuracy': [],
 }
 
+rngs = nnx.Rngs(0)
+
 for step, batch in enumerate(train_ds.as_numpy_iterator()):
   # Run the optimization for one step and make a stateful update to the following:
   # - The train state's model parameters
   # - The optimizer state
   # - The training loss and accuracy batch metrics
   model.train() # Switch to train mode
-  train_step(model, optimizer, metrics, batch)
+  train_step(model, optimizer, metrics, batch, rngs)
 
   if step > 0 and (step % eval_every == 0 or step == train_steps - 1):  # One training epoch has passed.
     # Log the training metrics.
@@ -229,18 +231,18 @@ for step, batch in enumerate(train_ds.as_numpy_iterator()):
 
 Create a `jit`-compiled model inference function (with `nnx.jit`) - `pred_step` - to generate predictions on the test set using the learned model parameters. This will enable you to visualize test images alongside their predicted labels for a qualitative assessment of model performance.
 
-```{code-cell}
+```{code-cell} ipython3
 model.eval() # Switch to evaluation mode.
 
 @nnx.jit
 def pred_step(model: CNN, batch):
-  logits = model(batch['image'])
+  logits = model(batch['image'], None)
   return logits.argmax(axis=1)
 ```
 
 We call .eval() before inference so Dropout is disabled and BatchNorm uses stored running stats. It is used during inference to suppress gradients and ensure deterministic, resource-efficient output.
 
-```{code-cell}
+```{code-cell} ipython3
 test_batch = test_ds.as_numpy_iterator().next()
 pred = pred_step(model, test_batch)
 
@@ -251,6 +253,40 @@ for i, ax in enumerate(axs.flatten()):
   ax.axis('off')
 ```
 
+# 8. Export the model
+
++++
+
+Flax models are great for research, but aren't meant to be deployed directly. Instead, high performance inference runtimes like LiteRT or TensorFlow Serving operate on a special [SavedModel](https://www.tensorflow.org/guide/saved_model) format. The [Orbax](https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html) library makes it easy to export Flax models to this format. First, we must create a `JaxModule` object wrapping a model and its prediction method.
+
+```{code-cell} ipython3
+from orbax.export import JaxModule, ExportManager, ServingConfig
+```
+
+```{code-cell} ipython3
+def exported_predict(model, y):
+    return model(y, None)
+    
+jax_module = JaxModule(model, exported_predict)
+```
+
+We also need to tell Tensorflow Serving what input type `exported_predict` expects in its second argument. The export machinery expects type signature arguments to be PyTrees of `tf.TensorSpec`.
+
+```{code-cell} ipython3
+sig = [tf.TensorSpec(shape=(1, 28, 28, 1), dtype=tf.float32)]
+```
+
+Finally, we can bundle up the input signature and the `JaxModule` together using the `ExportManager` class.
+
+```{code-cell} ipython3
+export_mgr = ExportManager(jax_module, [
+    ServingConfig('mnist_server', input_signature=sig)
+])
+
+output_dir='/tmp/mnist_export'
+export_mgr.save(output_dir)
+```
+
 Congratulations! You have learned how to use Flax NNX to build and train a simple classification model end-to-end on the MNIST dataset.
 
 Next, check out [Why Flax NNX?](https://flax.readthedocs.io/en/latest/why.html) and get started with a series of [Flax NNX Guides](https://flax.readthedocs.io/en/latest/guides/index.html).

docs_nnx/mnist_tutorial.md

@@ -20,8 +20,7 @@ https://colab.research.google.com/github/google/flax/blob/main/examples/mnist/mn
 [gs://flax_public/examples/mnist/default]: https://console.cloud.google.com/storage/browser/flax_public/examples/mnist/default
 
 ```
-I0828 08:51:41.821526 139971964110656 train.py:130] train epoch: 10, loss: 0.0097, accuracy: 99.69
-I0828 08:51:42.248714 139971964110656 train.py:180] eval epoch: 10, loss: 0.0299, accuracy: 99.14
+I1009 17:56:42.674334 3280981 train.py:175] epoch: 10, train_loss: 0.0073, train_accuracy: 99.75, test_loss: 0.0294, test_accuracy: 99.25
 ```
 
 ### How to run

examples/mnist/README.md

@@ -26,7 +26,7 @@
 from ml_collections import config_flags
 import tensorflow as tf
 
-import train
+import train  # pylint: disable=g-bad-import-order
 
 
 FLAGS = flags.FLAGS