notebooks/305-tensorflow-quantization-aware-training/305-tensorflow-quantization-aware-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f855a151",
   "metadata": {},
   "source": [
    "# Quantization Aware Training with NNCF, using TensorFlow Framework\n",
    "\n",
    "The goal of this notebook to demonstrate how to use the Neural Network Compression Framework [NNCF](https://github.com/openvinotoolkit/nncf) 8-bit quantization to optimize a TensorFlow model for inference with OpenVINO™ Toolkit. The optimization process contains the following steps:\n",
    "* Transforming the original `FP32` model to `INT8`\n",
    "* Using fine-tuning to restore the accuracy.\n",
    "* Exporting optimized and original models to Frozen Graph and then to OpenVINO.\n",
    "* Measuring and comparing the performance of models.\n",
    "\n",
    "For more advanced usage, refer to these [examples](https://github.com/openvinotoolkit/nncf/tree/develop/examples).\n",
    "\n",
    "This tutorial uses the ResNet-18 model with Imagenette dataset. Imagenette is a subset of 10 easily classified classes from the Imagenet dataset. Using the smaller model and dataset will speed up training and download time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddd1b744",
   "metadata": {},
   "source": [
    "## Imports and Settings\n",
    "Import NNCF and all auxiliary packages from your Python code.\n",
    "Set a name for the model, input image size, used batch size, and the learning rate. Also, define paths where Frozen Graph and OpenVINO IR versions of the models will be stored.\n",
    "\n",
    "> **NOTE**: All NNCF logging messages below ERROR level (INFO and WARNING) are disabled to simplify the tutorial. For production use, it is recommended to enable logging by removing ```set_log_level(logging.ERROR)```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1e95247",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import logging\n",
    "\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras import models\n",
    "\n",
    "from nncf import NNCFConfig\n",
    "from nncf.tensorflow.helpers.model_creation import create_compressed_model\n",
    "from nncf.tensorflow.initialization import register_default_init_args\n",
    "from nncf.common.utils.logger import set_log_level\n",
    "\n",
    "set_log_level(logging.ERROR)\n",
    "\n",
    "MODEL_DIR = Path(\"model\")\n",
    "OUTPUT_DIR = Path(\"output\")\n",
    "MODEL_DIR.mkdir(exist_ok=True)\n",
    "OUTPUT_DIR.mkdir(exist_ok=True)\n",
    "\n",
    "BASE_MODEL_NAME = \"ResNet-18\"\n",
    "\n",
    "fp32_h5_path = Path(MODEL_DIR / (BASE_MODEL_NAME + \"_fp32\")).with_suffix(\".h5\")\n",
    "fp32_sm_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + \"_fp32\"))\n",
    "fp32_ir_path = Path(OUTPUT_DIR / \"saved_model\").with_suffix(\".xml\")\n",
    "int8_pb_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + \"_int8\")).with_suffix(\".pb\")\n",
    "int8_pb_name = Path(BASE_MODEL_NAME + \"_int8\").with_suffix(\".pb\")\n",
    "int8_ir_path = int8_pb_path.with_suffix(\".xml\")\n",
    "\n",
    "BATCH_SIZE = 128\n",
    "IMG_SIZE = (64, 64)  # Default Imagenet image size\n",
    "NUM_CLASSES = 10  # For Imagenette dataset\n",
    "\n",
    "LR = 1e-5\n",
    "\n",
    "MEAN_RGB = (0.485 * 255, 0.456 * 255, 0.406 * 255)  # From Imagenet dataset\n",
    "STDDEV_RGB = (0.229 * 255, 0.224 * 255, 0.225 * 255)  # From Imagenet dataset\n",
    "\n",
    "fp32_pth_url = \"https://storage.openvinotoolkit.org/repositories/nncf/openvino_notebook_ckpts/305_resnet18_imagenette_fp32_v1.h5\"\n",
    "_ = tf.keras.utils.get_file(fp32_h5_path.resolve(), fp32_pth_url)\n",
    "print(f'Absolute path where the model weights are saved:\\n {fp32_h5_path.resolve()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f178c3c1",
   "metadata": {},
   "source": [
    "## Dataset Preprocessing\n",
    "\n",
    "Download and prepare Imagenette 160px dataset.\n",
    "- Number of classes: 10\n",
    "- Download size: 94.18 MiB\n",
    "| Split        | Examples |\n",
    "|--------------|----------|\n",
    "| 'train'      | 12,894   |\n",
    "| 'validation' | 500      |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dc3b1a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasets, datasets_info = tfds.load('imagenette/160px', shuffle_files=True, as_supervised=True, with_info=True,\n",
    "                                    read_config=tfds.ReadConfig(shuffle_seed=0))\n",
    "train_dataset, validation_dataset = datasets['train'], datasets['validation']\n",
    "fig = tfds.show_examples(train_dataset, datasets_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e18adc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocessing(image, label):\n",
    "    image = tf.image.resize(image, IMG_SIZE)\n",
    "    image = image - MEAN_RGB\n",
    "    image = image / STDDEV_RGB\n",
    "    label = tf.one_hot(label, NUM_CLASSES)\n",
    "    return image, label\n",
    "\n",
    "\n",
    "train_dataset = (train_dataset.map(preprocessing, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "                              .batch(BATCH_SIZE)\n",
    "                              .prefetch(tf.data.experimental.AUTOTUNE))\n",
    "\n",
    "validation_dataset = (validation_dataset.map(preprocessing, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
    "                                        .batch(BATCH_SIZE)\n",
    "                                        .prefetch(tf.data.experimental.AUTOTUNE))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "978e0fdb",
   "metadata": {},
   "source": [
    "## Define a Floating-Point Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9b5eccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def residual_conv_block(filters, stage, block, strides=(1, 1), cut='pre'):\n",
    "    def layer(input_tensor):\n",
    "        x = layers.BatchNormalization(epsilon=2e-5)(input_tensor)\n",
    "        x = layers.Activation('relu')(x)\n",
    "\n",
    "        # Defining shortcut connection.\n",
    "        if cut == 'pre':\n",
    "            shortcut = input_tensor\n",
    "        elif cut == 'post':\n",
    "            shortcut = layers.Conv2D(filters, (1, 1), strides=strides, kernel_initializer='he_uniform', \n",
    "                                     use_bias=False)(x)\n",
    "\n",
    "        # Continue with convolution layers.\n",
    "        x = layers.ZeroPadding2D(padding=(1, 1))(x)\n",
    "        x = layers.Conv2D(filters, (3, 3), strides=strides, kernel_initializer='he_uniform', use_bias=False)(x)\n",
    "\n",
    "        x = layers.BatchNormalization(epsilon=2e-5)(x)\n",
    "        x = layers.Activation('relu')(x)\n",
    "        x = layers.ZeroPadding2D(padding=(1, 1))(x)\n",
    "        x = layers.Conv2D(filters, (3, 3), kernel_initializer='he_uniform', use_bias=False)(x)\n",
    "\n",
    "        # Add residual connection.\n",
    "        x = layers.Add()([x, shortcut])\n",
    "        return x\n",
    "\n",
    "    return layer\n",
    "\n",
    "\n",
    "def ResNet18(input_shape=None):\n",
    "    \"\"\"Instantiates the ResNet18 architecture.\"\"\"\n",
    "    img_input = layers.Input(shape=input_shape, name='data')\n",
    "\n",
    "    # ResNet18 bottom\n",
    "    x = layers.BatchNormalization(epsilon=2e-5, scale=False)(img_input)\n",
    "    x = layers.ZeroPadding2D(padding=(3, 3))(x)\n",
    "    x = layers.Conv2D(64, (7, 7), strides=(2, 2), kernel_initializer='he_uniform', use_bias=False)(x)\n",
    "    x = layers.BatchNormalization(epsilon=2e-5)(x)\n",
    "    x = layers.Activation('relu')(x)\n",
    "    x = layers.ZeroPadding2D(padding=(1, 1))(x)\n",
    "    x = layers.MaxPooling2D((3, 3), strides=(2, 2), padding='valid')(x)\n",
    "\n",
    "    # ResNet18 body\n",
    "    repetitions = (2, 2, 2, 2)\n",
    "    for stage, rep in enumerate(repetitions):\n",
    "        for block in range(rep):\n",
    "            filters = 64 * (2 ** stage)\n",
    "            if block == 0 and stage == 0:\n",
    "                x = residual_conv_block(filters, stage, block, strides=(1, 1), cut='post')(x)\n",
    "            elif block == 0:\n",
    "                x = residual_conv_block(filters, stage, block, strides=(2, 2), cut='post')(x)\n",
    "            else:\n",
    "                x = residual_conv_block(filters, stage, block, strides=(1, 1), cut='pre')(x)\n",
    "    x = layers.BatchNormalization(epsilon=2e-5)(x)\n",
    "    x = layers.Activation('relu')(x)\n",
    "\n",
    "    # ResNet18 top\n",
    "    x = layers.GlobalAveragePooling2D()(x)\n",
    "    x = layers.Dense(NUM_CLASSES)(x)\n",
    "    x = layers.Activation('softmax')(x)\n",
    "\n",
    "    # Create the model.\n",
    "    model = models.Model(img_input, x)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d49af16",
   "metadata": {},
   "outputs": [],
   "source": [
    "IMG_SHAPE = IMG_SIZE + (3,)\n",
    "model = ResNet18(input_shape=IMG_SHAPE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97fcd3fc",
   "metadata": {},
   "source": [
    "## Pre-train a Floating-Point Model\n",
    "\n",
    "Using NNCF for model compression assumes that the user has a pre-trained model and a training pipeline.\n",
    "\n",
    "> **NOTE** For the sake of simplicity of the tutorial, it is recommended to skip `FP32` model training and load the weights that are provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6e34f08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the floating-point weights.\n",
    "model.load_weights(fp32_h5_path)\n",
    "\n",
    "# Compile the floating-point model.\n",
    "model.compile(loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1),\n",
    "              metrics=[tf.keras.metrics.CategoricalAccuracy(name='acc@1')])\n",
    "\n",
    "# Validate the floating-point model.\n",
    "test_loss, acc_fp32 = model.evaluate(validation_dataset,\n",
    "                                     callbacks=tf.keras.callbacks.ProgbarLogger(stateful_metrics=['acc@1']))\n",
    "print(f\"\\nAccuracy of FP32 model: {acc_fp32:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b80f67d6",
   "metadata": {},
   "source": [
    "Save the floating-point model to the saved model, which will be later used for conversion to OpenVINO IR and further performance measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "450cbcb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save(fp32_sm_path)\n",
    "print(f'Absolute path where the model is saved:\\n {fp32_sm_path.resolve()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13b81167",
   "metadata": {},
   "source": [
    "## Create and Initialize Quantization\n",
    "\n",
    "NNCF enables compression-aware training by integrating into regular training pipelines. The framework is designed so that modifications to your original training code are minor. Quantization is the simplest scenario and requires only 3 modifications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61b74e77",
   "metadata": {},
   "source": [
    "1. Configure NNCF parameters to specify compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cafb190",
   "metadata": {},
   "outputs": [],
   "source": [
    "nncf_config_dict = {\n",
    "    \"input_info\": {\"sample_size\": [1, 3] + list(IMG_SIZE)},\n",
    "    \"log_dir\": str(OUTPUT_DIR),  # The log directory for NNCF-specific logging outputs.\n",
    "    \"compression\": {\n",
    "        \"algorithm\": \"quantization\",  # Specify the algorithm here.\n",
    "    },\n",
    "}\n",
    "nncf_config = NNCFConfig.from_dict(nncf_config_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0fb9e4d",
   "metadata": {},
   "source": [
    "2. Provide a data loader to initialize the values of quantization ranges and determine which activation should be signed or unsigned from the collected statistics, using a given number of samples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54d541d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "nncf_config = register_default_init_args(nncf_config=nncf_config,\n",
    "                                         data_loader=train_dataset,\n",
    "                                         batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddc64e3e",
   "metadata": {},
   "source": [
    "3. Create a wrapped model ready for compression fine-tuning from a pre-trained `FP32` model and a configuration object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c9021c",
   "metadata": {},
   "outputs": [],
   "source": [
    "compression_ctrl, model = create_compressed_model(model, nncf_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffadfd04",
   "metadata": {},
   "source": [
    "Evaluate the new model on the validation set after initialization of quantization. The accuracy should be not far from the accuracy of the floating-point `FP32` model for a simple case like the one being demonstrated here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4aa6d88",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the INT8 model.\n",
    "model.compile(optimizer=tf.keras.optimizers.Adam(lr=LR),\n",
    "              loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1),\n",
    "              metrics=[tf.keras.metrics.CategoricalAccuracy(name='acc@1')])\n",
    "\n",
    "# Validate the INT8 model.\n",
    "test_loss, test_acc = model.evaluate(validation_dataset,\n",
    "                                     callbacks=tf.keras.callbacks.ProgbarLogger(stateful_metrics=['acc@1']))\n",
    "print(f\"\\nAccuracy of INT8 model after initialization: {test_acc:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16d0b8f8",
   "metadata": {},
   "source": [
    "## Fine-tune the Compressed Model\n",
    "\n",
    "At this step, a regular fine-tuning process is applied to further improve quantized model accuracy. Normally, several epochs of tuning are required with a small learning rate, the same that is usually used at the end of the training of the original model. No other changes in the training pipeline are required. Here is a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4047d88",
   "metadata": {
    "scrolled": true,
    "tags": [],
    "test_replace": {
     "fit(train_dataset,": "fit(validation_dataset,"
    }
   },
   "outputs": [],
   "source": [
    "# Train the INT8 model.\n",
    "model.fit(train_dataset,\n",
    "          epochs=2)\n",
    "\n",
    "# Validate the INT8 model.\n",
    "test_loss, acc_int8 = model.evaluate(validation_dataset,\n",
    "                                     callbacks=tf.keras.callbacks.ProgbarLogger(stateful_metrics=['acc@1']))\n",
    "print(f\"\\nAccuracy of INT8 model after fine-tuning: {acc_int8:.3f}\")\n",
    "print(f\"\\nAccuracy drop of tuned INT8 model over pre-trained FP32 model: {acc_fp32 - acc_int8:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7af453ef",
   "metadata": {},
   "source": [
    "Save the `INT8` model to the frozen graph (saved model does not work with quantized model for now). Frozen graph will be later used for conversion to OpenVINO IR and further performance measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b208b6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "compression_ctrl.export_model(int8_pb_path, 'frozen_graph')\n",
    "print(f'Absolute path where the int8 model is saved:\\n {int8_pb_path.resolve()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1248a563",
   "metadata": {},
   "source": [
    "## Export Frozen Graph Models to OpenVINO Intermediate Representation (IR)\n",
    "\n",
    "Use Model Optimizer to convert the Saved Model and Frozen Graph models to OpenVINO IR. The models are saved to the current directory.\n",
    "\n",
    "For more information about Model Optimizer, see the [Model Optimizer Developer Guide](https://docs.openvino.ai/latest/openvino_docs_MO_DG_Deep_Learning_Model_Optimizer_DevGuide.html).\n",
    "\n",
    "Executing this command may take a while. There may be some errors or warnings in the output. When Model Optimization successfully exports to OpenVINO IR, the last lines of the output will include: `[ SUCCESS ] Generated IR version 11 model`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92fda382",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo --input_shape=[1,64,64,3] --input=data --saved_model_dir=$fp32_sm_path --output_dir=$OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8adccc72",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mo --input_shape=[1,64,64,3] --input=Placeholder --input_model=$int8_pb_path --output_dir=$OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f048b410",
   "metadata": {},
   "source": [
    "## Benchmark Model Performance by Computing Inference Time\n",
    "\n",
    "Finally, measure the inference performance of the `FP32` and `INT8` models, using [Benchmark Tool](https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html) - an inference performance measurement tool in OpenVINO. By default, Benchmark Tool runs inference for 60 seconds in asynchronous mode on CPU. It returns inference speed as latency (milliseconds per image) and throughput (frames per second) values.\n",
    "\n",
    "> **NOTE**: This notebook runs `benchmark_app` for 15 seconds to give a quick indication of performance. For more accurate performance, it is recommended to 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": "63355744",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def parse_benchmark_output(benchmark_output):\n",
    "    parsed_output = [line for line in benchmark_output if not (line.startswith(r\"[\") or line.startswith(\"  \") or line == \"\")]\n",
    "    print(*parsed_output, sep='\\n')\n",
    "\n",
    "\n",
    "print('Benchmark FP32 model (IR)')\n",
    "benchmark_output = ! benchmark_app -m $fp32_ir_path -d CPU -api async -t 15\n",
    "parse_benchmark_output(benchmark_output)\n",
    "\n",
    "print('\\nBenchmark INT8 model (IR)')\n",
    "benchmark_output = ! benchmark_app -m $int8_ir_path -d CPU -api async -t 15\n",
    "parse_benchmark_output(benchmark_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f60b19e9",
   "metadata": {},
   "source": [
    "Show CPU Information for reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71a6407c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openvino.runtime import Core\n",
    "\n",
    "ie = Core()\n",
    "ie.get_property('CPU', \"FULL_DEVICE_NAME\")"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Отсутствует",
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}