tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_data_generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Data Generation Tutorial: Data-Free Quantization with the Model Compression Toolkit"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "74a56f1fe3c17fcf"
  },
  {
   "cell_type": "markdown",
   "source": [
    "[Run this tutorial in Google Colab](https://colab.research.google.com/github/sony/model_optimization/blob/main/tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_data_generation.ipynb)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "547eb47b9afe4dc0"
  },
  {
   "cell_type": "markdown",
   "source": [
    "In this tutorial, we will explore how to generate synthetic images using the Model Compression Toolkit (MCT) and the Data Generation Library. These generated images are based on the statistics stored in the model's batch normalization layers and can be usefull for various compression tasks, such as quantization and pruning. We will use the generated images as a representative dataset to quantize our model to 8-bit using MCT's Post Training Quantization (PTQ).\n",
    "\n",
    "We will cover the following steps:\n",
    "1. **Setup** Install and import necessary libraries and load a pre-trained model.\n",
    "2. **Configuration**: Define the data generation configuration.\n",
    "3. **Data Generation**: Generate synthetic images.\n",
    "4. **Visualization**: Visualize the generated images.\n",
    "5. **Quantization**: Quantize our model to 8-bit using PTQ with the generated images as a representative dataset. This is called **\"Data-Free Quantization\"** since no real data is used in the quantization process."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "25cb505c44118f02"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 1: Setup\n",
    "Install the necessary packages:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ce2d053e4b52db07"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q torch torchvision\n",
    "!pip install -q model-compression-toolkit"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "941089a3a8cbdf3b"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Imports:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "58d031b0b282dd59"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision.models import resnet18, ResNet18_Weights\n",
    "import model_compression_toolkit as mct\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "879b3f41ba5f6921"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Load the model from the torchvision library:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "a6a2a1f3e024127"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Load a pre-trained model (e.g., ResNet18)\n",
    "model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2b0e10d5f3471530"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 2: Define a Data Generation Configuration\n",
    "Next, we need to specify the configuration for data generation using 'get_pytorch_data_generation_config'. This configuration includes parameters such as the number of iterations, optimizer, batch size, and more. Customize these parameters according to your needs."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "2058e5c7f24c10f1"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "data_gen_config = mct.data_generation.get_pytorch_data_generation_config(\n",
    "    n_iter=500,                      # Number of iterations\n",
    "    optimizer=torch.optim.RAdam,     # Optimizer\n",
    "    data_gen_batch_size=32,          # Batch size for data generation\n",
    "    initial_lr=16,                   # Initial learning rate\n",
    "    output_loss_multiplier=1e-6,     # Multiplier for output loss\n",
    "    extra_pixels=32, \n",
    "    # ... (customize other parameters)\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "833dcf17ef4b49ac"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 3: Generate Synthetic Images\n",
    "\n",
    "Now, let's generate synthetic images using the 'pytorch_data_generation_experimental' function. Specify the number of images you want to generate and the output image size."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "4e5fde47532e18fa"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "n_images = 256              # Number of images to generate\n",
    "output_image_size = 224     # Size of output images\n",
    "\n",
    "generated_images = mct.data_generation.pytorch_data_generation_experimental(\n",
    "    model=model,\n",
    "    n_images=n_images,\n",
    "    output_image_size=output_image_size,\n",
    "    data_generation_config=data_gen_config\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "606e623b954db7c5"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 4: Visualization\n",
    "Lets begin by defining some functions to display the generated images in a grid:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "831e12849e9962a4"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def plot_image_grid(images, reverse_preprocess=False, titles=[], ncols=None, cmap='gray', mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):\n",
    "    images = [plot_image(img, reverse_preprocess, mean, std, plot_img=False) for img in images]\n",
    "    if len(titles) < len(images):\n",
    "        titles += ['_' for _ in range(len(images) - len(titles))]\n",
    "    '''Plot a grid of images'''\n",
    "    if not ncols:\n",
    "        factors = [i for i in range(1, len(images)+1) if len(images) % i == 0]\n",
    "        ncols = factors[len(factors) // 2] if len(factors) else len(images) // 4 + 1\n",
    "    nrows = int(len(images) / ncols) + int(len(images) % ncols)\n",
    "    imgs = [images[i] if len(images) > i else None for i in range(nrows * ncols)]\n",
    "    f, axes = plt.subplots(nrows, ncols, figsize=(3*ncols, 2*nrows))\n",
    "    axes = axes.flatten()[:len(imgs)]\n",
    "    for img, ax, t in zip(imgs, axes.flatten(), titles):\n",
    "        if np.any(img):\n",
    "            if len(img.shape) > 2 and img.shape[2] == 1:\n",
    "                img = img.squeeze()\n",
    "            ax.imshow(img, cmap=cmap)\n",
    "            ax.title.set_text(t)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_image(image, reverse_preprocess=False, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], plot_img=True):\n",
    "    image = image.detach().cpu().numpy()\n",
    "    if len(image.shape) == 4:\n",
    "        image = image[0, :, :, :]\n",
    "    if image.shape[0] == 3:\n",
    "        image = image.transpose(1, 2, 0)\n",
    "    if reverse_preprocess:\n",
    "        new_image = np.round(((image.astype(np.float32) * std) + mean) * 255).astype(np.uint8)\n",
    "    else:\n",
    "        new_image = image\n",
    "    if plot_img:\n",
    "        plt.imshow(new_image)\n",
    "        plt.show()\n",
    "    return new_image"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "efda94fc010113e3"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, lets visualize our generated images:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "22c54ff3d5b2ad"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "plot_image_grid(generated_images[69:71], True)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "6a7faa41b7481573"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Step 5: Post Training Quantization\n",
    "In order to evaulate our generated images, we will use them to quantize the model using MCT's PTQ.This is called **\"Data-Free Quantization\"** because no real data is used in the quantization process. "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "7b40f70b4132c5fb"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Setup for evaluation on the ImageNet dataset\n",
    "Here we define functions for evaluation on ImageNet:\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c4ccf92648d8bc20"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from torchvision import datasets, transforms\n",
    "from tqdm import tqdm\n",
    "\n",
    "# If GPU available, move the model to GPU\n",
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "# Load a pre-trained model (e.g., ResNet18)\n",
    "model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)\n",
    "\n",
    "def get_validation_loader(imagenet_validation_folder, batch_size=50):\n",
    "    preprocess = transforms.Compose([\n",
    "            transforms.Resize(256),\n",
    "            transforms.CenterCrop(224),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225]),\n",
    "        ])\n",
    "    data_loader = torch.utils.data.DataLoader(\n",
    "        datasets.ImageFolder(imagenet_validation_folder, preprocess),\n",
    "        batch_size=batch_size, shuffle=False,\n",
    "        num_workers=8, pin_memory=True)\n",
    "    return data_loader\n",
    "\n",
    "def eval(outputs, labels, topk=(1,)):\n",
    "    maxk = max(topk)\n",
    "\n",
    "    _, pred = outputs.topk(maxk, 1, True, True)\n",
    "    pred = pred.t()\n",
    "    correct = pred.eq(labels.view(1, -1).expand_as(pred))\n",
    "    return correct\n",
    "\n",
    "\n",
    "def accuracy(outputs, labels, topk=(1,)):\n",
    "    \"\"\"Computes the accuracy over the k top predictions for the specified values of k\"\"\"\n",
    "\n",
    "    correct = eval(outputs, labels)\n",
    "\n",
    "    batch_size = labels.size(0)\n",
    "\n",
    "    res = []\n",
    "    for k in topk:\n",
    "        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)\n",
    "        res.append(correct_k.mul_(100.0 / batch_size))\n",
    "\n",
    "    return res, correct\n",
    "\n",
    "\n",
    "def pytorch_model_accuracy_evaluation(model, val_data_loader) -> float:\n",
    "    model = model.to(DEVICE)\n",
    "    model.eval()\n",
    "    acc_top1 = 0\n",
    "\n",
    "    batch_cntr = 1\n",
    "    iterations = len(val_data_loader)\n",
    "    with torch.no_grad():\n",
    "        for input_data, target_data in tqdm(val_data_loader):\n",
    "            inputs_batch = input_data.to(DEVICE)\n",
    "            target_batch = target_data.to(DEVICE)\n",
    "\n",
    "\n",
    "            predicted_batch = model(inputs_batch)\n",
    "\n",
    "            batch_avg_top_1, correct_inds = accuracy(outputs=predicted_batch, labels=target_batch)\n",
    "            acc_top1 += batch_avg_top_1[0].item()\n",
    "    \n",
    "    \n",
    "            batch_cntr += 1\n",
    "            if batch_cntr > iterations:\n",
    "                break\n",
    "    acc_top1 /= iterations\n",
    "    return acc_top1    "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "ad895ccd05275eb9"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Here we define configurations for MCT's PTQ:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "e1a9b2df31324281"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "num_calibration_iter = 10\n",
    "batch_size=50\n",
    "target_platform_cap = mct.get_target_platform_capabilities(\"pytorch\", \"default\")\n",
    "core_config = mct.core.CoreConfig(quantization_config=mct.core.QuantizationConfig())"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "672ffbf357234def"
  },
  {
   "cell_type": "markdown",
   "source": [
    "Specify the path to the imagenet validation folder:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5aa6547351df38bd"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "imagenet_validation_folder = '/path/to/imagenet/validation/folder'\n",
    "val_loader = get_validation_loader(imagenet_validation_folder)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "5392d6e44eebb864"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Quantization with our generated images\n",
    "In this section we use our generated images as a representative dataset for PTQ:"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "97073eeea51b4dee"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "batches_inds = np.random.choice(len(generated_images),\n",
    "                                size=(int(len(generated_images) / batch_size),batch_size),\n",
    "                                replace=False)\n",
    "def representative_data_gen():\n",
    "    for nn in range(num_calibration_iter):\n",
    "        nn_mod = nn % len(batches_inds)\n",
    "        yield [np.concatenate([generated_images[b].detach().cpu().numpy() for b in batches_inds[nn_mod]], axis=0)]\n",
    "        \n",
    "# run post training quantization on the model to get the quantized model output\n",
    "quantized_model_generated_data, quantization_info = mct.ptq.pytorch_post_training_quantization(\n",
    "    in_module=model,\n",
    "    representative_data_gen=representative_data_gen,\n",
    "    core_config=core_config,\n",
    "    target_platform_capabilities=target_platform_cap\n",
    ")"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c7f57ae27466992e"
  },
  {
   "cell_type": "markdown",
   "source": [
    "### Evaluation of the quantized model's performance\n",
    "Here we evaluate our model's top 1 classification performance after PTQ on the ImageNet validation dataset."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8cb543dcc791ee24"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "accuracy_values = pytorch_model_accuracy_evaluation(quantized_model_generated_data, val_loader)\n",
    "print('Float model\\'s reported top 1 performance on ImageNet: 69.86')\n",
    "print(f'Data-Free quantized model\\'s top 1 performance on ImageNet: {accuracy_values}')"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "857b5d4111a42071"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Conclusion:\n",
    "In this tutorial we demonstrated how to generate synthetic images from a trained model and how to use these images for quantizing the model. The quantized model's size is x4 compressed compared to the original float model, however, its performance is similar to the repored float result. No real data was needed in this process. "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "902909871aff0db"
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Copyrights:\n",
    "Copyright 2023 Sony Semiconductor Israel, Inc. All rights reserved.\n",
    "Licensed under the Apache License, Version 2.0 (the “License”); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0\n",
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an “AS IS” BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "b2a030eb3ee565ef"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "44a80837535b9c08"
  }
 ],
 "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
}