{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# ONNX Quantization",
   "id": "b2d74a184211852e"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "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 two parameters:\n",
    "\n",
    "* Input model\n",
    "* Calibration dataset\n",
    "\n",
    "The calibration dataset is specified by mapping input tensor names to NumPy files.\n",
    "All of this is demonstrated in the example below.\n",
    "\n",
    "If you have your own calibration dataset, feel free to use it. Otherwise, you can use the example dataset provided."
   ],
   "id": "ec3a6936c6d945f0"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from pathlib import Path\n",
    "dataset_path = Path(\"your_calibration_dataset.zip\")"
   ],
   "id": "4d7686babaa67c44"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "import requests\n",
    "import zipfile\n",
    "\n",
    "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.\")"
   ],
   "id": "7fa0409ac511bbd3"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Next, let’s prepare the model for quantization. If you already have your own model, you can use that; otherwise, download the example model provided.",
   "id": "3f58f88736655821"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "model_path = Path(\"your_model.onnx\")",
   "id": "a8ec09efe35078e5"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "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}\")"
   ],
   "id": "cc6be4336169a01d"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "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`"
   ],
   "id": "80af3440d91763db"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "input_calib_dataset_directory = output_dataset_dir / \"images\"\n",
    "!onnx2quant $model_path -c \"images;$input_calib_dataset_directory\"\n",
    "\n",
    "print(f\"Quantized model saved to {model_path}_quant.onnx file\")"
   ],
   "id": "a7ff94d5c641b014"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "To list all configurable options, use the help command:",
   "id": "ccae0a036b676955"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "!onnx2quant -h",
   "id": "fdaeec507c23d56a"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}