{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b2d74a184211852e",
   "metadata": {},
   "source": [
    "# ONNX Quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3a6936c6d945f0",
   "metadata": {},
   "source": [
    "**ONNX2TFLite** provides two packages. The one used for quantizing an ONNX model is called **ONNX2Quant**.\n",
    "Quantization requires a calibration dataset to determine the value ranges within the model.\n",
    "\n",
    "To run quantization, you need to provide input ONNX model and calibration dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b8a2451f8b20d19",
   "metadata": {},
   "source": [
    "### 1. Setup and configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "656f54a6dc97459f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import requests\n",
    "import zipfile"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6657368ca907f844",
   "metadata": {},
   "source": [
    "### 2. Prepare your model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4d88a0a8c914fb6",
   "metadata": {},
   "source": [
    "Let’s prepare the model for quantization. If you already have your own model, you can use that; otherwise, download the example model provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c47ae446390311bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = Path(\"your_model.onnx\")\n",
    "example_model_url = \"https://eiq.nxp.com/training-materials/_misc/models/model.onnx\"\n",
    "\n",
    "with open(model_path, \"wb\") as f:\n",
    "    response = requests.get(\n",
    "        url=example_model_url\n",
    "    )\n",
    "    f.write(response.content)\n",
    "\n",
    "    if response.status_code != 200:\n",
    "        print(f\"Failed to download model: {response.content}\")\n",
    "    else:\n",
    "        print(f\"Model downloaded and stored in {model_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c4ce4ec6d057d4",
   "metadata": {},
   "source": [
    "### 3. Prepare your calibration dataset\n",
    "\n",
    "Calibration dataset has to have specific format and structure. Each model input should have a separate folder containing the input data files. Each file must be a numpy array saved as a .npy file. Following script demonstrates how to prepare calibration dataset for model with single FLOAT32 input:\n",
    "\n",
    "```python\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "\n",
    "# Configuration\n",
    "SOURCE_DIR = Path(\"images\")\n",
    "OUTPUT_DIR = Path(\"calibration_data\")\n",
    "OUTPUT_DIR.mkdir(exist_ok=True)\n",
    "\n",
    "# Convert images to .npy files\n",
    "for idx, img_path in enumerate(SOURCE_DIR.glob(\"*.*\")):\n",
    "    img = Image.open(img_path).resize((224, 224))\n",
    "    img_array = np.array(img, dtype=np.float32) / 255.0\n",
    "    np.save(OUTPUT_DIR / f\"{idx:04d}.npy\", img_array)\n",
    "```\n",
    "\n",
    "In case you don't have your calibration dataset prepared yet, you can use predefined one.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad86c739b6975057",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = Path(\"your_calibration_dataset.zip\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e85f157818624d",
   "metadata": {},
   "outputs": [],
   "source": [
    "example_dataset_url = \"https://eiq.nxp.com/training-materials/_misc/datasets/kws_calib.zip\"\n",
    "\n",
    "with open(dataset_path, \"wb\") as f:\n",
    "    response = requests.get(\n",
    "        url=example_dataset_url\n",
    "    )\n",
    "    f.write(response.content)\n",
    "\n",
    "    if response.status_code != 200:\n",
    "        print(f\"Failed to download model: {response.content}\")\n",
    "    else:\n",
    "        output_dataset_dir = Path(\"unzipped_dataset\")\n",
    "\n",
    "        # unzip dataset\n",
    "        with zipfile.ZipFile(dataset_path) as zip_file:\n",
    "            zip_file.extractall(output_dataset_dir)\n",
    "\n",
    "        print(f\"Dataset downloaded and unzipped in {output_dataset_dir} folder.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80af3440d91763db",
   "metadata": {},
   "source": [
    "### 4. Run Quantization\n",
    "\n",
    "Now we can use the **onnx2quant** tool to quantize the model using the calibration dataset.\n",
    "The calibration dataset is passed using the `-c` argument. You need to provide a dataset for each input tensor in the format:\n",
    "```\n",
    "<input>;<dataset_dir>\n",
    "```\n",
    "\n",
    "If the model has multiple inputs, you can repeat the `-c` argument multiple times.\n",
    "For example, if the model has inputs `x` and `y`, you would specify their calibration datasets like this:\n",
    "```\n",
    "-c \"x;<path_to_x_dataset>\" -c \"y;<path_to_y_dataset>\"\n",
    "```\n",
    "\n",
    "The example model has only one input, so we only need to set this argument once.\n",
    "If you’re using your own model, make sure to adjust these settings accordingly.\n",
    "Once quantization completes, the quantized model will appear in the current directory under the name `<input_model_name>_quant.onnx`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ff94d5c641b014",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_calib_dataset_directory = output_dataset_dir / \"images\"\n",
    "!onnx2quant $model_path -c \"images;$input_calib_dataset_directory\" -o output_model_quant.onnx\n",
    "\n",
    "print(f\"Quantized model saved to output_model_quant.onnx file\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b801e50e4f2af22",
   "metadata": {},
   "source": [
    "### 5. Verify Quantization Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b71d3a8dd1419f01",
   "metadata": {},
   "outputs": [],
   "source": [
    "quantized_model = Path(\"output_model_quant.onnx\")\n",
    "\n",
    "original_size = model_path.stat().st_size / (1024 * 1024)  # MB\n",
    "quantized_size = quantized_model.stat().st_size / (1024 * 1024)  # MB\n",
    "reduction = (1 - quantized_size / original_size) * 100\n",
    "\n",
    "print(f\"Original:  {original_size:.2f} MB\")\n",
    "print(f\"Quantized: {quantized_size:.2f} MB\")\n",
    "print(f\"Reduction: {reduction:.1f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccae0a036b676955",
   "metadata": {},
   "source": [
    "### 6. Additional resources\n",
    "To list all configurable options, use the help command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdaeec507c23d56a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!onnx2quant -h"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "training_materials",
   "language": "python",
   "name": "training_materials"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}