Tutorial/custom layer

tutorials/pytorch/custom_layer.ipynb

@@ -0,0 +1,243 @@
+{
+ "cells": [
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "# Custom Layer Usage Example\n",
+    "\n",
+    "## Overview\n",
+    "\n",
+    "In this tutorial we will illustrate how to integrate a custom layer with model quantization using the MCT library.\n",
+    "Using a simple object detection model as an example, we will apply post-training quantization, then incorporate a custom NMS layer into the quantized model."
+   ],
+   "id": "27d28731857f2991"
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "## Setup\n",
+    "\n",
+    "### Install & import relevant packages"
+   ],
+   "id": "29309a0291ff4f41"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "!pip install -q torch\n",
+    "!pip install onnx\n",
+    "!pip install -q model_compression_toolkit"
+   ],
+   "id": "initial_id",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from typing import Iterator, List\n",
+    "import model_compression_toolkit as mct\n",
+    "from sony_custom_layers.pytorch.nms import multiclass_nms"
+   ],
+   "id": "3a7da9c475f95aa9",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "## Model Quantization\n",
+    "\n",
+    "### Create Model Instance\n",
+    "\n",
+    "We will start with creating a simple object-detection model instance as an example. You can replace the model with your own model use a pre-trained model (Make sure the model is supported by MCT library). "
+   ],
+   "id": "cb2a4ee6cf3e5e98"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "class ObjectDetector(nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super(ObjectDetector, self).__init__()\n",
+    "        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)\n",
+    "        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
+    "        self.fc1 = nn.Linear(16 * 16 * 16, 128)\n",
+    "        self.fc_bbox = nn.Linear(128, 4)\n",
+    "        self.fc_class = nn.Linear(128, 2)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x = torch.relu(self.conv1(x))\n",
+    "        x = self.pool1(x)\n",
+    "        x = x.view(-1, 16 * 16 * 16)\n",
+    "        x = torch.relu(self.fc1(x))\n",
+    "        bbox = self.fc_bbox(x)\n",
+    "        class_probs = torch.softmax(self.fc_class(x), dim=1)\n",
+    "        return bbox, class_probs\n",
+    "\n",
+    "model = ObjectDetector()\n",
+    "model.eval()"
+   ],
+   "id": "cbe0031bb7f16986",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "### Post-Training Quantization using Model Compression Toolkit\n",
+    "\n",
+    "Now, we're all set to use MCT's post-training quantization. \n",
+    "To begin, we'll define a representative dataset generator. Please note that for demonstration purposes, we will generate random data of the desired image shape instead of using real images.  \n",
+    "Then, we will apply PTQ on our model using the dataset generator we have created."
+   ],
+   "id": "d653b898460b44e2"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "n_iters = 20\n",
+    "batch_size = 4\n",
+    "\n",
+    "def get_representative_dataset(n_iter: int):\n",
+    "    \"\"\"\n",
+    "    This function creates a representative dataset generator. The generator yields numpy\n",
+    "        arrays of batches of shape: [Batch, C, H, W].\n",
+    "    Args:\n",
+    "        n_iter: number of iterations for MCT to calibrate on\n",
+    "    Returns:\n",
+    "        A representative dataset generator\n",
+    "    \"\"\"\n",
+    "    def representative_dataset() -> Iterator[List]:\n",
+    "        for _ in range(n_iter):\n",
+    "            yield [torch.rand(1, 3, 64, 64)]\n",
+    "\n",
+    "    return representative_dataset\n",
+    "\n",
+    "representative_data_generator = get_representative_dataset(n_iter=n_iters)\n",
+    "\n",
+    "quant_model, _ = mct.ptq.pytorch_post_training_quantization(model, representative_data_gen=representative_data_generator)\n",
+    "print('Quantized model is ready')"
+   ],
+   "id": "72d25144f573ead3",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "## Post-process integration using custom layer\n",
+    "\n",
+    "Now that we have a quantized model, we can add it a custom layer. In our example we will add NMS layer by createing a model wrapper that apply NMS over the quantized model output."
+   ],
+   "id": "3a8d28fed3a87f65"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "class PostProcessWrapper(nn.Module):\n",
+    "    def __init__(self,\n",
+    "                 model: nn.Module,\n",
+    "                 score_threshold: float = 0.001,\n",
+    "                 iou_threshold: float = 0.7,\n",
+    "                 max_detections: int = 300):\n",
+    "\n",
+    "        super(PostProcessWrapper, self).__init__()\n",
+    "        self.model = model\n",
+    "        self.score_threshold = score_threshold\n",
+    "        self.iou_threshold = iou_threshold\n",
+    "        self.max_detections = max_detections\n",
+    "\n",
+    "    def forward(self, images):\n",
+    "        # model inference\n",
+    "        outputs = self.model(images)\n",
+    "\n",
+    "        boxes = outputs[0]\n",
+    "        scores = outputs[1]\n",
+    "        nms = multiclass_nms(boxes=boxes, scores=scores, score_threshold=self.score_threshold,\n",
+    "                             iou_threshold=self.iou_threshold, max_detections=self.max_detections)\n",
+    "        return nms\n",
+    "\n",
+    "\n",
+    "\n",
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "quant_model_with_nms = PostProcessWrapper(model=quant_model,\n",
+    "                                    score_threshold=0.001,\n",
+    "                                    iou_threshold=0.7,\n",
+    "                                    max_detections=300).to(device=device)\n",
+    "print('Quantized model with NMS is ready')"
+   ],
+   "id": "baa386a04a8dd664",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "### Model Export\n",
+    "\n",
+    "Now, we can export the quantized model, ready for deployment, into a .onnx format file. Please ensure that the save_model_path has been set correctly."
+   ],
+   "id": "e7ca57539bdc7239"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "mct.exporter.pytorch_export_model(model=quant_model_with_nms,\n",
+    "                                  save_model_path='./qmodel_with_nms.onnx',\n",
+    "                                  repr_dataset=representative_data_generator)"
+   ],
+   "id": "776a6f99bd0a6efe",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": [
+    "Copyright 2025 Sony Semiconductor Israel, Inc. All rights reserved.\n",
+    "\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\n",
+    "\n",
+    "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."
+   ],
+   "id": "bb7c13e41a012f3"
+  }
+ ],
+ "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
+}