vision-transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accelerate Vision Transformer (ViT) with Quantization using Optimum\n",
    "\n",
    "In this session, you will learn how to optimize Vision Transformers models using Optimum. The session will show you how to dynamically quantize and optimize a ViT model using [Hugging Face Optimum](https://huggingface.co/docs/optimum/index) and [ONNX Runtime](https://onnxruntime.ai/). Hugging Face Optimum is an extension of 🤗 Transformers, providing a set of performance optimization tools enabling maximum efficiency to train and run models on targeted hardware.\n",
    "\n",
    "Note: dynamic quantization is currently only supported for CPUs, so we will not be utilizing GPUs / CUDA in this session.\n",
    "\n",
    "By the end of this session, you see how quantization and optimization with Hugging Face Optimum can result in significant increase in model latency while keeping almost 100% of the full-precision model. Furthermore, you’ll see how to easily apply some advanced quantization and optimization techniques shown here so that your models take much less of an accuracy hit than they would otherwise. \n",
    "\n",
    "You will learn how to:\n",
    "1. Setup Development Environment\n",
    "2. Convert a Hugging Face `Transformers` model to ONNX for inference\n",
    "3. Apply dynamic quantization using `ORTQuantizer` from Optimum\n",
    "4. Test inference with the quantized model\n",
    "5. Evaluate the performance and speed\n",
    "\n",
    "Let's get started! 🚀\n",
    "\n",
    "_This tutorial was created and run on an c6i.xlarge AWS EC2 Instance._\n",
    "\n",
    "---\n",
    "\n",
    "## Quick intro: Vision Transformer (ViT) by Google Brain\n",
    "\n",
    "The Vision Transformer (ViT) is basically BERT, but applied to images. It attains excellent results compared to state-of-the-art convolutional networks. In order to provide images to the model, each image is split into a sequence of fixed-size patches (typically of resolution 16x16 or 32x32), which are linearly embedded. One also adds a [CLS] token at the beginning of the sequence in order to classify images. Next, one adds absolute position embeddings and provides this sequence to the Transformer encoder.\n",
    "\n",
    "![vision-transformer-architecture](./assets/vision-transformer-architecture.webp)\n",
    "\n",
    "- Paper: https://arxiv.org/abs/2010.11929\n",
    "- Official repo (in JAX): https://github.com/google-research/vision_transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup Development Environment\n",
    "\n",
    "Our first step is to install Optimum, along with  Evaluate and some other libraries. Running the following cell will install all the required packages for us including Transformers, PyTorch, and ONNX Runtime utilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"optimum[onnxruntime]==1.5.0\" evaluate[evaluator] sklearn mkl-include mkl --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> If you want to run inference on a GPU, you can install 🤗 Optimum with `pip install optimum[onnxruntime-gpu]`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Convert a Hugging Face `Transformers` model to ONNX for inference\n",
    "\n",
    "Before we can start qunatizing we need to convert our vanilla `transformers` model to the `onnx` format. To do this we will use the new [ORTModelForImageClassification](https://huggingface.co/docs/optimum/onnxruntime/modeling_ort#optimum.onnxruntime.ORTModelForImageClassification) class calling the `from_pretrained()` method with the `from_transformers` attribute. The model we are using is the a fine-tuned version of [google/vit-base-patch16-224-in21k](https://huggingface.co/google/vit-base-patch16-224-in21k) on the [beans](https://huggingface.co/datasets/beans) dataset ([nateraw/vit-base-beans](https://huggingface.co/nateraw/vit-base-beans)) achieving an accuracy of 96.88%.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/miniconda3/lib/python3.9/site-packages/transformers/models/vit/modeling_vit.py:172: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if num_channels != self.num_channels:\n",
      "/home/ubuntu/miniconda3/lib/python3.9/site-packages/transformers/models/vit/modeling_vit.py:177: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if height != self.image_size[0] or width != self.image_size[1]:\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['onnx/preprocessor_config.json']"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from optimum.onnxruntime import ORTModelForImageClassification\n",
    "from transformers import AutoFeatureExtractor\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "model_id=\"nateraw/vit-base-beans\"\n",
    "onnx_path = Path(\"onnx\")\n",
    "\n",
    "# load vanilla transformers and convert to onnx\n",
    "model = ORTModelForImageClassification.from_pretrained(model_id, from_transformers=True)\n",
    "preprocessor = AutoFeatureExtractor.from_pretrained(model_id)\n",
    "\n",
    "# save onnx checkpoint and tokenizer\n",
    "model.save_pretrained(onnx_path)\n",
    "preprocessor.save_pretrained(onnx_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One neat thing about 🤗 Optimum, is that allows you to run ONNX models with the `pipeline()` function from 🤗 Transformers. This means that you get all the pre- and post-processing features for free, without needing to re-implement them for each model! Here's how you can run inference with our vanilla ONNX model:\n",
    "\n",
    "`https://datasets-server.huggingface.co/assets/beans/--/default/validation/30/image/image.jpg`  \n",
    "![test-image](https://datasets-server.huggingface.co/assets/beans/--/default/validation/30/image/image.jpg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'score': 0.9405876994132996, 'label': 'angular_leaf_spot'}, {'score': 0.03206056356430054, 'label': 'bean_rust'}, {'score': 0.02735181152820587, 'label': 'healthy'}]\n"
     ]
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "vanilla_clf = pipeline(\"image-classification\", model=model, feature_extractor=preprocessor)\n",
    "print(vanilla_clf(\"https://datasets-server.huggingface.co/assets/beans/--/default/validation/30/image/image.jpg\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to learn more about exporting transformers model check-out [Convert Transformers to ONNX with Hugging Face Optimum](https://www.philschmid.de/convert-transformers-to-onnx) blog post\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Apply dynamic quantization using `ORTQuantizer` from Optimum\n",
    "\n",
    "The `ORTQuantizer` can be used to apply dynamic quantization to decrease the size of the model size and accelerate latency and inference.\n",
    "\n",
    "_We use the `avx512_vnni` config since the instance is powered by an intel ice-lake CPU supporting avx512._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimum.onnxruntime import ORTQuantizer\n",
    "from optimum.onnxruntime.configuration import AutoQuantizationConfig\n",
    "\n",
    "# create ORTQuantizer and define quantization configuration\n",
    "dynamic_quantizer = ORTQuantizer.from_pretrained(model)\n",
    "dqconfig = AutoQuantizationConfig.avx512_vnni(is_static=False, per_channel=False)\n",
    "\n",
    "# apply the quantization configuration to the model\n",
    "model_quantized_path = dynamic_quantizer.quantize(\n",
    "    save_dir=onnx_path,\n",
    "    quantization_config=dqconfig,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets quickly check the new model size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model file size: 327.40 MB\n",
      "Quantized Model file size: 84.49 MB\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "# get model file size\n",
    "size = os.path.getsize(onnx_path / \"model.onnx\")/(1024*1024)\n",
    "quantized_model = os.path.getsize(onnx_path / \"model_quantized.onnx\")/(1024*1024)\n",
    "\n",
    "print(f\"Model file size: {size:.2f} MB\")\n",
    "print(f\"Quantized Model file size: {quantized_model:.2f} MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test inference with the quantized model\n",
    "\n",
    "[Optimum](https://huggingface.co/docs/optimum/main/en/pipelines#optimizing-with-ortoptimizer) has built-in support for [transformers pipelines](https://huggingface.co/docs/transformers/main/en/main_classes/pipelines#pipelines). This allows us to leverage the same API that we know from using PyTorch and TensorFlow models.\n",
    "Therefore we can load our quantized model with `ORTModelForImageClassification` class and transformers `pipeline`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'score': 0.9412522912025452, 'label': 'angular_leaf_spot'}, {'score': 0.031623296439647675, 'label': 'bean_rust'}, {'score': 0.027124471962451935, 'label': 'healthy'}]\n"
     ]
    }
   ],
   "source": [
    "from optimum.onnxruntime import ORTModelForImageClassification\n",
    "from transformers import pipeline, AutoFeatureExtractor\n",
    "\n",
    "model = ORTModelForImageClassification.from_pretrained(onnx_path, file_name=\"model_quantized.onnx\")\n",
    "preprocessor = AutoFeatureExtractor.from_pretrained(onnx_path)\n",
    "\n",
    "q8_clf = pipeline(\"image-classification\", model=model, feature_extractor=preprocessor)\n",
    "print(q8_clf(\"https://datasets-server.huggingface.co/assets/beans/--/default/validation/30/image/image.jpg\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate the performance and speed\n",
    "\n",
    "To evaluate the model performance and speed are we going to use a the `test` split of the [beans](https://huggingface.co/datasets/beans) dataset containing only 3 classes ('angular_leaf_spot', 'bean_rust', 'healthy') and 128 images. The evaluation was done by using [Huggingface/evaluate](https://huggingface.co/docs/evaluate/index) a library for easily evaluating machine learning models and datasets.\n",
    "\n",
    "We evaluated the vanilla model outside of this example using the same `evaluator` with the vanilla model achieving an accuraccy of `96.88%` on our dataset. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0917986e5ae4f428372289004f62a3c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading builder script:   0%|          | 0.00/1.43k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd587a5366c54de3bbeec329be1e9d50",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading metadata:   0%|          | 0.00/1.05k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:datasets.builder:Using custom data configuration default\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset beans/default (download: 171.69 MiB, generated: 467.19 KiB, post-processed: Unknown size, total: 172.14 MiB) to /home/ubuntu/.cache/huggingface/datasets/beans/default/0.0.0/90c755fb6db1c0ccdad02e897a37969dbf070bed3755d4391e269ff70642d791...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "657c3501617d44149d3e7380391cc20d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "56875d7f0626494599e8557340b6ef63",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/144M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da8547a320d84340ad6a7331819702ab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/18.5M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f1baea0ba719458abcc077d5b84a8ce6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data:   0%|          | 0.00/17.7M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "98d168b6d6894dd493df584c7d9a9235",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00d6209e063745e2bd231f87b119de4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split:   0%|          | 0/1034 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8250a2b900f04e60928f5c57614699b1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating validation split:   0%|          | 0/133 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2fa4e8cc3e048c1b6f141dae5ebb740",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating test split:   0%|          | 0/128 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset beans downloaded and prepared to /home/ubuntu/.cache/huggingface/datasets/beans/default/0.0.0/90c755fb6db1c0ccdad02e897a37969dbf070bed3755d4391e269ff70642d791. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9f05c78e8f074bae838a654e0f65912c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vanilla model: 96.88%\n",
      "Quantized model: 96.88%\n",
      "The quantized model achieves 99.99% accuracy of the fp32 model\n"
     ]
    }
   ],
   "source": [
    "from evaluate import evaluator\n",
    "from datasets import load_dataset\n",
    "\n",
    "e = evaluator(\"image-classification\")\n",
    "eval_dataset = load_dataset(\"beans\",split=[\"test\"])[0]\n",
    "\n",
    "results = e.compute(\n",
    "    model_or_pipeline=q8_clf,\n",
    "    data=eval_dataset,\n",
    "    metric=\"accuracy\",\n",
    "    input_column=\"image\",\n",
    "    label_column=\"labels\",\n",
    "    label_mapping=model.config.label2id,\n",
    "    strategy=\"simple\",\n",
    ")\n",
    "\n",
    "print(f\"Vanilla model: 96.88%\")\n",
    "print(f\"Quantized model: {results['accuracy']*100:.2f}%\")\n",
    "print(f\"The quantized model achieves {round(results['accuracy']/0.9688,4)*100:.2f}% accuracy of the fp32 model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Okay, now let's test the performance (latency) of our quantized model. We are going to use a the [beans sample](https://datasets-server.huggingface.co/assets/beans/--/default/validation/30/image/image.jpg) for the benchmark. To keep it simple, we are going to use a python loop and calculate the avg,mean & p95 latency for our vanilla model and for the quantized model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vanilla model: P95 latency (ms) - 165.06651640004284; Average latency (ms) - 149.00 +\\- 11.22;\n",
      "Quantized model: P95 latency (ms) - 63.56140074997256; Average latency (ms) - 62.81 +\\- 2.18;\n",
      "Improvement through quantization: 2.6x\n"
     ]
    }
   ],
   "source": [
    "from time import perf_counter\n",
    "import numpy as np \n",
    "from PIL import Image\n",
    "import requests\n",
    "\n",
    "payload=\"https://datasets-server.huggingface.co/assets/beans/--/default/validation/30/image/image.jpg\"\n",
    "\n",
    "def measure_latency(pipe):\n",
    "    # prepare date\n",
    "    image = Image.open(requests.get(payload, stream=True).raw)\n",
    "    inputs = pipe.feature_extractor(images=image, return_tensors=\"pt\")\n",
    "    latencies = []\n",
    "    # warm up\n",
    "    for _ in range(10):\n",
    "        _ = pipe.model(**inputs)\n",
    "    # Timed run\n",
    "    for _ in range(200):\n",
    "        start_time = perf_counter()\n",
    "        _ =  pipe.model(**inputs)\n",
    "        latency = perf_counter() - start_time\n",
    "        latencies.append(latency)\n",
    "    # Compute run statistics\n",
    "    time_avg_ms = 1000 * np.mean(latencies)\n",
    "    time_std_ms = 1000 * np.std(latencies)\n",
    "    time_p95_ms = 1000 * np.percentile(latencies,95)\n",
    "    return f\"P95 latency (ms) - {time_p95_ms}; Average latency (ms) - {time_avg_ms:.2f} +\\- {time_std_ms:.2f};\", time_p95_ms\n",
    "\n",
    "\n",
    "vanilla_model=measure_latency(vanilla_clf)\n",
    "quantized_model=measure_latency(q8_clf)\n",
    "\n",
    "print(f\"Vanilla model: {vanilla_model[0]}\")\n",
    "print(f\"Quantized model: {quantized_model[0]}\")\n",
    "print(f\"Improvement through quantization: {round(vanilla_model[1]/quantized_model[1],2)}x\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We managed to accelerate our model latency from 165ms to 64ms or 2.6x while keeping 99.99% of the accuracy. \n",
    "\n",
    "![performance](assets/vit-performance.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "We successfully quantized our vanilla Transformers model with Hugging Face and managed to accelerate our model latency 165ms to 64ms or 2.6x while keeping 99.99% of the accuracy. \n",
    "\n",
    "But I have to say that this isn't a plug and play process you can transfer to any Transformers model, task or dataset. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "7a2c4b191d1ae843dde5cb5f4d1f62fa892f6b79b0f9392a84691e890e33c5a4"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('base')",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}