Guide: Exporting Keras models to LiteRT with PyTorch backend

guides/ipynb/litert_export.ipynb

@@ -0,0 +1,380 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "# Exporting Keras models to LiteRT\n",
+    "\n",
+    "**Author:** [Rahul Kumar](https://github.com/pctablet505)<br>\n",
+    "**Date created:** 2025/12/10<br>\n",
+    "**Last modified:** 2026/06/02<br>\n",
+    "**Description:** Learn how to export Keras models to LiteRT for mobile and edge deployment using the PyTorch backend."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "## Introduction\n",
+    "\n",
+    "[LiteRT](https://ai.google.dev/edge/litert) (formerly TensorFlow Lite) lets you run\n",
+    "machine learning models on mobile devices, embedded systems, and browsers with\n",
+    "low latency and small binary size.\n",
+    "\n",
+    "This guide shows how to export a Keras model to LiteRT format using the\n",
+    "built-in `model.export()` API, run inference with the new LiteRT interpreter,\n",
+    "and apply quantization for smaller file sizes.\n",
+    "\n",
+    "We recommend the **PyTorch backend** for LiteRT export because:\n",
+    "\n",
+    "1. **No flex ops** \u2014 the TensorFlow backend path uses\n",
+    "   `tf.lite.TFLiteConverter.from_saved_model()`, which enables\n",
+    "   `SELECT_TF_OPS` (flex ops) by default. Flex ops are not supported by the\n",
+    "   new LiteRT Android runtime and the `ai_edge_litert` interpreter.\n",
+    "2. **Future-proof interpreter** \u2014 `tf.lite.Interpreter` is deprecated and\n",
+    "   scheduled for removal. The new `ai_edge_litert.interpreter.Interpreter` is\n",
+    "   the supported path, and it requires models without flex ops.\n",
+    "3. **Native PyTorch-to-LiteRT pipeline** \u2014 `litert-torch` converts Keras models\n",
+    "   through PyTorch's `ExportedProgram` directly to the LiteRT flatbuffer format.\n",
+    "\n",
+    "We will use a tiny transformer as our runnable example, but the same steps work\n",
+    "for any Keras model \u2014 including **Gemma 3 270M**, a lightweight language model\n",
+    "that is small enough to export quickly and run on edge devices.\n",
+    "\n",
+    "## Setup\n",
+    "\n",
+    "Install the required packages:\n",
+    "\n",
+    "```shell\n",
+    "pip install -q keras keras-hub ai-edge-litert\n",
+    "```\n",
+    "\n",
+    "> **Note:** LiteRT export with the PyTorch backend requires `litert-torch`.\n",
+    "> If it is not already installed, run:\n",
+    "> ```shell\n",
+    "> pip install -q litert-torch\n",
+    "> ```\n",
+    "\n",
+    "Set the PyTorch backend before importing Keras:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"KERAS_BACKEND\"] = \"torch\"\n",
+    "\n",
+    "import numpy as np\n",
+    "import keras\n",
+    "\n",
+    "print(\"Keras version:\", keras.__version__)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "## Export a Keras model\n",
+    "\n",
+    "Build a tiny transformer and export it to LiteRT. The same API works for\n",
+    "KerasHub presets such as `gemma3_270m`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "vocab_size = 256\n",
+    "seq_length = 8\n",
+    "embed_dim = 16\n",
+    "num_heads = 2\n",
+    "dim_feedforward = 32\n",
+    "num_layers = 2\n",
+    "\n",
+    "inputs = keras.Input(shape=(seq_length,), dtype=\"int32\")\n",
+    "x = keras.layers.Embedding(vocab_size, embed_dim)(inputs)\n",
+    "for _ in range(num_layers):\n",
+    "    attn = keras.layers.MultiHeadAttention(\n",
+    "        num_heads=num_heads, key_dim=embed_dim // num_heads\n",
+    "    )(x, x)\n",
+    "    x = keras.layers.LayerNormalization()(x + attn)\n",
+    "    ffn = keras.layers.Dense(dim_feedforward, activation=\"relu\")(x)\n",
+    "    ffn = keras.layers.Dense(embed_dim)(ffn)\n",
+    "    x = keras.layers.LayerNormalization()(x + ffn)\n",
+    "outputs = keras.layers.Dense(vocab_size)(x)\n",
+    "model = keras.Model(inputs, outputs)\n",
+    "model.compile()\n",
+    "\n",
+    "# Build weights with a sample input\n",
+    "sample_input = np.zeros((1, seq_length), dtype=\"int32\")\n",
+    "_ = model(sample_input)\n",
+    "\n",
+    "# Export to LiteRT\n",
+    "model.export(\"model.tflite\", format=\"litert\")\n",
+    "print(\"Exported to model.tflite\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "## Run inference with the LiteRT model\n",
+    "\n",
+    "Load the exported `.tflite` file with `ai_edge_litert` and run inference.\n",
+    "\n",
+    "> **Important:** `tf.lite.Interpreter` is deprecated and scheduled for deletion.\n",
+    "> Always use `ai_edge_litert.interpreter.Interpreter` for new code."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "from ai_edge_litert.interpreter import Interpreter\n",
+    "\n",
+    "interpreter = Interpreter(model_path=\"model.tflite\")\n",
+    "interpreter.allocate_tensors()\n",
+    "\n",
+    "input_details = interpreter.get_input_details()\n",
+    "output_details = interpreter.get_output_details()\n",
+    "\n",
+    "print(\"Input shape :\", input_details[0][\"shape\"])\n",
+    "print(\"Output shape:\", output_details[0][\"shape\"])\n",
+    "\n",
+    "interpreter.set_tensor(input_details[0][\"index\"], sample_input)\n",
+    "interpreter.invoke()\n",
+    "output = interpreter.get_tensor(output_details[0][\"index\"])\n",
+    "print(\"Inference output shape:\", output.shape)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "## Quantization for smaller models\n",
+    "\n",
+    "Quantization reduces model size and can speed up inference on edge devices.\n",
+    "\n",
+    "### Post-export quantization with ai-edge-quantizer (recommended)\n",
+    "\n",
+    "For the PyTorch backend, we recommend quantizing the model **after** export\n",
+    "using the **AI Edge Quantizer**. It provides finer-grained control (e.g.\n",
+    "channel-wise symmetric INT8 weights) that preserves quality better for\n",
+    "generative models, and does not require passing `optimizations` to the\n",
+    "exporter.\n",
+    "\n",
+    "Install it:\n",
+    "\n",
+    "```shell\n",
+    "pip install -q ai-edge-quantizer\n",
+    "```\n",
+    "\n",
+    "Quantize the already-exported `.tflite` file:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "from ai_edge_quantizer import quantizer, qtyping\n",
+    "\n",
+    "qt = quantizer.Quantizer(\"model.tflite\")\n",
+    "\n",
+    "# Recipe: channel-wise symmetric INT8 weights, FP32 activations.\n",
+    "# This is similar to dynamic-range quantization but uses channel-wise\n",
+    "# scaling which is more accurate for transformer attention weights.\n",
+    "recipe = [\n",
+    "    {\n",
+    "        \"regex\": \".*\",\n",
+    "        \"operation\": qtyping.TFLOperationName.ALL_SUPPORTED,\n",
+    "        \"algorithm_key\": quantizer.AlgorithmName.MIN_MAX_UNIFORM_QUANT,\n",
+    "        \"op_config\": {\n",
+    "            \"weight_tensor_config\": {\n",
+    "                \"dtype\": qtyping.TensorDataType.INT,\n",
+    "                \"num_bits\": 8,\n",
+    "                \"granularity\": qtyping.QuantGranularity.CHANNELWISE,\n",
+    "                \"symmetric\": True,\n",
+    "            },\n",
+    "            \"compute_precision\": qtyping.ComputePrecision.FLOAT,\n",
+    "            \"explicit_dequantize\": False,\n",
+    "        },\n",
+    "    },\n",
+    "]\n",
+    "\n",
+    "qt.load_quantization_recipe(recipe)\n",
+    "\n",
+    "# No calibration needed for this weight-only recipe.\n",
+    "quant_result = qt.quantize()\n",
+    "quant_result.save(\".\", model_name=\"model_aieq\")\n",
+    "\n",
+    "print(\"Exported ai-edge-quantizer model\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "The AI Edge Quantizer typically achieves:\n",
+    "- **Similar file size** to TFLite dynamic-range quantization\n",
+    "- **Better perplexity / BLEU scores** on generative tasks thanks to\n",
+    "  channel-wise symmetric scaling\n",
+    "- **No calibration required** for weight-only recipes\n",
+    "\n",
+    "### Why not use `optimizations` with the PyTorch backend?\n",
+    "\n",
+    "Passing `optimizations=[tf.lite.Optimize.DEFAULT]` to `model.export()` works on\n",
+    "the **TensorFlow** backend but can fail on the **PyTorch** backend due to\n",
+    "differences in how `litert-torch` handles quantized dtypes during the\n",
+    "MLIR pipeline. Until this is fully stabilized, post-export quantization with\n",
+    "`ai-edge-quantizer` is the safer and more accurate path.\n",
+    "\n",
+    "## Compare file sizes"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 0,
+   "metadata": {
+    "colab_type": "code"
+   },
+   "outputs": [],
+   "source": [
+    "\n",
+    "def file_size_mb(path):\n",
+    "    return os.path.getsize(path) / (1024 * 1024)\n",
+    "\n",
+    "\n",
+    "print(f\"\\nOriginal     : {file_size_mb('model.tflite'):.1f} MB\")\n",
+    "print(f\"AI Edge Quant: {file_size_mb('model_aieq.tflite'):.1f} MB\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "colab_type": "text"
+   },
+   "source": [
+    "## Exporting a KerasHub model\n",
+    "\n",
+    "The same API works for any Keras model, including KerasHub presets.\n",
+    "Here is how you would export **Gemma 3 270M**:\n",
+    "\n",
+    "```python\n",
+    "import keras_hub\n",
+    "\n",
+    "preset = \"gemma3_270m\"\n",
+    "preprocessor = keras_hub.models.Gemma3CausalLMPreprocessor.from_preset(\n",
+    "    preset, sequence_length=32\n",
+    ")\n",
+    "model = keras_hub.models.Gemma3CausalLM.from_preset(\n",
+    "    preset, preprocessor=preprocessor\n",
+    ")\n",
+    "\n",
+    "sample_text = {\"prompts\": [\"Hello\"], \"responses\": [\"world\"]}\n",
+    "sample_input = preprocessor(sample_text)[0]\n",
+    "_ = model(sample_input)\n",
+    "\n",
+    "model.export(\"gemma3_270m.tflite\", format=\"litert\")\n",
+    "```\n",
+    "\n",
+    "## Backend comparison\n",
+    "\n",
+    "| Feature | TensorFlow backend | PyTorch backend (recommended) |\n",
+    "|---|---|---|\n",
+    "| Converter path | `tf.lite.TFLiteConverter` | `litert-torch` via `ExportedProgram` |\n",
+    "| Flex ops | Enabled by default | **Not generated** |\n",
+    "| Android runtime | May crash on new LiteRT | Fully compatible |\n",
+    "| Interpreter | `tf.lite.Interpreter` (deprecated) | `ai_edge_litert.interpreter.Interpreter` |\n",
+    "| Quantization via `optimizations` | Supported | Can fail for some models |\n",
+    "| Post-export `ai-edge-quantizer` | Supported | **Supported** |\n",
+    "\n",
+    "## Best practices\n",
+    "\n",
+    "1. **Always test the exported model** before deploying \u2014 run inference and\n",
+    "   compare outputs with the original Keras model.\n",
+    "2. **Use the PyTorch backend** for LiteRT export to avoid flex ops and ensure\n",
+    "   compatibility with the new LiteRT Android runtime.\n",
+    "3. **Quantize with `ai-edge-quantizer`** after export for the best size/accuracy\n",
+    "   trade-off, especially for LLMs.\n",
+    "4. **Build subclassed models before exporting** \u2014 call them on sample data so\n",
+    "   Keras knows the input signature.\n",
+    "5. **Keep models under 2 GB per TFLite file** \u2014 LiteRT uses a flatbuffer format\n",
+    "   with a 2 GB limit per file. If your model is larger, use the\n",
+    "   `litert-lm` / `litert-lm-builder` pipeline which shards the model into\n",
+    "   multiple sub-models inside a `.litertlm` container.\n",
+    "\n",
+    "## Troubleshooting\n",
+    "\n",
+    "| Issue | Solution |\n",
+    "|---|---|\n",
+    "| `ImportError` for `ai_edge_litert` | Run `pip install ai-edge-litert` |\n",
+    "| `ImportError` for `litert_torch` | Run `pip install litert-torch` |\n",
+    "| Shape mismatch at inference | Verify the input shape matches what the model expects |\n",
+    "| Subclassed model fails to export | Call the model on sample data before `export()` to build weights |\n",
+    "| Out of memory during export | Try quantization or export on a machine with more RAM |\n",
+    "| Quantization via `optimizations` fails on PyTorch | Use `ai-edge-quantizer` as a post-export step instead |\n",
+    "| Flex ops on Android | Re-export with `KERAS_BACKEND=torch` |"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "None",
+  "colab": {
+   "collapsed_sections": [],
+   "name": "litert_export",
+   "private_outputs": false,
+   "provenance": [],
+   "toc_visible": true
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.7.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}