Thanks for tuning in to Google I/O. View all sessions on demandWatch on demand

Jax Model Conversion For TFLite

Overview

This CodeLab demonstrates how to build a model for MNIST recognition using Jax, and how to convert it to TensorFlow Lite. This codelab will also demonstrate how to optimize the Jax-converted TFLite model with post-training quantiztion.

View on TensorFlow.org Run in Google Colab View source on GitHub Download notebook

Prerequisites

It's recommended to try this feature with the newest TensorFlow nightly pip build.

pip install tf-nightly --upgrade
pip install jax --upgrade
pip install jaxlib --upgrade

Data Preparation

Download the MNIST data with Keras dataset and pre-process.

import numpy as np
import tensorflow as tf
import functools

import time
import itertools

import numpy.random as npr

import jax.numpy as jnp
from jax import jit, grad, random
from jax.example_libraries import optimizers
from jax.example_libraries import stax

2023-05-11 11:13:45.244834: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:7704] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-05-11 11:13:45.244880: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-05-11 11:13:45.244889: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1520] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

def _one_hot(x, k, dtype=np.float32):
  """Create a one-hot encoding of x of size k."""
  return np.array(x[:, None] == np.arange(k), dtype)

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_images, test_images = train_images / 255.0, test_images / 255.0
train_images = train_images.astype(np.float32)
test_images = test_images.astype(np.float32)

train_labels = _one_hot(train_labels, 10)
test_labels = _one_hot(test_labels, 10)

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11490434/11490434 [==============================] - 0s 0us/step

Build the MNIST model with Jax

def loss(params, batch):
  inputs, targets = batch
  preds = predict(params, inputs)
  return -jnp.mean(jnp.sum(preds * targets, axis=1))

def accuracy(params, batch):
  inputs, targets = batch
  target_class = jnp.argmax(targets, axis=1)
  predicted_class = jnp.argmax(predict(params, inputs), axis=1)
  return jnp.mean(predicted_class == target_class)

init_random_params, predict = stax.serial(
    stax.Flatten,
    stax.Dense(1024), stax.Relu,
    stax.Dense(1024), stax.Relu,
    stax.Dense(10), stax.LogSoftmax)

rng = random.PRNGKey(0)

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

Train & Evaluate the model

step_size = 0.001
num_epochs = 10
batch_size = 128
momentum_mass = 0.9


num_train = train_images.shape[0]
num_complete_batches, leftover = divmod(num_train, batch_size)
num_batches = num_complete_batches + bool(leftover)

def data_stream():
  rng = npr.RandomState(0)
  while True:
    perm = rng.permutation(num_train)
    for i in range(num_batches):
      batch_idx = perm[i * batch_size:(i + 1) * batch_size]
      yield train_images[batch_idx], train_labels[batch_idx]
batches = data_stream()

opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)

@jit
def update(i, opt_state, batch):
  params = get_params(opt_state)
  return opt_update(i, grad(loss)(params, batch), opt_state)

_, init_params = init_random_params(rng, (-1, 28 * 28))
opt_state = opt_init(init_params)
itercount = itertools.count()

print("\nStarting training...")
for epoch in range(num_epochs):
  start_time = time.time()
  for _ in range(num_batches):
    opt_state = update(next(itercount), opt_state, next(batches))
  epoch_time = time.time() - start_time

  params = get_params(opt_state)
  train_acc = accuracy(params, (train_images, train_labels))
  test_acc = accuracy(params, (test_images, test_labels))
  print("Epoch {} in {:0.2f} sec".format(epoch, epoch_time))
  print("Training set accuracy {}".format(train_acc))
  print("Test set accuracy {}".format(test_acc))

Starting training...
Epoch 0 in 4.22 sec
Training set accuracy 0.8728833198547363
Test set accuracy 0.880299985408783
Epoch 1 in 2.30 sec
Training set accuracy 0.8983833193778992
Test set accuracy 0.9047999978065491
Epoch 2 in 2.25 sec
Training set accuracy 0.9102333188056946
Test set accuracy 0.9138000011444092
Epoch 3 in 2.26 sec
Training set accuracy 0.9172333478927612
Test set accuracy 0.9218999743461609
Epoch 4 in 2.25 sec
Training set accuracy 0.9224833250045776
Test set accuracy 0.9253999590873718
Epoch 5 in 2.25 sec
Training set accuracy 0.9272000193595886
Test set accuracy 0.9309999942779541
Epoch 6 in 2.25 sec
Training set accuracy 0.9328166842460632
Test set accuracy 0.9334999918937683
Epoch 7 in 2.25 sec
Training set accuracy 0.9360166788101196
Test set accuracy 0.9370999932289124
Epoch 8 in 2.24 sec
Training set accuracy 0.939050018787384
Test set accuracy 0.939300000667572
Epoch 9 in 2.24 sec
Training set accuracy 0.9425666928291321
Test set accuracy 0.9429000020027161

Convert to TFLite model.

Note here, we

  1. Inline the params to the Jax predict func with functools.partial.
  2. Build a jnp.zeros, this is a "placeholder" tensor used for Jax to trace the model.
  3. Call experimental_from_jax: > * The serving_func is wrapped in a list. > * The input is associated with a given name and passed in as an array wrapped in a list.
serving_func = functools.partial(predict, params)
x_input = jnp.zeros((1, 28, 28))
converter = tf.lite.TFLiteConverter.experimental_from_jax(
    [serving_func], [[('input1', x_input)]])
tflite_model = converter.convert()
with open('jax_mnist.tflite', 'wb') as f:
  f.write(tflite_model)

2023-05-11 11:14:20.786709: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:364] Ignored output_format.
2023-05-11 11:14:20.786758: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:367] Ignored drop_control_dependency.
2023-05-11 11:14:20.786765: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:373] Ignored change_concat_input_ranges.

Check the Converted TFLite Model

Compare the converted model's results with the Jax model.

expected = serving_func(train_images[0:1])

# Run the model with TensorFlow Lite
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]["index"], train_images[0:1, :, :])
interpreter.invoke()
result = interpreter.get_tensor(output_details[0]["index"])

# Assert if the result of TFLite model is consistent with the JAX model.
np.testing.assert_almost_equal(expected, result, 1e-5)

INFO: Created TensorFlow Lite XNNPACK delegate for CPU.

Optimize the Model

We will provide a representative_dataset to do post-training quantiztion to optimize the model.

def representative_dataset():
  for i in range(1000):
    x = train_images[i:i+1]
    yield [x]

converter = tf.lite.TFLiteConverter.experimental_from_jax(
    [serving_func], [[('x', x_input)]])
tflite_model = converter.convert()
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_dataset
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
tflite_quant_model = converter.convert()
with open('jax_mnist_quant.tflite', 'wb') as f:
  f.write(tflite_quant_model)

2023-05-11 11:14:21.182441: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:364] Ignored output_format.
2023-05-11 11:14:21.182479: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:367] Ignored drop_control_dependency.
2023-05-11 11:14:21.182485: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:373] Ignored change_concat_input_ranges.
2023-05-11 11:14:21.391493: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:364] Ignored output_format.
2023-05-11 11:14:21.391536: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:367] Ignored drop_control_dependency.
2023-05-11 11:14:21.391544: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:373] Ignored change_concat_input_ranges.
fully_quantize: 0, inference_type: 6, input_inference_type: FLOAT32, output_inference_type: FLOAT32

Evaluate the Optimized Model

expected = serving_func(train_images[0:1])

# Run the model with TensorFlow Lite
interpreter = tf.lite.Interpreter(model_content=tflite_quant_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]["index"], train_images[0:1, :, :])
interpreter.invoke()
result = interpreter.get_tensor(output_details[0]["index"])

# Assert if the result of TFLite model is consistent with the Jax model.
np.testing.assert_almost_equal(expected, result, 1e-5)

Compare the Quantized Model size

We should be able to see the quantized model is four times smaller than the original model.

du -h jax_mnist.tflite
du -h jax_mnist_quant.tflite

7.2M    jax_mnist.tflite
1.8M    jax_mnist_quant.tflite