Add instructions

demo.ipynb

@@ -15,7 +15,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%env CUDA_VISIBLE_DEVICES=0"
+    "# You may need to uncomment and change this variable to match free GPU index\n",
+    "# %env CUDA_VISIBLE_DEVICES=0"
    ]
   },
   {
@@ -35,7 +36,7 @@
     "    'batch_size': 3, # 32\n",
     "}\n",
     "\n",
-    "run_val(**opt);\n"
+    "run_val(**opt);"
    ]
   },
   {
@@ -46,13 +47,21 @@
     "## Optimize model with ENOT"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "7697bfbb-37b8-4c24-abcb-8c999a073a9f",
+   "metadata": {},
+   "source": [
+    "Metrics after optimization may be very low, but it will recover after tuning phase. Depending on `target-latency-fraction`, drop of the metrics can be different. The complexity of your detection task also affects possible drop of the metrics after pruning.\n",
+    "\n",
+    "It is always the trade-off between how fast you want your model to be and the final metrics after fine-tuning."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "58eaeb32",
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "import torch\n",
@@ -63,13 +72,13 @@
     "    'data': 'data/coco128.yaml',\n",
     "    'weights': 'yolov5s.pt',\n",
     "    'half': True,\n",
-    "    'batch_size': 3, # 32\n",
+    "    'batch_size': 3, # Better use batch_size > 8 because of batch norm\n",
     "    'imgsz': 640,\n",
     "    'hyp': 'data/hyps/hyp.coco_pruning.yaml',\n",
     "    'name': 'prune_yolov5s_coco',\n",
     "    'save_before_prune': True,\n",
     "    'n_search_steps': 3, # This value is just for demo, in production we recommend to use more than 200 steps.\n",
-    "    'target-latency-fraction': 0.5, # It means that optimized model will be 2 times faster that baseline.\n",
+    "    'target-latency-fraction': 0.5, # It means that optimized model will be 2 times faster than baseline.\n",
     "}\n",
     "\n",
     "run_prune(**opt);\n",
@@ -84,6 +93,14 @@
     "## Make onnx for original and optimized models"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "4e366058-95f9-4b0f-b630-df0e782e8a9d",
+   "metadata": {},
+   "source": [
+    "You can open ONNXes using [Netron](https://netron.app/) and visualize model architectures. You can see that pruned model has different number of channels compared to the original one."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -118,6 +135,14 @@
     "## Run optimized model tuning"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "d119bbed-b4a3-4015-9a42-ef9f5ce42c19",
+   "metadata": {},
+   "source": [
+    "This phase is required because just after pruning with low `target-latency-fraction` metrics will be close to zero. But just after the first epoch of tuning metrics will start to recover."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -131,10 +156,10 @@
     "opt = {\n",
     "    'data': 'data/coco128.yaml',\n",
     "    'weights': 'runs/prune/prune_yolov5s_coco/weights/pruned_model.pt',\n",
-    "    'batch_size': 3, # 32\n",
+    "    'batch_size': 3, # Better use batch_size > 8 because of batch norm\n",
     "    'imgsz': 640,\n",
     "    'from_pruned': True,\n",
-    "    'epochs': 1,\n",
+    "    'epochs': 1, # Just for demo, in the real task it needs to be a half of baseline training epochs or lower.\n",
     "    'device': 0,\n",
     "    'name': 'tune_pruned_model',\n",
     "}\n",
@@ -165,7 +190,7 @@
     "    'data': 'data/coco128.yaml',\n",
     "    'weights': 'runs/train/tune_pruned_model/weights/best.pt',\n",
     "    'half': True,\n",
-    "    'batch_size': 3, # 32\n",
+    "    'batch_size': 3, # Better use batch_size > 8 because of batch norm\n",
     "    'imgsz': 640,\n",
     "}\n",
     "\n",
@@ -207,7 +232,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Uncomment this if you want to show results\n",
+    "# Show detection results for original and optimized models\n",
     "\n",
     "%matplotlib inline\n",
     "import cv2\n",
@@ -449,7 +474,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.10.12"
   }
  },
  "nbformat": 4,