wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite.ipynb
+485
View File
@@ -0,0 +1,485 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "8vD3L4qeREvg"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qLCxmWRyRMZE"
},
"outputs": [],
"source": [
"#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8LYgHRFPRpS1"
},
"source": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "4k5PoHrgJQOU"
},
"source": [
"# Jax Model Conversion For TFLite\n",
"## Overview\n",
"Note: This API is new and we recommend using via pip install tf-nightly. Also, the API is still experimental and subject to changes.\n",
"\n",
"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."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "i8cfOBcjSByO"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/examples/jax_conversion/jax_to_tflite\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lq-T8XZMJ-zv"
},
"source": [
"## Prerequisites\n",
"It's recommended to try this feature with the newest TensorFlow nightly pip build."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "EV04hKdrnE4f"
},
"outputs": [],
"source": [
"!pip install tf-nightly --upgrade\n",
"!pip install jax --upgrade"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "vsilblGuGQa2"
},
"outputs": [],
"source": [
"# Make sure your JAX version is at least 0.4.20 or above.\n",
"import jax\n",
"jax.__version__"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "PJeQhMUwH0oX"
},
"outputs": [],
"source": [
"!pip install orbax-export --upgrade"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "j9_CVA0THQNc"
},
"outputs": [],
"source": [
"from orbax.export import ExportManager\n",
"from orbax.export import JaxModule\n",
"from orbax.export import ServingConfig"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QAeY43k9KM55"
},
"source": [
"## Data Preparation\n",
"Download the MNIST data with Keras dataset and pre-process."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qSOPSZJn1_Tj"
},
"outputs": [],
"source": [
"import numpy as np\n",
"import tensorflow as tf\n",
"import functools\n",
"\n",
"import time\n",
"import itertools\n",
"\n",
"import numpy.random as npr\n",
"\n",
"import jax.numpy as jnp\n",
"from jax import jit, grad, random\n",
"from jax.example_libraries import optimizers\n",
"from jax.example_libraries import stax\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hdJIt3Da2Qn1"
},
"outputs": [],
"source": [
"def _one_hot(x, k, dtype=np.float32):\n",
" \"\"\"Create a one-hot encoding of x of size k.\"\"\"\n",
" return np.array(x[:, None] == np.arange(k), dtype)\n",
"\n",
"(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
"train_images, test_images = train_images / 255.0, test_images / 255.0\n",
"train_images = train_images.astype(np.float32)\n",
"test_images = test_images.astype(np.float32)\n",
"\n",
"train_labels = _one_hot(train_labels, 10)\n",
"test_labels = _one_hot(test_labels, 10)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0eFhx85YKlEY"
},
"source": [
"## Build the MNIST model with Jax"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mi3TKB9nnQdK"
},
"outputs": [],
"source": [
"def loss(params, batch):\n",
" inputs, targets = batch\n",
" preds = predict(params, inputs)\n",
" return -jnp.mean(jnp.sum(preds * targets, axis=1))\n",
"\n",
"def accuracy(params, batch):\n",
" inputs, targets = batch\n",
" target_class = jnp.argmax(targets, axis=1)\n",
" predicted_class = jnp.argmax(predict(params, inputs), axis=1)\n",
" return jnp.mean(predicted_class == target_class)\n",
"\n",
"init_random_params, predict = stax.serial(\n",
" stax.Flatten,\n",
" stax.Dense(1024), stax.Relu,\n",
" stax.Dense(1024), stax.Relu,\n",
" stax.Dense(10), stax.LogSoftmax)\n",
"\n",
"rng = random.PRNGKey(0)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "bRtnOBdJLd63"
},
"source": [
"## Train & Evaluate the model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "SWbYRyj7LYZt"
},
"outputs": [],
"source": [
"step_size = 0.001\n",
"num_epochs = 10\n",
"batch_size = 128\n",
"momentum_mass = 0.9\n",
"\n",
"\n",
"num_train = train_images.shape[0]\n",
"num_complete_batches, leftover = divmod(num_train, batch_size)\n",
"num_batches = num_complete_batches + bool(leftover)\n",
"\n",
"def data_stream():\n",
" rng = npr.RandomState(0)\n",
" while True:\n",
" perm = rng.permutation(num_train)\n",
" for i in range(num_batches):\n",
" batch_idx = perm[i * batch_size:(i + 1) * batch_size]\n",
" yield train_images[batch_idx], train_labels[batch_idx]\n",
"batches = data_stream()\n",
"\n",
"opt_init, opt_update, get_params = optimizers.momentum(step_size, mass=momentum_mass)\n",
"\n",
"@jit\n",
"def update(i, opt_state, batch):\n",
" params = get_params(opt_state)\n",
" return opt_update(i, grad(loss)(params, batch), opt_state)\n",
"\n",
"_, init_params = init_random_params(rng, (-1, 28 * 28))\n",
"opt_state = opt_init(init_params)\n",
"itercount = itertools.count()\n",
"\n",
"print(\"\\nStarting training...\")\n",
"for epoch in range(num_epochs):\n",
" start_time = time.time()\n",
" for _ in range(num_batches):\n",
" opt_state = update(next(itercount), opt_state, next(batches))\n",
" epoch_time = time.time() - start_time\n",
"\n",
" params = get_params(opt_state)\n",
" train_acc = accuracy(params, (train_images, train_labels))\n",
" test_acc = accuracy(params, (test_images, test_labels))\n",
" print(\"Epoch {} in {:0.2f} sec\".format(epoch, epoch_time))\n",
" print(\"Training set accuracy {}\".format(train_acc))\n",
" print(\"Test set accuracy {}\".format(test_acc))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7Y1OZBhfQhOj"
},
"source": [
"## Convert to TFLite model.\n",
"Note here, we\n",
"1. Export the `JAX` model to `TF SavedModel` using `orbax`.\n",
"2. Call TFLite converter API to convert the `TF SavedModel` to `.tflite` model:\n",
"\n",
"\n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6pcqKZqdNTmn"
},
"outputs": [],
"source": [
"jax_module = JaxModule(params, predict, input_polymorphic_shape='b, ...')\n",
"converter = tf.lite.TFLiteConverter.from_concrete_functions(\n",
" [\n",
" jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(\n",
" tf.TensorSpec(shape=(1, 28, 28), dtype=tf.float32, name=\"input\")\n",
" )\n",
" ]\n",
")\n",
"\n",
"tflite_model = converter.convert()\n",
"with open('jax_mnist.tflite', 'wb') as f:\n",
" f.write(tflite_model)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "sqEhzaJPSPS1"
},
"source": [
"## Check the Converted TFLite Model\n",
"Compare the converted model's results with the Jax model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "acj2AYzjSlaY"
},
"outputs": [],
"source": [
"serving_func = functools.partial(predict, params)\n",
"expected = serving_func(train_images[0:1])\n",
"\n",
"# Run the model with TensorFlow Lite\n",
"interpreter = tf.lite.Interpreter(model_content=tflite_model)\n",
"interpreter.allocate_tensors()\n",
"input_details = interpreter.get_input_details()\n",
"output_details = interpreter.get_output_details()\n",
"interpreter.set_tensor(input_details[0][\"index\"], train_images[0:1, :, :])\n",
"interpreter.invoke()\n",
"result = interpreter.get_tensor(output_details[0][\"index\"])\n",
"\n",
"# Assert if the result of TFLite model is consistent with the JAX model.\n",
"np.testing.assert_almost_equal(expected, result, 1e-5)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Qy9Gp4H2SjBL"
},
"source": [
"## Optimize the Model\n",
"We will provide a `representative_dataset` to do post-training quantiztion to optimize the model.\n",
"\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "KI0rLV-Meg-2"
},
"outputs": [],
"source": [
"def representative_dataset():\n",
" for i in range(1000):\n",
" x = train_images[i:i+1]\n",
" yield [x]\n",
"x_input = jnp.zeros((1, 28, 28))\n",
"converter = tf.lite.TFLiteConverter.experimental_from_jax(\n",
" [serving_func], [[('x', x_input)]])\n",
"tflite_model = converter.convert()\n",
"converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
"converter.representative_dataset = representative_dataset\n",
"converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
"tflite_quant_model = converter.convert()\n",
"with open('jax_mnist_quant.tflite', 'wb') as f:\n",
" f.write(tflite_quant_model)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "15xQR3JZS8TV"
},
"source": [
"## Evaluate the Optimized Model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "X3oOm0OaevD6"
},
"outputs": [],
"source": [
"expected = serving_func(train_images[0:1])\n",
"\n",
"# Run the model with TensorFlow Lite\n",
"interpreter = tf.lite.Interpreter(model_content=tflite_quant_model)\n",
"interpreter.allocate_tensors()\n",
"input_details = interpreter.get_input_details()\n",
"output_details = interpreter.get_output_details()\n",
"interpreter.set_tensor(input_details[0][\"index\"], train_images[0:1, :, :])\n",
"interpreter.invoke()\n",
"result = interpreter.get_tensor(output_details[0][\"index\"])\n",
"\n",
"# Assert if the result of TFLite model is consistent with the Jax model.\n",
"np.testing.assert_almost_equal(expected, result, 1e-5)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "QqHXCNa3myor"
},
"source": [
"## Compare the Quantized Model size\n",
"We should be able to see the quantized model is four times smaller than the original model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "imFPw007juVG"
},
"outputs": [],
"source": [
"!du -h jax_mnist.tflite\n",
"!du -h jax_mnist_quant.tflite"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "N5WdO18wNfyn"
},
"outputs": [],
"source": []
}
],
"metadata": {
"colab": {
"private_outputs": true,
"provenance": [
{
"file_id": "1UlzbQspn2an2kzlLWZBWhP_JEqvShxmi",
"timestamp": 1705015454450
},
{
"file_id": "https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/examples/jax_conversion/overview.ipynb",
"timestamp": 1698963811786
}
],
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite_resnet50.ipynb
+517
View File
@@ -0,0 +1,517 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "eB-dvsMI09O4"
},
"source": [
"##### Copyright 2024 The AI Edge Authors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "mwvC53CC1K3n"
},
"outputs": [],
"source": [
"# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0WlR2-Y3v1mf"
},
"source": [
"End-to-end Example to show how a pre-trained FLAX model can be downloaded, exported and ran.\n",
"===============\n",
"1. The model used is Resnet50, downloaded from huggingface pre-trained models\n",
"2. Test image is of a CAT\n",
"3. `orbax-export` API is used to export the JAX Module to a TF Saved Model, along with image pre/post-processing functions.\n",
"4. TFLite Converter & Runtime are used to convert the TF Saved Model to .tflite and run on the test image.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zQXWQ7y11eIR"
},
"source": [
"<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://www.tensorflow.org/lite/examples/jax_conversion/jax_to_tflite_resnet50\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite_resnet50.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
" </td>\n",
" <td>\n",
" <a target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite_resnet50.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
" </td>\n",
" <td>\n",
" <a href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite_resnet50.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
" </td>\n",
"</table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gq1373VwVlKV"
},
"source": [
"# Setup"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TUOym6hJnTxC"
},
"source": [
"```\n",
"!pip install orbax-export\n",
"\n",
"!pip install tf-nightly\n",
"\n",
"!pip install --upgrade jax jaxlib\n",
"\n",
"!pip install transformers flax\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hDzBKqEfrafW"
},
"source": [
"## Show Test Image"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2IOV4p_jripo"
},
"outputs": [],
"source": [
"from PIL import Image\n",
"import jax\n",
"import requests\n",
"\n",
"url = \"https://storage.googleapis.com/download.tensorflow.org/example_images/astrid_l_shaped.jpg\"\n",
"image = Image.open(requests.get(url, stream=True).raw)\n",
"image"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SPcMyI8ERmIB"
},
"source": [
"## Download and test pre-trained Resnet50 FLAX model from HuggingFace"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "a__LrY62PUtg"
},
"outputs": [],
"source": [
"from transformers import ConvNextImageProcessor, FlaxResNetForImageClassification\n",
"\n",
"image_processor = ConvNextImageProcessor.from_pretrained(\"microsoft/resnet-50\")\n",
"model = FlaxResNetForImageClassification.from_pretrained(\"microsoft/resnet-50\")\n",
"\n",
"inputs = image_processor(images=image, return_tensors=\"np\")\n",
"outputs = model(**inputs)\n",
"logits = outputs.logits\n",
"\n",
"# model predicts one of the 1000 ImageNet classes\n",
"predicted_class_idx = jax.numpy.argmax(logits, axis=-1)\n",
"print(\"Predicted class:\", model.config.id2label[predicted_class_idx.item()])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "MKJHRRjGR1rC"
},
"source": [
"# Create a JAX wrapper for the model\n",
"Wrapper is needed in order to comply with TFLite accepts inputs. TFLite accets a tensor or a tuple-of-tensors."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Af2zv1qFSOR3"
},
"outputs": [],
"source": [
"import flax.linen as nn\n",
"from transformers import FlaxResNetForImageClassification\n",
"\n",
"\n",
"class Resnet50Wrapper(nn.Module):\n",
" pretrained_model_name: str = \"microsoft/resnet-50\" # Pre-trained model name\n",
"\n",
" def setup(self):\n",
" # Initialize the pre-trained ResNet50 model\n",
" self.model = FlaxResNetForImageClassification.from_pretrained(\n",
" self.pretrained_model_name\n",
" )\n",
"\n",
" def __call__(self, inputs):\n",
" # Process input images through the ResNet50 model\n",
" outputs = self.model(pixel_values=inputs)\n",
"\n",
" # Return logits or directly apply softmax for probabilities (optional)\n",
" return outputs.logits"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mrKtRfkhf73M"
},
"source": [
"## Tensorflow image pre-processor function\n",
"This essentialliy implements the underlying logic of the `ConvNextImageProcessor` class in huggingface transformers:\n",
"> ```\n",
"> image_processor = ConvNextImageProcessor.from_pretrained(\"microsoft/resnet-50\")\n",
"> inputs = image_processor(images=image, return_tensors=\"np\")\n",
"> ```\n",
"\n",
"This utility can be reused later during orbax-export for `tf_preprocessing`.\n",
"\n",
"Note: We can perfectly use the result of `ConvNextImageProcessor` to run a TFLite model. But this example would like to showcase how `orbax-export` helps handle input/output pre/post-processing."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "h85-Qb1PgANA"
},
"outputs": [],
"source": [
"import tensorflow as tf\n",
"import numpy as np\n",
"\n",
"\n",
"def resnet_image_processor(image_tensor):\n",
" # 1. Resize and Cast to Float32\n",
" image_resized = tf.image.resize(\n",
" image_tensor, (224, 224), method=tf.image.ResizeMethod.BILINEAR\n",
" )\n",
" image_float = tf.cast(image_resized, tf.float32)\n",
"\n",
" # 2. Normalize (Using TensorFlow Constants)\n",
" mean = tf.constant([0.485, 0.456, 0.406])\n",
" std = tf.constant([0.229, 0.224, 0.225])\n",
" image_normalized = (image_float / 255.0 - mean) / std\n",
"\n",
" # 3. Transpose for Channel-First Format\n",
" image_transposed = tf.transpose(image_normalized, perm=[2, 0, 1])\n",
"\n",
" # 4. Add Batch Dimension\n",
" return tf.expand_dims(image_transposed, axis=0)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3GD4Wb2K3VT_"
},
"source": [
"## Initialize the Jax model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2wHvNFpiThVL"
},
"outputs": [],
"source": [
"# Initialize the JAX Model\n",
"jax_model = Resnet50Wrapper()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "uwP7XjS43M7W"
},
"source": [
"## Validate the wrapped JAX Model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Gl8mypra3MGw"
},
"outputs": [],
"source": [
"# Convert the raw image values to RGB tensor\n",
"raw_image_tensor = tf.convert_to_tensor(np.array(image, dtype=np.float32))\n",
"\n",
"# Appy the above TF imape preprocessing to get an input tensor supported by Resnet50\n",
"input_tensor = resnet_image_processor(raw_image_tensor)\n",
"\n",
"# Run the JAX model\n",
"jax_logits = jax_model.apply({}, input_tensor.numpy())\n",
"\n",
"jax_predicted_class_idx = jax.numpy.argmax(jax_logits, axis=-1)\n",
"print(\"Predicted class:\", model.config.id2label[jax_predicted_class_idx.item()])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "lKZst83Mu6TD"
},
"outputs": [],
"source": [
"raw_image_tensor.shape"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zuAl8ubTTKxZ"
},
"source": [
"# Export to TFLite model and run"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Pz61IKnwjhbn"
},
"source": [
"## Export the JAX to TF Saved Model using orbax-export"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Cwq1rxEZjpoU"
},
"outputs": [],
"source": [
"from orbax.export import ExportManager, JaxModule, ServingConfig\n",
"\n",
"# Wrap the model params and function into a JaxModule.\n",
"jax_module = JaxModule({}, jax_model.apply, trainable=False)\n",
"\n",
"# Specify the serving configuration and export the model.\n",
"serving_config = ServingConfig(\n",
" \"serving_default\",\n",
" input_signature=[tf.TensorSpec([480, 640, 3], tf.float32, name=\"inputs\")],\n",
" tf_preprocessor=resnet_image_processor,\n",
" tf_postprocessor=lambda x: tf.argmax(x, axis=-1),\n",
")\n",
"\n",
"export_manager = ExportManager(jax_module, [serving_config])\n",
"\n",
"saved_model_dir = \"resnet50_saved_model\"\n",
"export_manager.save(saved_model_dir)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "O00xe8MevVLj"
},
"source": [
"### Convert to a TFLite Model"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Uf2-ZMECmtJ3"
},
"outputs": [],
"source": [
"converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)\n",
"tflite_model = converter.convert()"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "n-fBAd64vZse"
},
"source": [
"### Apply TFLite Runtime API on the `raw_image_tensor`"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "unFFby0Zmup-"
},
"outputs": [],
"source": [
"def run_tflite_model(tflite_model_content, input_tensor):\n",
" interpreter = tf.lite.Interpreter(model_content=tflite_model_content)\n",
" interpreter.allocate_tensors()\n",
"\n",
" input_details = interpreter.get_input_details()[0]\n",
" interpreter.set_tensor(input_details[\"index\"], input_tensor)\n",
" interpreter.invoke()\n",
"\n",
" output_details = interpreter.get_output_details()\n",
" return interpreter.get_tensor(output_details[0][\"index\"])"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Z9Rwk8iFvSBO"
},
"outputs": [],
"source": [
"output_data = run_tflite_model(tflite_model, raw_image_tensor)\n",
"print(\"Predicted class:\", model.config.id2label[output_data[0]])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "NbtCrNLYKpK8"
},
"source": [
"## Export to TF Saved Model without pre/post-processing"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qPLmCDniKvV3"
},
"outputs": [],
"source": [
"saved_model_dir_2 = \"resnet50_saved_model_1\"\n",
"\n",
"tf.saved_model.save(\n",
" jax_module,\n",
" saved_model_dir_2,\n",
" signatures=jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(\n",
" tf.TensorSpec([1, 3, 224, 224], tf.float32, name=\"inputs\")\n",
" ),\n",
" options=tf.saved_model.SaveOptions(experimental_custom_gradients=True),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Wzp8M2ebLYoH"
},
"outputs": [],
"source": [
"converter_1 = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir_2)\n",
"tflite_model_1 = converter_1.convert()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "3Di5ikJYMcKy"
},
"outputs": [],
"source": [
"output_data_1 = run_tflite_model(tflite_model_1, input_tensor)\n",
"tfl_predicted_class_idx_1 = tf.argmax(output_data_1, axis=-1).numpy()\n",
"print(\"Predicted class:\", model.config.id2label[tfl_predicted_class_idx_1[0]])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_SwsbUH_Smt7"
},
"source": [
"## Export from TF Function"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2MbcN1vjSt5E"
},
"outputs": [],
"source": [
"converter_2 = tf.lite.TFLiteConverter.from_concrete_functions(\n",
" [\n",
" jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(\n",
" tf.TensorSpec([1, 3, 224, 224], tf.float32, name=\"inputs\")\n",
" )\n",
" ]\n",
")\n",
"tflite_model_2 = converter_2.convert()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WdM3p03AS-xq"
},
"outputs": [],
"source": [
"output_data_2 = run_tflite_model(tflite_model_2, input_tensor)\n",
"tfl_predicted_class_idx_2 = tf.argmax(output_data_2, axis=-1).numpy()\n",
"print(\"Predicted class:\", model.config.id2label[tfl_predicted_class_idx_2[0]])"
]
}
],
"metadata": {
"colab": {
"name": "jax_to_tflite_resnet50.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
tensorflow/lite/g3doc/examples/jax_conversion/overview.md
+98
View File
@@ -0,0 +1,98 @@
# JAX models with TensorFlow Lite
This page provides a path for users who want to train models in JAX and deploy
to mobile for inference ([example colab](https://colab.sandbox.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/examples/jax_conversion/jax_to_tflite.ipynb)).
The methods in this guide produce a `tflite_model` which can be used directly
with the TFLite interpreter code example or saved to a TFLite FlatBuffer file.
## Prerequisite
It's recommended to try this feature with the newest TensorFlow nightly Python
package.
```
pip install tf-nightly --upgrade
```
We will use the
[Orbax Export](https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html)
library to export JAX models. Make sure your JAX version is at least 0.4.20 or
above.
```
pip install jax --upgrade
pip install orbax-export --upgrade
```
## Convert JAX models to TensorFlow Lite
We use the TensorFlow
[SavedModel](https://www.tensorflow.org/guide/saved_model) as the intermediate
format between JAX and TensorFlow Lite. Once you have a SavedModel then
existing TensorFlow Lite APIs can be used to complete the conversion process.
```py
# This code snippet converts a JAX model to TFLite through TF SavedModel.
from orbax.export import ExportManager
from orbax.export import JaxModule
from orbax.export import ServingConfig
import tensorflow as tf
import jax.numpy as jnp
def model_fn(_, x):
return jnp.sin(jnp.cos(x))
jax_module = JaxModule({}, model_fn, input_polymorphic_shape='b, ...')
# Option 1: Simply save the model via `tf.saved_model.save` if no need for pre/post
# processing.
tf.saved_model.save(
jax_module,
'/some/directory',
signatures=jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
),
options=tf.saved_model.SaveOptions(experimental_custom_gradients=True),
)
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 2: Define pre/post processing TF functions (e.g. (de)?tokenize).
serving_config = ServingConfig(
'Serving_default',
# Corresponds to the input signature of `tf_preprocessor`
input_signature=[tf.TensorSpec(shape=(None,), dtype=tf.float32, name='input')],
tf_preprocessor=lambda x: x,
tf_postprocessor=lambda out: {'output': out}
)
export_mgr = ExportManager(jax_module, [serving_config])
export_mgr.save('/some/directory')
converter = tf.lite.TFLiteConverter.from_saved_model('/some/directory')
tflite_model = converter.convert()
# Option 3: Convert from TF concrete function directly
converter = tf.lite.TFLiteConverter.from_concrete_functions(
[
jax_module.methods[JaxModule.DEFAULT_METHOD_KEY].get_concrete_function(
tf.TensorSpec(shape=(None,), dtype=tf.float32, name="input")
)
]
)
tflite_model = converter.convert()
```
## Check the converted TFLite model
After the model is converted to TFLite, you can run TFLite interpreter APIs to
check model outputs.
```py
# 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"], input_data)
interpreter.invoke()
result = interpreter.get_tensor(output_details[0]["index"])
```