notebooks/105-language-quantize-bert/105-language-quantize-bert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cacbe6b4",
   "metadata": {
    "id": "rQc-wXjqrEuR"
   },
   "source": [
    "# Quantize NLP models with Post-Training Quantization ​in NNCF\n",
    "This tutorial demonstrates how to apply `INT8` quantization to the Natural Language Processing model known as [BERT](https://en.wikipedia.org/wiki/BERT_(language_model)), using the [Post-Training Quantization API](https://docs.openvino.ai/latest/nncf_ptq_introduction.html) (NNCF library). A fine-tuned [HuggingFace BERT](https://huggingface.co/transformers/model_doc/bert.html) [PyTorch](https://pytorch.org/) model, trained on the [Microsoft Research Paraphrase Corpus (MRPC)](https://www.microsoft.com/en-us/download/details.aspx?id=52398), will be used. The tutorial is designed to be extendable to custom models and datasets. It consists of the following steps:\n",
    "\n",
    "- Download and prepare the BERT model and MRPC dataset.\n",
    "- Define data loading and accuracy validation functionality.\n",
    "- Prepare the model for quantization.\n",
    "- Run optimization pipeline.\n",
    "- Load and test quantized model.\n",
    "- Compare the performance of the original, converted and quantized models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "694d9fc1-501c-4b86-a747-637e2aad64ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q nncf datasets evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d6b41e6-132b-40da-b3b9-91bacba29e31",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "771388d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import time\n",
    "from pathlib import Path\n",
    "from zipfile import ZipFile\n",
    "from typing import Iterable\n",
    "from typing import Any\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from openvino import runtime as ov\n",
    "from openvino.tools import mo\n",
    "from openvino.runtime import serialize, Model\n",
    "import nncf\n",
    "from nncf.parameters import ModelType\n",
    "from transformers import BertForSequenceClassification, BertTokenizer\n",
    "import datasets\n",
    "import evaluate\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "from notebook_utils import download_file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9e66896-d439-4065-868a-65b44d31525a",
   "metadata": {},
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "284e9a4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the data and model directories, source URL and the filename of the model.\n",
    "DATA_DIR = \"data\"\n",
    "MODEL_DIR = \"model\"\n",
    "MODEL_LINK = \"https://download.pytorch.org/tutorial/MRPC.zip\"\n",
    "FILE_NAME = MODEL_LINK.split(\"/\")[-1]\n",
    "PRETRAINED_MODEL_DIR = os.path.join(MODEL_DIR, \"MRPC\")\n",
    "\n",
    "os.makedirs(DATA_DIR, exist_ok=True)\n",
    "os.makedirs(MODEL_DIR, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44dc335d",
   "metadata": {
    "id": "YytHDzLE0uOJ",
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Prepare the Model\n",
    "Perform the following:\n",
    "- Download and unpack pre-trained BERT model for MRPC by PyTorch.\n",
    "- Convert the model to the ONNX.\n",
    "- Run Model Optimizer to convert the model from the ONNX representation to the OpenVINO Intermediate Representation (OpenVINO IR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be9fc64c",
   "metadata": {},
   "outputs": [],
   "source": [
    "download_file(MODEL_LINK, directory=MODEL_DIR, show_progress=True)\n",
    "with ZipFile(f\"{MODEL_DIR}/{FILE_NAME}\", \"r\") as zip_ref:\n",
    "    zip_ref.extractall(MODEL_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dca2fa0",
   "metadata": {
    "id": "ehX7F6KB0uPu"
   },
   "source": [
    "Convert the original PyTorch model to the ONNX representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb2f6d66",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 1\n",
    "MAX_SEQ_LENGTH = 128\n",
    "\n",
    "\n",
    "def export_model_to_onnx(model, path):\n",
    "    with torch.no_grad():\n",
    "        default_input = torch.ones(1, MAX_SEQ_LENGTH, dtype=torch.int64)\n",
    "        inputs = {\n",
    "            \"input_ids\": default_input,\n",
    "            \"attention_mask\": default_input,\n",
    "            \"token_type_ids\": default_input,\n",
    "        }\n",
    "        torch.onnx.export(\n",
    "            model,\n",
    "            (inputs[\"input_ids\"], inputs[\"attention_mask\"], inputs[\"token_type_ids\"]),\n",
    "            path,\n",
    "            opset_version=11,\n",
    "            input_names=[\"input_ids\", \"attention_mask\", \"token_type_ids\"],\n",
    "            output_names=[\"output\"]\n",
    "        )\n",
    "        print(\"ONNX model saved to {}\".format(path))\n",
    "\n",
    "\n",
    "torch_model = BertForSequenceClassification.from_pretrained(PRETRAINED_MODEL_DIR)\n",
    "onnx_model_path = Path(MODEL_DIR) / \"bert_mrpc.onnx\"\n",
    "if not onnx_model_path.exists():\n",
    "    export_model_to_onnx(torch_model, onnx_model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6911365",
   "metadata": {
    "id": "sNWDAGGd0uRt"
   },
   "source": [
    "## Convert the ONNX Model to OpenVINO IR\n",
    "Use Model Optimizer Python API to convert the model to OpenVINO IR with `FP32` precision. For more information about Model Optimizer Python API, see the [Model Optimizer Developer Guide](https://docs.openvino.ai/latest/openvino_docs_MO_DG_Python_API.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20aeea80",
   "metadata": {},
   "outputs": [],
   "source": [
    "ir_model_xml = onnx_model_path.with_suffix(\".xml\")\n",
    "\n",
    "# Convert the ONNX model to OpenVINO IR FP32.\n",
    "if not ir_model_xml.exists():\n",
    "    model = mo.convert_model(onnx_model_path)\n",
    "    serialize(model, str(ir_model_xml))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17f79b5f",
   "metadata": {},
   "source": [
    "## Prepare the Dataset\n",
    "We download the General Language Understanding Evaluation (GLUE) dataset for the MRPC task from HuggingFace datasets.\n",
    "Then, we tokenize the data with a pre-trained BERT tokenizer from HuggingFace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "632fb1fc",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def create_data_source():\n",
    "    raw_dataset = datasets.load_dataset('glue', 'mrpc', split='validation')\n",
    "    tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_DIR)\n",
    "\n",
    "    def _preprocess_fn(examples):\n",
    "        texts = (examples['sentence1'], examples['sentence2'])\n",
    "        result = tokenizer(*texts, padding='max_length', max_length=MAX_SEQ_LENGTH, truncation=True)\n",
    "        result['labels'] = examples['label']\n",
    "        return result\n",
    "    processed_dataset = raw_dataset.map(_preprocess_fn, batched=True, batch_size=1)\n",
    "\n",
    "    return processed_dataset\n",
    "\n",
    "\n",
    "data_source = create_data_source()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e082b01d",
   "metadata": {},
   "source": [
    "## Optimize model using NNCF Post-training Quantization API\n",
    "\n",
    "[NNCF](https://github.com/openvinotoolkit/nncf) provides a suite of advanced algorithms for Neural Networks inference optimization in OpenVINO with minimal accuracy drop.\n",
    "We will use 8-bit quantization in post-training mode (without the fine-tuning pipeline) to optimize BERT.\n",
    "\n",
    "> **Note**: NNCF Post-training Quantization is available as a preview feature in OpenVINO 2022.3 release.\n",
    "Fully functional support will be provided in the next releases.\n",
    "\n",
    "The optimization process contains the following steps:\n",
    "\n",
    "1. Create a Dataset for quantization\n",
    "2. Run `nncf.quantize` for getting an optimized model\n",
    "3. Serialize OpenVINO IR model using `openvino.runtime.serialize` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f4e7fbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the network in OpenVINO Runtime.\n",
    "core = ov.Core()\n",
    "model = core.read_model(ir_model_xml)\n",
    "INPUT_NAMES = [x.any_name for x in model.inputs]\n",
    "\n",
    "\n",
    "def transform_fn(data_item):\n",
    "    \"\"\"\n",
    "    Extract the model's input from the data item.\n",
    "    The data item here is the data item that is returned from the data source per iteration.\n",
    "    This function should be passed when the data item cannot be used as model's input.\n",
    "    \"\"\"\n",
    "    inputs = {\n",
    "        name: np.asarray(data_item[name], dtype=np.int64) for name in INPUT_NAMES\n",
    "    }\n",
    "    return inputs\n",
    "\n",
    "\n",
    "calibration_dataset = nncf.Dataset(data_source, transform_fn)\n",
    "# Quantize the model. By specifying model_type, we specify additional transformer patterns in the model.\n",
    "quantized_model = nncf.quantize(model, calibration_dataset,\n",
    "                                model_type=ModelType.TRANSFORMER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da83574c-7abc-40a8-ae30-431c1b2bd823",
   "metadata": {},
   "outputs": [],
   "source": [
    "compressed_model_xml = 'quantized_bert_mrpc.xml'\n",
    "ov.serialize(quantized_model, compressed_model_xml)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c30ab44",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Load and Test OpenVINO Model\n",
    "\n",
    "To load and test converted model, perform the following:\n",
    "* Load the model and compile it for CPU.\n",
    "* Prepare the input.\n",
    "* Run the inference.\n",
    "* Get the answer from the model output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d79b1a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "core = ov.Core()\n",
    "\n",
    "# Read the model from files.\n",
    "model = core.read_model(model=compressed_model_xml)\n",
    "\n",
    "# Assign dynamic shapes to every input layer.\n",
    "for input_layer in model.inputs:\n",
    "    input_shape = input_layer.partial_shape\n",
    "    input_shape[1] = -1\n",
    "    model.reshape({input_layer: input_shape})\n",
    "\n",
    "# Compile the model for a specific device.\n",
    "compiled_model_int8 = core.compile_model(model=model, device_name=\"CPU\")\n",
    "\n",
    "output_layer = compiled_model_int8.outputs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef1d846e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The Data Source returns a pair of sentences (indicated by `sample_idx`) and the inference compares these sentences and outputs whether their meaning is the same. You can test other sentences by changing `sample_idx` to another value (from 0 to 407)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e72504b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_idx = 5\n",
    "sample = data_source[sample_idx]\n",
    "inputs = {k: torch.unsqueeze(torch.tensor(sample[k]), 0) for k in ['input_ids', 'token_type_ids', 'attention_mask']}\n",
    "\n",
    "result = compiled_model_int8(inputs)[output_layer]\n",
    "result = np.argmax(result)\n",
    "\n",
    "print(f\"Text 1: {sample['sentence1']}\")\n",
    "print(f\"Text 2: {sample['sentence2']}\")\n",
    "print(f\"The same meaning: {'yes' if result == 1 else 'no'}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89920c37-dc2f-4177-b25f-bd8b1d0e34d3",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Compare F1-score of FP32 and INT8 models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeea7cc8-3eed-4474-8f59-ae63197368d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate(model: Model, dataset: Iterable[Any]) -> float:\n",
    "    \"\"\"\n",
    "    Evaluate the model on GLUE dataset. \n",
    "    Returns F1 score metric.\n",
    "    \"\"\"\n",
    "    compiled_model = core.compile_model(model, device_name='CPU')\n",
    "    output_layer = compiled_model.output(0)\n",
    "\n",
    "    metric = evaluate.load('glue', 'mrpc')\n",
    "    INPUT_NAMES = [x.any_name for x in compiled_model.inputs]\n",
    "    for batch in dataset:\n",
    "        inputs = [\n",
    "            np.expand_dims(np.asarray(batch[key], dtype=np.int64), 0) for key in INPUT_NAMES\n",
    "        ]\n",
    "        outputs = compiled_model(inputs)[output_layer]\n",
    "        predictions = outputs[0].argmax(axis=-1)\n",
    "        metric.add_batch(predictions=[predictions], references=[batch['labels']])\n",
    "    metrics = metric.compute()\n",
    "    f1_score = metrics['f1']\n",
    "\n",
    "    return f1_score\n",
    "\n",
    "\n",
    "print('Checking the accuracy of the original model:')\n",
    "metric = validate(model, data_source)\n",
    "print(f'F1 score: {metric:.4f}')\n",
    "\n",
    "print('Checking the accuracy of the quantized model:')\n",
    "metric = validate(quantized_model, data_source)\n",
    "print(f'F1 score: {metric:.4f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f67f6a2",
   "metadata": {
    "id": "vQACMfAUo52V",
    "tags": []
   },
   "source": [
    "## Compare Performance of the Original, Converted and Quantized Models\n",
    "\n",
    "Compare the original PyTorch model with OpenVINO converted and quantized models (`FP32`, `INT8`) to see the difference in performance. It is expressed in Sentences Per Second (SPS) measure, which is the same as Frames Per Second (FPS) for images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "734ae69a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = core.read_model(model=ir_model_xml)\n",
    "\n",
    "# Assign dynamic shapes to every input layer.\n",
    "dynamic_shapes = {}\n",
    "for input_layer in model.inputs:\n",
    "    input_shape = input_layer.partial_shape\n",
    "    input_shape[1] = -1\n",
    "\n",
    "    dynamic_shapes[input_layer] = input_shape\n",
    "\n",
    "model.reshape(dynamic_shapes)\n",
    "\n",
    "# Compile the model for a specific device.\n",
    "compiled_model_fp32 = core.compile_model(model=model, device_name=\"CPU\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f484fff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_samples = 50\n",
    "sample = data_source[0]\n",
    "inputs = {k: torch.unsqueeze(torch.tensor(sample[k]), 0) for k in ['input_ids', 'token_type_ids', 'attention_mask']}\n",
    "\n",
    "with torch.no_grad():\n",
    "    start = time.perf_counter()\n",
    "    for _ in range(num_samples):\n",
    "        torch_model(torch.vstack(list(inputs.values())))\n",
    "    end = time.perf_counter()\n",
    "    time_torch = end - start\n",
    "print(\n",
    "    f\"PyTorch model on CPU: {time_torch / num_samples:.3f} seconds per sentence, \"\n",
    "    f\"SPS: {num_samples / time_torch:.2f}\"\n",
    ")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for _ in range(num_samples):\n",
    "    compiled_model_fp32(inputs)\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(\n",
    "    f\"IR FP32 model in OpenVINO Runtime/CPU: {time_ir / num_samples:.3f} \"\n",
    "    f\"seconds per sentence, SPS: {num_samples / time_ir:.2f}\"\n",
    ")\n",
    "\n",
    "start = time.perf_counter()\n",
    "for _ in range(num_samples):\n",
    "    compiled_model_int8(inputs)\n",
    "end = time.perf_counter()\n",
    "time_ir = end - start\n",
    "print(\n",
    "    f\"OpenVINO IR INT8 model in OpenVINO Runtime/CPU: {time_ir / num_samples:.3f} \"\n",
    "    f\"seconds per sentence, SPS: {num_samples / time_ir:.2f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "add78af0",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Finally, measure the inference performance of OpenVINO `FP32` and `INT8` models. For this purpose, use [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) in OpenVINO.\n",
    "\n",
    "> **Note**: The `benchmark_app` tool is able to measure the performance of the OpenVINO Intermediate Representation (OpenVINO IR) models only. For more accurate performance, run `benchmark_app` in a terminal/command prompt after closing other applications. Run `benchmark_app -m model.xml -d CPU` to benchmark async inference on CPU for one minute. Change `CPU` to `GPU` to benchmark on GPU. Run `benchmark_app --help` to see an overview of all command-line options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f71b38a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference FP32 model (OpenVINO IR)\n",
    "! benchmark_app -m $ir_model_xml -d CPU -api sync"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdf41525",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference INT8 model (OpenVINO IR)\n",
    "! benchmark_app -m $compressed_model_xml -d CPU -api sync"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('notebooks_env')",
   "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.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "cec18e25feb9469b5ff1085a8097bdcd86db6a4ac301d6aeff87d0f3e7ce4ca5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}