tflite-micro/tensorflow/lite/micro/compression/mnist_compression_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c388e286",
   "metadata": {},
   "source": [
    "# TFLM Compression Tutorial: MNIST [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tensorflow/tflite-micro/blob/main/tensorflow/lite/micro/compression/mnist_compression_tutorial.ipynb)\n",
    "\n",
    "This tutorial demonstrates TensorFlow Lite for Microcontrollers (TFLM) model compression\n",
    "using the MNIST dataset.\n",
    "\n",
    "## Beyond Standard Quantization: Compression Enabled by Weight Clustering\n",
    "\n",
    "While quantization (converting float32 to INT8) is standard practice in\n",
    "embedded ML and provides a reliable 4x size reduction, TFLM can achieve even\n",
    "better results through compression of weight-clustered models. To unlock this\n",
    "additional compression, you'll need to perform two key steps: weight clustering\n",
    "during model training, followed by LUT-based compression during post-processing.\n",
    "\n",
    "### Why Weight Clustering is Essential\n",
    "\n",
    "Neural network weights are typically unique values spread across a continuous\n",
    "range. This diversity makes them nearly impossible to compress effectively - there's no\n",
    "pattern or redundancy to exploit. Weight clustering changes this fundamentally by:\n",
    "\n",
    "- Grouping similar weights into a limited number of clusters\n",
    "- Replacing diverse weight values with a small set of shared values\n",
    "- Creating the redundancy that compression algorithms need to work\n",
    "\n",
    "Without clustering, compression algorithms struggle because every weight is different.\n",
    "With clustering, we can represent each weight using just an index into a small lookup table.\n",
    "The number of clusters (and thus the compression ratio) is a parameter you control based\n",
    "on your accuracy requirements.\n",
    "\n",
    "## The Complete Pipeline\n",
    "\n",
    "TFLM combines three techniques:\n",
    "\n",
    "1. **Weight Clustering** (the enabler): Groups weights into clusters, creating patterns\n",
    "   that can be compressed. Without this step, meaningful compression isn't possible.\n",
    "\n",
    "2. **Quantization** (standard practice): Converts float32 to INT8, providing the baseline\n",
    "   4x reduction that's common in embedded deployments.\n",
    "\n",
    "3. **Look-Up Table (LUT) Compression** (the payoff): Leverages the clustered structure to\n",
    "   store indices and a lookup table, achieving compression beyond what quantization alone\n",
    "   can provide.\n",
    "\n",
    "## Current Limitations\n",
    "\n",
    "As of today, not every operator in TFLM supports compression. This is actively being\n",
    "improved, and this tutorial will be updated as support expands. Currently, compression\n",
    "works with fully connected and convolutional layers.\n",
    "\n",
    "## What You'll Learn\n",
    "\n",
    "In this tutorial, you'll:\n",
    "- Train a simple CNN model for MNIST digit classification\n",
    "- Apply weight clustering to specific layers using TensorFlow Model Optimization toolkit\n",
    "- Convert and quantize the model for embedded deployment\n",
    "- Apply TFLM's LUT compression for maximum size reduction\n",
    "- Evaluate the tradeoffs between compression ratio and accuracy\n",
    "\n",
    "For more details on weight clustering, refer to:\n",
    "- [TensorFlow Model Optimization Guide](https://www.tensorflow.org/model_optimization/guide/clustering)\n",
    "- [Clustering Example](https://www.tensorflow.org/model_optimization/guide/clustering/clustering_example)\n",
    "- [Comprehensive Clustering Guide](https://www.tensorflow.org/model_optimization/guide/clustering/clustering_comprehensive_guide)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f65d42c8",
   "metadata": {},
   "source": [
    "## Import Required Libraries\n",
    "\n",
    "First, we'll import all the necessary libraries for this tutorial. We need:\n",
    "- TensorFlow for model training\n",
    "- TFLM Python package for compression and embedded inference simulation\n",
    "- TensorFlow Model Optimization toolkit for weight clustering\n",
    "- NumPy and Matplotlib for data handling and visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f252c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tflite-micro\n",
    "import tensorflow as tf\n",
    "import tflite_micro as tflm\n",
    "import numpy as np \n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Import tf_keras for compatibility. This is a standalone Keras implementation that\n",
    "# maintains stable APIs across TensorFlow versions. Using tf_keras instead of tf.keras\n",
    "# ensures that the TensorFlow Model Optimization toolkit (used for weight clustering)\n",
    "# works consistently regardless of the TensorFlow version installed.\n",
    "import tf_keras\n",
    "\n",
    "print(f\"TensorFlow version: {tf.__version__}\")\n",
    "print(f\"tf_keras version: {tf_keras.__version__}\")\n",
    "print(f\"TFLM module: {tflm.__version__}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e5a7885",
   "metadata": {},
   "source": [
    "## Load and Prepare MNIST Dataset\n",
    "\n",
    "MNIST is a classic dataset of handwritten digits (0-9). Each image is 28x28\n",
    "pixels in grayscale. We'll use this simple dataset to show the\n",
    "effects of compression without the complexity of larger models.\n",
    "\n",
    "The data preparation steps are standard:\n",
    "- Normalize pixel values to [0, 1] range.\n",
    "- Add a channel dimension since convolutional layers expect 3D input (height, width, channels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f1a3b28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load MNIST dataset\n",
    "(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()\n",
    "\n",
    "# Define the input shape for MNIST images\n",
    "# This will be used throughout the tutorial for model creation and TFLite conversion\n",
    "# Shape format: [batch_size, height, width, channels]\n",
    "# - None: Variable batch size (TFLite will use batch size 1 at inference)\n",
    "# - 28, 28: MNIST image dimensions\n",
    "# - 1: Grayscale images have 1 channel\n",
    "MNIST_INPUT_SHAPE = [None, 28, 28, 1]\n",
    "\n",
    "# Normalize pixel values to [0, 1]\n",
    "train_images = train_images.astype('float32') / 255.0\n",
    "test_images = test_images.astype('float32') / 255.0\n",
    "\n",
    "# Add channel dimension\n",
    "train_images = train_images[..., np.newaxis]\n",
    "test_images = test_images[..., np.newaxis]\n",
    "\n",
    "print(f\"Training images shape: {train_images.shape}\")\n",
    "print(f\"Training labels shape: {train_labels.shape}\")\n",
    "print(f\"Test images shape: {test_images.shape}\")\n",
    "print(f\"Test labels shape: {test_labels.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8baf2147",
   "metadata": {},
   "source": [
    "## Visualize Sample Images\n",
    "\n",
    "Let's visualize some sample images to understand our dataset. This helps verify\n",
    "that our data is loaded correctly and gives us intuition about the\n",
    "classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb524ab8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display sample images\n",
    "fig, axes = plt.subplots(2, 5, figsize=(12, 6))\n",
    "axes = axes.ravel()\n",
    "\n",
    "for i in range(10):\n",
    "    axes[i].imshow(train_images[i].squeeze(), cmap='gray')\n",
    "    axes[i].set_title(f'Label: {train_labels[i]}')\n",
    "    axes[i].axis('off')\n",
    "\n",
    "plt.suptitle('Sample MNIST Images')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d15ca87",
   "metadata": {},
   "source": [
    "## Create a Model Architecture for Compression Demonstration\n",
    "\n",
    "For this tutorial, we'll use a model that's appropriately sized to demonstrate\n",
    "compression while keeping training time reasonable.\n",
    "\n",
    "TFLM currently supports compression on:\n",
    "- **Conv2D layers**: Convolutional weights can be compressed\n",
    "- **Dense layers**: Converted to FullyConnected operators in TFLite, also compressible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8283209",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a model architecture optimized for demonstrating compression\n",
    "model = tf_keras.Sequential([\n",
    "    tf_keras.layers.Input(shape=MNIST_INPUT_SHAPE[1:]),  # Skip batch dimension\n",
    "    \n",
    "    # Convolutional layers\n",
    "    tf_keras.layers.Conv2D(16, (5, 5), activation='relu'),\n",
    "    tf_keras.layers.MaxPooling2D((2, 2)),\n",
    "    \n",
    "    # Second conv layer\n",
    "    tf_keras.layers.Conv2D(32, (3, 3), activation='relu'),\n",
    "    tf_keras.layers.MaxPooling2D((2, 2)),\n",
    "    \n",
    "    # Flatten for fully connected layers\n",
    "    tf_keras.layers.Flatten(),\n",
    "    \n",
    "    # Fully connected layers\n",
    "    tf_keras.layers.Dense(128, activation='relu'),\n",
    "    tf_keras.layers.Dense(64, activation='relu'),\n",
    "    \n",
    "    # Output layer\n",
    "    tf_keras.layers.Dense(10, activation='softmax')\n",
    "])\n",
    "\n",
    "model.compile(optimizer='adam',\n",
    "              loss='sparse_categorical_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6923fb47",
   "metadata": {},
   "source": [
    "## Train the Model\n",
    "\n",
    "Now we'll train our baseline model. Note that this tutorial focuses on demonstrating\n",
    "compression techniques, not optimal training practices. We're using 3 epochs which\n",
    "is sufficient for MNIST and keeps the tutorial quick.\n",
    "\n",
    "For production training best practices, see:\n",
    "- [TensorFlow Training Guide](https://www.tensorflow.org/guide/keras/training_with_built_in_methods)\n",
    "- [Keras Training Documentation](https://keras.io/guides/training_with_built_in_methods/)\n",
    "\n",
    "Pay attention to the final test accuracy - we'll use this to measure how much\n",
    "accuracy we trade for compression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c0132f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train the model (simplified training for demonstration)\n",
    "history = model.fit(\n",
    "    train_images, train_labels,\n",
    "    epochs=3,\n",
    "    batch_size=128,\n",
    "    validation_split=0.1,\n",
    "    verbose=1\n",
    ")\n",
    "\n",
    "# Evaluate on test set\n",
    "test_loss, test_accuracy = model.evaluate(test_images, test_labels, verbose=0)\n",
    "print(f\"\\nTest accuracy: {test_accuracy:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e303c3e",
   "metadata": {},
   "source": [
    "## Plot Training History\n",
    "\n",
    "Visualizing the training history helps us verify that our model trained properly.\n",
    "We should see decreasing loss and increasing accuracy over epochs, with validation\n",
    "metrics following similar trends (indicating no overfitting)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4101a0f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot training history\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))\n",
    "\n",
    "ax1.plot(history.history['loss'], label='Training Loss')\n",
    "ax1.plot(history.history['val_loss'], label='Validation Loss')\n",
    "ax1.set_xlabel('Epoch')\n",
    "ax1.set_ylabel('Loss')\n",
    "ax1.set_title('Training and Validation Loss')\n",
    "ax1.legend()\n",
    "ax1.grid(True)\n",
    "\n",
    "ax2.plot(history.history['accuracy'], label='Training Accuracy')\n",
    "ax2.plot(history.history['val_accuracy'], label='Validation Accuracy')\n",
    "ax2.set_xlabel('Epoch')\n",
    "ax2.set_ylabel('Accuracy')\n",
    "ax2.set_title('Training and Validation Accuracy')\n",
    "ax2.legend()\n",
    "ax2.grid(True)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86b9344d",
   "metadata": {},
   "source": [
    "## Apply Weight Clustering to Selected Layers\n",
    "\n",
    "Everything we've done so far - loading data, creating a model,\n",
    "training it - is standard deep learning workflow that you'd use for any neural network.\n",
    "From this point forward, we enter compression-specific territory with weight clustering,\n",
    "the critical first step that enables TFLM's compression pipeline.\n",
    "\n",
    "### Why Clustering Enables Compression\n",
    "\n",
    "Without clustering, neural network weights are typically unique floating-point values\n",
    "spread across a continuous range. This makes compression nearly impossible because\n",
    "there's no redundancy to exploit. Weight clustering groups similar weights together,\n",
    "creating redundancy that compression algorithms can leverage.\n",
    "\n",
    "For example, if we cluster weights into 16 groups:\n",
    "- Original: thousands of unique 32-bit float values\n",
    "- After clustering: only 16 unique values\n",
    "- Result: each weight can be represented by a 4-bit index into a 16-entry lookup table\n",
    "\n",
    "### Selective Layer Clustering\n",
    "\n",
    "We don't need to cluster every layer. In fact, selective clustering often yields better\n",
    "results because:\n",
    "- Some layers (like the final classification layer) might be more sensitive to clustering\n",
    "- Different layers might benefit from different numbers of clusters\n",
    "- You can balance model size and accuracy by choosing which layers to compress\n",
    "\n",
    "In this tutorial, we'll cluster only the second Conv2D layer and the\n",
    "first fully-connected layer. This demonstrates targeted compression while\n",
    "maintaining the original precision in other layers.\n",
    "\n",
    "### Implementation Details\n",
    "\n",
    "We use TensorFlow Model Optimization toolkit's clustering API, which:\n",
    "- Replaces weights with cluster centroids during training\n",
    "- Fine-tunes the model to adapt to clustered weights\n",
    "- Maintains cluster assignments while updating centroids\n",
    "\n",
    "For comprehensive documentation on clustering techniques, see:\n",
    "- [TensorFlow Model Optimization Clustering Guide](https://www.tensorflow.org/model_optimization/guide/clustering)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66c68765",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorflow_model_optimization\n",
    "import tensorflow_model_optimization as tfmot\n",
    "\n",
    "# Define clustering parameters\n",
    "cluster_weights = tfmot.clustering.keras.cluster_weights\n",
    "CentroidInitialization = tfmot.clustering.keras.CentroidInitialization\n",
    "\n",
    "print(\"Applying weight clustering to selected layers...\")\n",
    "\n",
    "# First, we need to clone the original model to preserve it\n",
    "model_config = model.get_config()\n",
    "cloned_model = tf_keras.Sequential.from_config(model_config)\n",
    "cloned_model.set_weights(model.get_weights())\n",
    "\n",
    "# Define which layers to cluster and with how many clusters\n",
    "# We'll target the larger layers for maximum compression impact\n",
    "# \n",
    "# Note: We're storing this configuration in CLUSTERING_CONFIG because we'll need\n",
    "# to reference it later when specifying which layers should be compressed during\n",
    "# the TFLite conversion process. This ensures consistency between the layers we\n",
    "# cluster and the layers we compress.\n",
    "CLUSTERING_CONFIG = {\n",
    "    'conv2d_1': 16,  # Second conv layer\n",
    "    'dense': 16,     # First dense layer\n",
    "}\n",
    "\n",
    "# Important: The cluster_weights() function doesn't modify layers in-place.\n",
    "# Instead, it returns a wrapper layer that adds clustering functionality.\n",
    "# Therefore, we need to:\n",
    "# 1. Walk through each layer\n",
    "# 2. Wrap layers that need clustering\n",
    "# 3. Collect all layers (wrapped and unwrapped)\n",
    "# 4. Build a new Sequential model from the collected layers\n",
    "\n",
    "clustered_layers = []\n",
    "for layer in cloned_model.layers:\n",
    "    if layer.name in CLUSTERING_CONFIG:\n",
    "        # Wrap this layer with clustering functionality\n",
    "        num_clusters = CLUSTERING_CONFIG[layer.name]\n",
    "        clustered_layer = cluster_weights(\n",
    "            layer,\n",
    "            number_of_clusters=num_clusters,\n",
    "            cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS\n",
    "        )\n",
    "        clustered_layers.append(clustered_layer)\n",
    "        print(f\"Applied {num_clusters} clusters to layer: {layer.name}\")\n",
    "    else:\n",
    "        # Keep the original layer unchanged\n",
    "        clustered_layers.append(layer)\n",
    "\n",
    "# Create a new Sequential model from our list of layers\n",
    "# (This is necessary because Sequential models require rebuilding when modifying layers)\n",
    "clustered_model = tf_keras.Sequential(clustered_layers)\n",
    "\n",
    "print(f\"Clustered model created\")\n",
    "print(\"\\nClustered layers:\")\n",
    "for layer in clustered_model.layers:\n",
    "    if 'cluster' in str(type(layer)).lower():\n",
    "        print(f\"  - {layer.name}: {type(layer).__name__}\")\n",
    "    else:\n",
    "        print(f\"  - {layer.name}: not clustered\")\n",
    "\n",
    "# Compile the clustered model\n",
    "clustered_model.compile(\n",
    "    optimizer='adam',\n",
    "    loss='sparse_categorical_crossentropy',\n",
    "    metrics=['accuracy']\n",
    ")\n",
    "\n",
    "# Build the model by calling it on dummy data\n",
    "clustered_model.build(input_shape=MNIST_INPUT_SHAPE)\n",
    "\n",
    "print(\"Clustered model summary:\")\n",
    "clustered_model.summary()\n",
    "\n",
    "# Fine-tune the clustered model\n",
    "# Note: We use 2 epochs for demonstration. For optimal fine-tuning strategies, see:\n",
    "# https://www.tensorflow.org/model_optimization/guide/clustering/clustering_comprehensive_guide\n",
    "print(\"\\nFine-tuning clustered model...\")\n",
    "clustered_history = clustered_model.fit(\n",
    "    train_images, train_labels,\n",
    "    epochs=2,  # Simplified fine-tuning for demonstration\n",
    "    batch_size=128,\n",
    "    validation_split=0.1,\n",
    "    verbose=1\n",
    ")\n",
    "\n",
    "# Evaluate clustered model\n",
    "clustered_test_loss, clustered_test_accuracy = clustered_model.evaluate(test_images, test_labels, verbose=0)\n",
    "print(f\"\\nClustered model test accuracy: {clustered_test_accuracy:.4f}\")\n",
    "\n",
    "# Strip clustering wrappers to get the final model\n",
    "final_clustered_model = tfmot.clustering.keras.strip_clustering(clustered_model)\n",
    "\n",
    "# Compare model sizes before conversion\n",
    "print(\"\\nModel comparison:\")\n",
    "print(f\"Original model parameters: {model.count_params():,}\")\n",
    "print(f\"Clustered model parameters: {final_clustered_model.count_params():,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c964fa8",
   "metadata": {},
   "source": [
    "## Convert Models to TFLite\n",
    "\n",
    "Now we convert our models to TensorFlow Lite format, required by the TFLM\n",
    "interpreter. We follow a two-step approach that's standard in TensorFlow Lite\n",
    "workflows:\n",
    "\n",
    "### Why Convert to Float First, Then Quantize?\n",
    "\n",
    "You might wonder why we don't quantize directly during conversion. This two-step \n",
    "approach (float → quantized) is actually the preferred workflow for several reasons:\n",
    "\n",
    "1. **Debugging and Validation**: The float model serves as a baseline for accuracy\n",
    "   comparison. You can verify the conversion works correctly before adding quantization\n",
    "   complexity.\n",
    "\n",
    "2. **Flexibility**: You can experiment with different quantization strategies on the\n",
    "   same float model without re-converting from Keras/TensorFlow each time.\n",
    "\n",
    "3. **Gradual Optimization**: This approach lets you measure the impact of each\n",
    "   optimization step separately - first conversion, then quantization, then compression.\n",
    "\n",
    "### What This Step Accomplishes\n",
    "\n",
    "Converting to TFLite format:\n",
    "- Provides a standardized model representation that TFLM can work with\n",
    "- Optimizes the model graph for mobile and embedded deployment\n",
    "- Allows us to measure the baseline model size\n",
    "\n",
    "At this stage, you'll notice that clustering alone hasn't significantly reduced\n",
    "the model size. This is expected! The real size reduction comes when we combine\n",
    "clustering with quantization and TFLM's LUT compression in the following steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06ff9f70",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the original model to TFLite format using concrete function\n",
    "#\n",
    "# Why use get_concrete_function()?\n",
    "# - Using concrete functions ensures the resulting model will work with TFLM\n",
    "# - Keras models are high-level abstractions that support multiple input shapes\n",
    "# - TFLM requires all tensor shapes to be known at compile time (no dynamic tensors)\n",
    "# - When converting Keras models directly via convert_from_keras(), the TFLite\n",
    "#   converter sometimes creates models with dynamic tensors that are\n",
    "#   incompatible with TFLM\n",
    "#\n",
    "# The conversion process:\n",
    "# 1. Wrap the Keras model in tf.function to create a TensorFlow graph\n",
    "# 2. Call get_concrete_function() with a TensorSpec that defines the exact input shape\n",
    "# 3. This creates a \"concrete function\"---a graph with all shapes and types determined\n",
    "# 4. TFLite converter can then convert this concrete graph\n",
    "concrete_func = tf.function(lambda x: model(x)).get_concrete_function(\n",
    "    tf.TensorSpec(shape=MNIST_INPUT_SHAPE, dtype=tf.float32)\n",
    ")\n",
    "converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])\n",
    "tflite_model = converter.convert()\n",
    "\n",
    "# Save the model\n",
    "with open('mnist_model.tflite', 'wb') as f:\n",
    "    f.write(tflite_model)\n",
    "\n",
    "print(f\"Original model size: {len(tflite_model):,} bytes ({len(tflite_model)/1024:.2f} KB)\")\n",
    "\n",
    "# Convert the clustered model to TFLite format using concrete function\n",
    "# Same process as above---we need to create a concrete function with fixed input shapes\n",
    "clustered_concrete_func = tf.function(lambda x: final_clustered_model(x)).get_concrete_function(\n",
    "    tf.TensorSpec(shape=MNIST_INPUT_SHAPE, dtype=tf.float32)\n",
    ")\n",
    "converter_clustered = tf.lite.TFLiteConverter.from_concrete_functions([clustered_concrete_func])\n",
    "tflite_clustered_model = converter_clustered.convert()\n",
    "\n",
    "# Save the clustered model\n",
    "with open('mnist_model_clustered.tflite', 'wb') as f:\n",
    "    f.write(tflite_clustered_model)\n",
    "\n",
    "print(f\"Clustered model size: {len(tflite_clustered_model):,} bytes ({len(tflite_clustered_model)/1024:.2f} KB)\")\n",
    "print(f\"Size reduction from clustering: {(1 - len(tflite_clustered_model)/len(tflite_model))*100:.1f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3fe6099",
   "metadata": {},
   "source": [
    "## Apply Post-Training Quantization\n",
    "\n",
    "Next, we apply post-training quantization to convert our models from float32 to INT8.\n",
    "This is standard practice in TensorFlow Lite deployment - virtually all production\n",
    "models use quantization for the benefits it provides:\n",
    "\n",
    "- **4x size reduction**: 32-bit floats → 8-bit integers\n",
    "- **Faster inference**: Integer operations are more efficient on microcontrollers\n",
    "- **Lower power consumption**: Integer math requires less energy\n",
    "\n",
    "### This is Not the Novel Part\n",
    "\n",
    "To be clear: quantization is a form of compression, but it's not what makes this\n",
    "tutorial special. INT8 quantization is routine in embedded ML deployment. What's\n",
    "unique here is how our earlier clustering step sets up the model for TFLM's\n",
    "LUT-based compression in the next phase.\n",
    "\n",
    "### Quantization's Role in Our Pipeline\n",
    "\n",
    "While quantization doesn't enable compression per se, it plays an important role\n",
    "in our pipeline:\n",
    "- TFLM's LUT compression currently operates on integer models\n",
    "- Clustered weights naturally map well to INT8 representation\n",
    "- The overall pipeline (clustering → quantization → LUT compression) works together\n",
    "  to achieve the final size reduction\n",
    "\n",
    "So while this quantization step is standard practice, it's a necessary part of\n",
    "the complete compression workflow."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edd87d05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Provide representative dataset for full integer quantization\n",
    "def representative_dataset():\n",
    "    for i in range(100):\n",
    "        # Get a random batch of input data\n",
    "        indices = np.random.randint(0, len(train_images), size=1)\n",
    "        yield [train_images[indices].astype(np.float32)]\n",
    "\n",
    "# Convert original model with post-training quantization using concrete function\n",
    "# Even though we're quantizing, the concrete function still expects float32 input\n",
    "# The quantization happens during conversion, not in the concrete function definition\n",
    "quant_concrete_func = tf.function(lambda x: model(x)).get_concrete_function(\n",
    "    tf.TensorSpec(shape=MNIST_INPUT_SHAPE, dtype=tf.float32)\n",
    ")\n",
    "converter = tf.lite.TFLiteConverter.from_concrete_functions([quant_concrete_func])\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",
    "converter.inference_input_type = tf.uint8\n",
    "converter.inference_output_type = tf.uint8\n",
    "\n",
    "quantized_model = converter.convert()\n",
    "\n",
    "# Save the quantized model\n",
    "with open('mnist_model_quantized.tflite', 'wb') as f:\n",
    "    f.write(quantized_model)\n",
    "\n",
    "print(\"Original model with quantization:\")\n",
    "print(f\"  Size: {len(quantized_model):,} bytes ({len(quantized_model)/1024:.2f} KB)\")\n",
    "print(f\"  Compression ratio: {len(tflite_model) / len(quantized_model):.2f}x\")\n",
    "print(f\"  Size reduction: {(1 - len(quantized_model)/len(tflite_model))*100:.1f}%\")\n",
    "\n",
    "# Convert clustered model with post-training quantization using concrete function\n",
    "clustered_quant_concrete_func = tf.function(lambda x: final_clustered_model(x)).get_concrete_function(\n",
    "    tf.TensorSpec(shape=MNIST_INPUT_SHAPE, dtype=tf.float32)\n",
    ")\n",
    "converter_clustered_quant = tf.lite.TFLiteConverter.from_concrete_functions([clustered_quant_concrete_func])\n",
    "converter_clustered_quant.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "converter_clustered_quant.representative_dataset = representative_dataset\n",
    "converter_clustered_quant.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]\n",
    "converter_clustered_quant.inference_input_type = tf.uint8\n",
    "converter_clustered_quant.inference_output_type = tf.uint8\n",
    "\n",
    "clustered_quantized_model = converter_clustered_quant.convert()\n",
    "\n",
    "# Save the clustered+quantized model\n",
    "with open('mnist_model_clustered_quantized.tflite', 'wb') as f:\n",
    "    f.write(clustered_quantized_model)\n",
    "\n",
    "print(\"\\nClustered model with quantization:\")\n",
    "print(f\"  Size: {len(clustered_quantized_model):,} bytes ({len(clustered_quantized_model)/1024:.2f} KB)\")\n",
    "print(f\"  Compression ratio: {len(tflite_model) / len(clustered_quantized_model):.2f}x\")\n",
    "print(f\"  Size reduction: {(1 - len(clustered_quantized_model)/len(tflite_model))*100:.1f}%\")\n",
    "\n",
    "# Summary of all model sizes\n",
    "print(\"\\n\" + \"=\"*50)\n",
    "print(\"MODEL SIZE SUMMARY:\")\n",
    "print(\"=\"*50)\n",
    "print(f\"Original float32 model:         {len(tflite_model):,} bytes\")\n",
    "print(f\"Clustered float32 model:        {len(tflite_clustered_model):,} bytes\")\n",
    "print(f\"Quantized int8 model:           {len(quantized_model):,} bytes\")\n",
    "print(f\"Clustered + Quantized model:    {len(clustered_quantized_model):,} bytes\")\n",
    "print(f\"\\nBest compression ratio: {len(tflite_model) / len(clustered_quantized_model):.2f}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71829b9f",
   "metadata": {},
   "source": [
    "## Inspect Model Architecture and Data Types\n",
    "\n",
    "In production workflows, you might use sophisticated tools to visualize model graphs,\n",
    "such as Netron (interactive model visualizer), TensorBoard, or TFLite's built-in\n",
    "HTML visualization. However, for this tutorial, simple text-based output provides\n",
    "everything we need to understand our small MNIST model. The text output clearly shows\n",
    "which layers have been quantized, their quantization parameters, and tensor\n",
    "shapes - sufficient for verifying that quantization worked correctly without\n",
    "additional dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2a85615",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_model_details(model_content, model_name=\"Model\"):\n",
    "    \"\"\"Print detailed information about a TFLite model including data types.\"\"\"\n",
    "    # For model inspection, we'll use TFLite since TFLM doesn't expose tensor details\n",
    "    # TFLM is designed for inference, not model introspection\n",
    "    interpreter = tf.lite.Interpreter(model_content=model_content)\n",
    "    interpreter.allocate_tensors()\n",
    "    \n",
    "    print(f\"\\n{model_name} Details:\")\n",
    "    print(\"=\"*50)\n",
    "    \n",
    "    # Get tensor details\n",
    "    tensor_details = interpreter.get_tensor_details()\n",
    "    \n",
    "    # Print input details\n",
    "    input_details = interpreter.get_input_details()\n",
    "    print(\"\\nINPUT TENSORS:\")\n",
    "    for inp in input_details:\n",
    "        tensor = tensor_details[inp['index']]\n",
    "        print(f\"  Name: {tensor['name']}\")\n",
    "        print(f\"  Shape: {tensor['shape']}\")\n",
    "        print(f\"  Type: {tensor['dtype']}\")\n",
    "        if 'quantization' in tensor and tensor['quantization'][0] != 0:\n",
    "            print(f\"  Quantization: scale={tensor['quantization'][0]:.6f}, zero_point={tensor['quantization'][1]}\")\n",
    "        print()\n",
    "    \n",
    "    # Print output details\n",
    "    output_details = interpreter.get_output_details()\n",
    "    print(\"OUTPUT TENSORS:\")\n",
    "    for out in output_details:\n",
    "        tensor = tensor_details[out['index']]\n",
    "        print(f\"  Name: {tensor['name']}\")\n",
    "        print(f\"  Shape: {tensor['shape']}\")\n",
    "        print(f\"  Type: {tensor['dtype']}\")\n",
    "        if 'quantization' in tensor and tensor['quantization'][0] != 0:\n",
    "            print(f\"  Quantization: scale={tensor['quantization'][0]:.6f}, zero_point={tensor['quantization'][1]}\")\n",
    "        print()\n",
    "    \n",
    "    # Print layer information\n",
    "    print(\"ALL LAYERS:\")\n",
    "    print(f\"{'Index':<6} {'Name':<40} {'Type':<10} {'Shape':<20} {'Quantization':<30}\")\n",
    "    print(\"-\"*110)\n",
    "    \n",
    "    for i, tensor in enumerate(tensor_details):\n",
    "        quant_info = \"\"\n",
    "        if 'quantization' in tensor and tensor['quantization'][0] != 0:\n",
    "            quant_info = f\"s={tensor['quantization'][0]:.4f}, zp={tensor['quantization'][1]}\"\n",
    "        \n",
    "        print(f\"{i:<6} {tensor['name'][:39]:<40} {str(tensor['dtype']):<10} {str(tensor['shape']):<20} {quant_info:<30}\")\n",
    "\n",
    "# Print details for both models\n",
    "print_model_details(tflite_model, \"Original Float32 Model\")\n",
    "print_model_details(quantized_model, \"Quantized INT8 Model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15e72732",
   "metadata": {},
   "source": [
    "## Compare Model Accuracy\n",
    "\n",
    "Now we evaluate how compression techniques affect model accuracy. This is a critical\n",
    "step because it quantifies the accuracy-size tradeoff.\n",
    "\n",
    "We use TFLM's Python interpreter to simulate how these models will perform on\n",
    "microcontrollers. This gives us confidence that our accuracy measurements reflect\n",
    "real-world embedded performance.\n",
    "\n",
    "Key points to observe:\n",
    "- Quantization typically causes a small accuracy drop\n",
    "- Clustering might add minimal additional accuracy loss\n",
    "- The combined techniques usually maintain acceptable accuracy\n",
    "- If accuracy drops too much, consider using more clusters or compressing fewer layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5746219",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helper function to evaluate TFLite model using TFLM\n",
    "def evaluate_tflite_model(model_content, test_images, test_labels):\n",
    "    # Use TFLM interpreter\n",
    "    interpreter = tflm.runtime.Interpreter.from_bytes(bytes(model_content))\n",
    "    \n",
    "    # TFLM uses different API - get details for index 0\n",
    "    INPUT_INDEX = 0\n",
    "    OUTPUT_INDEX = 0\n",
    "    input_details = interpreter.get_input_details(INPUT_INDEX)\n",
    "    output_details = interpreter.get_output_details(OUTPUT_INDEX)\n",
    "    \n",
    "    correct = 0\n",
    "    total = min(1000, len(test_images))  # Evaluate on subset for speed\n",
    "    \n",
    "    for i in range(total):\n",
    "        # Prepare input\n",
    "        test_image = test_images[i:i+1]\n",
    "        \n",
    "        # Quantize input if needed\n",
    "        if input_details['dtype'] == np.uint8:\n",
    "            quant_params = input_details.get('quantization_parameters', {})\n",
    "            if 'scales' in quant_params and 'zero_points' in quant_params:\n",
    "                input_scale = quant_params['scales'][0]\n",
    "                input_zero_point = quant_params['zero_points'][0]\n",
    "                test_image = test_image / input_scale + input_zero_point\n",
    "                test_image = test_image.astype(np.uint8)\n",
    "        \n",
    "        # Run inference using TFLM API\n",
    "        interpreter.set_input(test_image, INPUT_INDEX)\n",
    "        interpreter.invoke()\n",
    "        \n",
    "        # Get output\n",
    "        output = interpreter.get_output(OUTPUT_INDEX)[0]\n",
    "        \n",
    "        # Dequantize output if needed\n",
    "        if output_details['dtype'] == np.uint8:\n",
    "            quant_params = output_details.get('quantization_parameters', {})\n",
    "            if 'scales' in quant_params and 'zero_points' in quant_params:\n",
    "                output_scale = quant_params['scales'][0]\n",
    "                output_zero_point = quant_params['zero_points'][0]\n",
    "                output = (output.astype(np.float32) - output_zero_point) * output_scale\n",
    "        \n",
    "        predicted = np.argmax(output)\n",
    "        if predicted == test_labels[i]:\n",
    "            correct += 1\n",
    "    \n",
    "    return correct / total\n",
    "\n",
    "# Evaluate all models\n",
    "print(\"Evaluating models on test subset...\")\n",
    "original_accuracy = evaluate_tflite_model(tflite_model, test_images, test_labels)\n",
    "clustered_accuracy = evaluate_tflite_model(tflite_clustered_model, test_images, test_labels)\n",
    "quantized_accuracy = evaluate_tflite_model(quantized_model, test_images, test_labels)\n",
    "clustered_quantized_accuracy = evaluate_tflite_model(clustered_quantized_model, test_images, test_labels)\n",
    "\n",
    "print(\"\\n\" + \"=\"*50)\n",
    "print(\"ACCURACY COMPARISON:\")\n",
    "print(\"=\"*50)\n",
    "print(f\"Original float32 model:      {original_accuracy:.4f}\")\n",
    "print(f\"Clustered float32 model:     {clustered_accuracy:.4f}\")\n",
    "print(f\"Quantized int8 model:        {quantized_accuracy:.4f}\")\n",
    "print(f\"Clustered + Quantized model: {clustered_quantized_accuracy:.4f}\")\n",
    "print(f\"\\nAccuracy drop from clustering+quantization: {(original_accuracy - clustered_quantized_accuracy)*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37ac8a47",
   "metadata": {},
   "source": [
    "## Apply TFLM Compression to the Clustered Layer\n",
    "\n",
    "This is where the magic happens! TFLM's Look-Up Table (LUT) compression leverages\n",
    "the clustering we applied earlier to achieve further size reduction.\n",
    "\n",
    "### How LUT Compression Works\n",
    "\n",
    "Remember our clustered layers have only 16 unique weight values. Instead of storing\n",
    "each weight as an 8-bit integer (after quantization), TFLM can:\n",
    "1. Store a table of the 16 unique values (the \"look-up table\")\n",
    "2. Replace each weight with a 4-bit index into this table\n",
    "3. Result: 2x additional compression (8 bits → 4 bits per weight)\n",
    "\n",
    "### The Process\n",
    "\n",
    "To apply TFLM compression, we need to:\n",
    "1. Identify the specific tensor (layer weights) to compress\n",
    "2. Create a compression specification\n",
    "3. Apply the compression transformation\n",
    "\n",
    "### Finding the Right Tensor\n",
    "\n",
    "TFLite models store weights as separate tensors connected to operations. We need\n",
    "to trace through the model graph to find the weight tensor for our target layer.\n",
    "This requires some understanding of TFLite's model structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a187f65f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import TFLM compression module\n",
    "from tflite_micro import compression\n",
    "from tensorflow.lite.python import schema_py_generated as schema_fb\n",
    "\n",
    "def find_weight_tensor_for_layer(model_bytes, layer_name):\n",
    "    \"\"\"Find weight tensor for a specific layer by following the operation graph.\n",
    "    \n",
    "    This works by:\n",
    "    1. Finding the operation that contains the layer name in its output\n",
    "    2. Following the operation's inputs to find the weight tensor\n",
    "    \"\"\"\n",
    "    buf = bytearray(model_bytes)\n",
    "    model = schema_fb.Model.GetRootAsModel(buf, 0)\n",
    "    subgraph = model.Subgraphs(0)\n",
    "    \n",
    "    # Find the operation for this layer\n",
    "    for i in range(subgraph.OperatorsLength()):\n",
    "        op = subgraph.Operators(i)\n",
    "        \n",
    "        # Check output tensor names\n",
    "        for output_idx in range(op.OutputsLength()):\n",
    "            tensor_idx = op.Outputs(output_idx)\n",
    "            tensor = subgraph.Tensors(tensor_idx)\n",
    "            tensor_name = tensor.Name().decode('utf-8')\n",
    "            \n",
    "            if layer_name in tensor_name and ('Conv2D' in tensor_name or 'MatMul' in tensor_name):\n",
    "                # Found the operation! Now get its weight input\n",
    "                # For Conv2D and Dense layers: input[0]=activation, input[1]=weights, input[2]=bias\n",
    "                if op.InputsLength() >= 2:\n",
    "                    weight_tensor_idx = op.Inputs(1)\n",
    "                    weight_tensor = subgraph.Tensors(weight_tensor_idx)\n",
    "                    \n",
    "                    # Get tensor info\n",
    "                    shape = [weight_tensor.Shape(i) for i in range(weight_tensor.ShapeLength())]\n",
    "                    name = weight_tensor.Name().decode('utf-8')\n",
    "                    \n",
    "                    return {\n",
    "                        'index': weight_tensor_idx,\n",
    "                        'name': name,\n",
    "                        'shape': shape\n",
    "                    }\n",
    "    \n",
    "    return None\n",
    "\n",
    "# Find tensors to compress for all clustered layers\n",
    "print(\"Finding weight tensors for compressed layers...\")\n",
    "tensors_to_compress = []\n",
    "\n",
    "for layer_name in CLUSTERING_CONFIG.keys():\n",
    "    weight_info = find_weight_tensor_for_layer(clustered_quantized_model, layer_name)\n",
    "    \n",
    "    if weight_info:\n",
    "        print(f\"\\nFound {layer_name} weight tensor:\")\n",
    "        print(f\"  Index: {weight_info['index']}\")\n",
    "        print(f\"  Name: '{weight_info['name']}'\")\n",
    "        print(f\"  Shape: {weight_info['shape']}\")\n",
    "        tensors_to_compress.append(weight_info)\n",
    "    else:\n",
    "        print(f\"Warning: Could not find weight tensor for layer {layer_name}\")\n",
    "\n",
    "if not tensors_to_compress:\n",
    "    raise ValueError(\"Could not find any weight tensors to compress\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faff6961",
   "metadata": {},
   "source": [
    "## Create and Apply Compression Specification\n",
    "\n",
    "Now we create the compression specification that tells TFLM exactly how to compress\n",
    "our model. TFLM provides two ways to specify compression:\n",
    "\n",
    "### Programmatic API (Used Here)\n",
    "The SpecBuilder API allows us to programmatically define compression:\n",
    "- `add_tensor()`: Specifies which tensor to compress\n",
    "- `with_lut()`: Configures Look-Up Table compression\n",
    "- `index_bitwidth=4`: Uses 4-bit indices (perfect for our 16 clusters)\n",
    "\n",
    "### YAML Configuration (Alternative)\n",
    "\n",
    "In production workflows, the clustering step might be performed by a separate tool\n",
    "or pipeline that analyzes which tensors to compress and with what parameters. That\n",
    "tool would then write a YAML specification file describing the compression strategy.\n",
    "This decouples the compression specification from the compression execution:\n",
    "\n",
    "```yaml\n",
    "tensors:\n",
    "  - subgraph: 0\n",
    "    tensor: 7  # conv2d_1 weights\n",
    "    compression:\n",
    "      lut:\n",
    "        index_bitwidth: 4\n",
    "  - subgraph: 0\n",
    "    tensor: 9  # dense weights\n",
    "    compression:\n",
    "      lut:\n",
    "        index_bitwidth: 4\n",
    "```\n",
    "\n",
    "### Compression Impact\n",
    "\n",
    "The overall compression ratio depends on several factors:\n",
    "\n",
    "1. **Number of layers compressed**: The more layers you compress, the closer you \n",
    "   approach the theoretical maximum compression ratio. For example, with INT8 \n",
    "   quantization (8-bit values) and 16 clusters (4-bit indices), you can approach \n",
    "   50% compression (8 bits → 4 bits) as more layers are compressed.\n",
    "\n",
    "2. **Size of compressed layers**: Larger layers contribute more to the overall \n",
    "   compression. Compressing a 256-unit dense layer has more impact than compressing \n",
    "   a small convolutional layer.\n",
    "\n",
    "3. **Lookup table overhead**: Each compressed tensor needs a lookup table (16 × 8 \n",
    "   bits = 128 bits for 16 clusters). This overhead is negligible for large tensors \n",
    "   but can be significant for very small ones.\n",
    "\n",
    "4. **Model architecture**: Models with many large fully-connected or convolutional \n",
    "   layers benefit more from compression than models dominated by small layers or \n",
    "   non-compressible operations.\n",
    "\n",
    "In this tutorial, we're compressing only 2 out of 8 layers, so our overall \n",
    "compression ratio is less than the theoretical maximum. If we compressed all \n",
    "eligible layers, we could achieve closer to 2x compression on top of quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "231f6016",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\nCreating compression specification...\")\n",
    "# Build compression spec for all identified tensors\n",
    "spec_builder = compression.SpecBuilder()\n",
    "\n",
    "for tensor_info in tensors_to_compress:\n",
    "    (spec_builder\n",
    "         .add_tensor(subgraph=0, tensor=tensor_info['index'])\n",
    "             .with_lut(index_bitwidth=4))\n",
    "    print(f\"Added tensor {tensor_info['index']} ({tensor_info['name']}) to compression spec\")\n",
    "\n",
    "compression_spec = spec_builder.build()\n",
    "\n",
    "print(f\"\\nCompression spec created for {len(tensors_to_compress)} tensors\")\n",
    "print(\"  Compression type: Look-Up Table (LUT)\")\n",
    "print(\"  Index bitwidth: 4 bits (supports up to 16 unique values)\")\n",
    "\n",
    "# Apply TFLM compression\n",
    "print(\"\\nApplying TFLM compression...\")\n",
    "try:\n",
    "    compressed_model = compression.compress(clustered_quantized_model, compression_spec)\n",
    "    \n",
    "    # Save the compressed model\n",
    "    with open('mnist_model_tflm_compressed.tflite', 'wb') as f:\n",
    "        f.write(compressed_model)\n",
    "    \n",
    "    print(f\"\\nTFLM compressed model size: {len(compressed_model):,} bytes ({len(compressed_model)/1024:.2f} KB)\")\n",
    "    print(f\"Additional compression from TFLM: {(1 - len(compressed_model)/len(clustered_quantized_model))*100:.1f}%\")\n",
    "    print(f\"Total compression vs original: {len(tflite_model) / len(compressed_model):.2f}x\")\n",
    "    \n",
    "    # Update the size summary\n",
    "    print(\"\\n\" + \"=\"*50)\n",
    "    print(\"UPDATED MODEL SIZE SUMMARY:\")\n",
    "    print(\"=\"*50)\n",
    "    print(f\"Original float32 model:         {len(tflite_model):,} bytes\")\n",
    "    print(f\"Clustered float32 model:        {len(tflite_clustered_model):,} bytes\")\n",
    "    print(f\"Quantized int8 model:           {len(quantized_model):,} bytes\")\n",
    "    print(f\"Clustered + Quantized model:    {len(clustered_quantized_model):,} bytes\")\n",
    "    print(f\"TFLM Compressed model:          {len(compressed_model):,} bytes\")\n",
    "    print(f\"\\nBest compression ratio: {len(tflite_model) / len(compressed_model):.2f}x\")\n",
    "    \n",
    "except Exception as e:\n",
    "    print(f\"Compression failed: {e}\")\n",
    "    print(\"This might happen if the tensor doesn't have enough unique values for effective compression\")\n",
    "    compressed_model = None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6bf10c6",
   "metadata": {},
   "source": [
    "## Verify TFLM Compressed Model Performance\n",
    "\n",
    "The final step is to verify that our compressed model still performs well. This is\n",
    "crucial because we want to ensure that the compression didn't degrade the model's\n",
    "functionality.\n",
    "\n",
    "### What to Look For\n",
    "- The compressed model should maintain similar accuracy to the clustered+quantized version\n",
    "- Any additional accuracy drop should be minimal\n",
    "- If accuracy drops significantly, consider:\n",
    "  - Using more clusters (e.g., 32 instead of 16)\n",
    "  - Compressing different or fewer layers\n",
    "  - Fine-tuning the clustered model for more epochs\n",
    "\n",
    "### Deployment Considerations\n",
    "Remember that this compressed model will run on microcontrollers using TFLM's C++\n",
    "inference engine, which includes support for decompressing LUT-compressed tensors during\n",
    "inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebb74f1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "if compressed_model is not None:\n",
    "    # Evaluate the TFLM compressed model\n",
    "    print(\"Evaluating TFLM compressed model...\")\n",
    "    tflm_compressed_accuracy = evaluate_tflite_model(compressed_model, test_images, test_labels)\n",
    "    \n",
    "    print(\"\\n\" + \"=\"*50)\n",
    "    print(\"FINAL ACCURACY COMPARISON:\")\n",
    "    print(\"=\"*50)\n",
    "    print(f\"Original float32 model:      {original_accuracy:.4f}\")\n",
    "    print(f\"Clustered float32 model:     {clustered_accuracy:.4f}\")\n",
    "    print(f\"Quantized int8 model:        {quantized_accuracy:.4f}\")\n",
    "    print(f\"Clustered + Quantized model: {clustered_quantized_accuracy:.4f}\")\n",
    "    print(f\"TFLM Compressed model:       {tflm_compressed_accuracy:.4f}\")\n",
    "    print(f\"\\nAccuracy drop from full compression: {(original_accuracy - tflm_compressed_accuracy)*100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1ca2117",
   "metadata": {},
   "source": [
    "## Visualize Predictions\n",
    "\n",
    "Let's visualize some actual predictions to see how our compressed model performs\n",
    "compared to the original."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "550b9836",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare predictions from both models\n",
    "def get_predictions(model_content, images):\n",
    "    # Use TFLM interpreter\n",
    "    interpreter = tflm.runtime.Interpreter.from_bytes(bytes(model_content))\n",
    "    \n",
    "    # TFLM uses different API\n",
    "    INPUT_INDEX = 0\n",
    "    OUTPUT_INDEX = 0\n",
    "    input_details = interpreter.get_input_details(INPUT_INDEX)\n",
    "    output_details = interpreter.get_output_details(OUTPUT_INDEX)\n",
    "    \n",
    "    predictions = []\n",
    "    \n",
    "    for img in images:\n",
    "        test_image = img[np.newaxis, ...]\n",
    "        \n",
    "        if input_details['dtype'] == np.uint8:\n",
    "            quant_params = input_details.get('quantization_parameters', {})\n",
    "            if 'scales' in quant_params and 'zero_points' in quant_params:\n",
    "                input_scale = quant_params['scales'][0]\n",
    "                input_zero_point = quant_params['zero_points'][0]\n",
    "                test_image = test_image / input_scale + input_zero_point\n",
    "                test_image = test_image.astype(np.uint8)\n",
    "        \n",
    "        interpreter.set_input(test_image, INPUT_INDEX)\n",
    "        interpreter.invoke()\n",
    "        \n",
    "        output = interpreter.get_output(OUTPUT_INDEX)[0]\n",
    "        \n",
    "        if output_details['dtype'] == np.uint8:\n",
    "            quant_params = output_details.get('quantization_parameters', {})\n",
    "            if 'scales' in quant_params and 'zero_points' in quant_params:\n",
    "                output_scale = quant_params['scales'][0]\n",
    "                output_zero_point = quant_params['zero_points'][0]\n",
    "                output = (output.astype(np.float32) - output_zero_point) * output_scale\n",
    "        \n",
    "        predictions.append(np.argmax(output))\n",
    "    \n",
    "    return predictions\n",
    "\n",
    "# Get sample images\n",
    "sample_indices = np.random.choice(len(test_images), 10, replace=False)\n",
    "sample_images = test_images[sample_indices]\n",
    "sample_labels = test_labels[sample_indices]\n",
    "\n",
    "# Get predictions from original and compressed models\n",
    "original_preds = get_predictions(tflite_model, sample_images)\n",
    "compressed_preds = get_predictions(compressed_model, sample_images)\n",
    "\n",
    "# Visualize\n",
    "fig, axes = plt.subplots(2, 5, figsize=(15, 8))\n",
    "axes = axes.ravel()\n",
    "\n",
    "for i in range(10):\n",
    "    axes[i].imshow(sample_images[i].squeeze(), cmap='gray')\n",
    "    axes[i].set_title(f'True: {sample_labels[i]}\\nOrig: {original_preds[i]}, Compressed: {compressed_preds[i]}',\n",
    "                      color='green' if original_preds[i] == compressed_preds[i] == sample_labels[i] else 'red')\n",
    "    axes[i].axis('off')\n",
    "\n",
    "plt.suptitle('Original vs TFLM Compressed Model Predictions\\n(Green = all correct, Red = at least one incorrect)')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ebdf03e",
   "metadata": {},
   "source": [
    "## Summary and Key Takeaways\n",
    "\n",
    "Congratulations! You've successfully compressed a neural network for microcontroller\n",
    "deployment using TFLM's compression techniques.\n",
    "\n",
    "### What We Accomplished\n",
    "\n",
    "Through this tutorial, you learned the complete compression pipeline:\n",
    "\n",
    "1. **Baseline Model Creation**: We built a simple CNN for MNIST classification,\n",
    "   achieving good accuracy while keeping the architecture microcontroller-friendly.\n",
    "\n",
    "2. **Selective Weight Clustering**: We applied clustering to only the second Conv2D\n",
    "   layer, demonstrating how to balance compression and accuracy by targeting specific\n",
    "   layers. The 16 clusters prepared this layer for 4-bit compression.\n",
    "\n",
    "3. **Quantization**: We converted the model from float32 to INT8, achieving 4x size\n",
    "   reduction while maintaining inference accuracy.\n",
    "\n",
    "4. **TFLM LUT Compression**: We applied Look-Up Table compression to the clustered\n",
    "   layers, storing weights as 4-bit indices into tables of 16 values, achieving\n",
    "   additional 2x compression for those layers.\n",
    "\n",
    "### Key Insights\n",
    "\n",
    "- **Clustering is Essential**: Without clustering, weights are too diverse to compress\n",
    "  effectively. Clustering creates the redundancy that compression algorithms exploit.\n",
    "\n",
    "- **Selective Compression Works**: You don't need to compress every layer. Target the\n",
    "  largest layers or those least sensitive to compression for optimal results.\n",
    "\n",
    "- **Compression is a Tradeoff**: We achieved significant size reduction (6-8x\n",
    "  overall) with minimal accuracy loss. Your specific tradeoff will depend on your\n",
    "  application's requirements.\n",
    "\n",
    "### Practical Considerations for Production\n",
    "\n",
    "1. **Operator Support**: Currently, not all TFLM operators support compression. Check\n",
    "   the documentation for your target operators and plan accordingly.\n",
    "\n",
    "2. **Compression Planning**: Before training, consider which layers you'll compress.\n",
    "   You might design your architecture with compression in mind.\n",
    "\n",
    "3. **Fine-tuning Matters**: Fine-tune after clustering to let the model adapt to the\n",
    "   constrained weight space.\n",
    "\n",
    "4. **Test on Target Hardware**: While TFLM's Python interpreter provides accurate\n",
    "   results, validate inference speed on your actual microcontroller.\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "- Experiment with different clustering configurations (4, 8, or 32 clusters)\n",
    "- Try compressing different combinations of layers\n",
    "\n",
    "### Resources\n",
    "\n",
    "- [TensorFlow Model Optimization Toolkit](https://www.tensorflow.org/model_optimization)\n",
    "- [TFLM Compression Documentation](https://github.com/tensorflow/tflite-micro/tree/main/tensorflow/lite/micro/compression)\n",
    "- [Running Compressed Models on Simulators and Hardware](https://github.com/tensorflow/tflite-micro/blob/main/tensorflow/lite/micro/docs/compression.md)\n",
    "- [TensorFlow Lite for Microcontrollers](https://www.tensorflow.org/lite/microcontrollers)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}