02-model_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2ae4ec38-5a70-4288-941a-7f7b7f712659",
   "metadata": {},
   "source": [
    "# Model training with YOLOv5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63e58114-57d5-4731-aa88-7e04ecd7f505",
   "metadata": {},
   "source": [
    "## Prepare YOLO\n",
    "\n",
    "This step is pretty easy, we only have to download YOLOv5 from Ultralytics and install the necessary dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c09241b6-8f13-4e20-bfd2-74afee1170b4",
   "metadata": {},
   "source": [
    "### Download YOLOv5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68197c4e-f6c3-48d0-b066-5726e6a055e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/ultralytics/yolov5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e89c6eb-98c7-4e9d-8756-e85a18af21c0",
   "metadata": {},
   "source": [
    "### We must slightly edit the requirements to have a container-compatible version of openCV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07282081-61c0-499a-959d-710b477e9415",
   "metadata": {},
   "outputs": [],
   "source": [
    "!sed -i 's/opencv-python/opencv-python-headless/' yolov5/requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "164e215b-7473-45a4-a91a-97b335546320",
   "metadata": {},
   "source": [
    "### Then install the requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8123c9c-80a7-4831-9bc7-b67a309fcd33",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qr yolov5/requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "514d72a7-04eb-4f22-b618-84d18f3e05d5",
   "metadata": {},
   "source": [
    "### Unfortunately, some subdependencies may have messed up with OpenCV. So let's make sure again we have the right version again..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b358c19d-c4d2-4562-bfdc-5884d437d46b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip uninstall -qy opencv-python\n",
    "!pip uninstall -qy opencv-python-headless\n",
    "!pip install -q opencv-python-headless"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3a490a-ad68-4bc5-afc5-cfdd048a93e9",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fae012df-92f7-4fc0-b012-2159ccdd2139",
   "metadata": {},
   "source": [
    "### Create the configuration file\n",
    "\n",
    "There is a `configuration.yaml` file already present in the folder. Verify that it has the right number and names of the classes you want to use and that you downloaded."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "446ddbf6-0e6e-49c0-ba40-eb9108beabd3",
   "metadata": {},
   "source": [
    "### Freeze the YOLOv5 Backbone\n",
    "The backbone means the layers that extract input image features. We will freeze the backbone so the weights in the backbone layers will not change during YOLOv5 transfer learning. We will then only train the last layers (the head layers).\n",
    "\n",
    "In this example, we will use the smallest available base model, yolov5n. If we open the file `yolov5/models/yolov5n.yaml` we can see the following structure:"
   ]
  },
  {
   "cell_type": "raw",
   "id": "9865eca3-13ab-4c93-bf49-d08442a58925",
   "metadata": {},
   "source": [
    "# YOLOv5 🚀 by Ultralytics, GPL-3.0 license\n",
    "\n",
    "# Parameters\n",
    "nc: 80  # number of classes\n",
    "depth_multiple: 0.33  # model depth multiple\n",
    "width_multiple: 0.25  # layer channel multiple\n",
    "anchors:\n",
    "  - [10,13, 16,30, 33,23]  # P3/8\n",
    "  - [30,61, 62,45, 59,119]  # P4/16\n",
    "  - [116,90, 156,198, 373,326]  # P5/32\n",
    "\n",
    "# YOLOv5 v6.0 backbone\n",
    "backbone:\n",
    "  # [from, number, module, args]\n",
    "  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2\n",
    "   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4\n",
    "   [-1, 3, C3, [128]],\n",
    "   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8\n",
    "   [-1, 6, C3, [256]],\n",
    "   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16\n",
    "   [-1, 9, C3, [512]],\n",
    "   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32\n",
    "   [-1, 3, C3, [1024]],\n",
    "   [-1, 1, SPPF, [1024, 5]],  # 9\n",
    "  ]\n",
    "\n",
    "# YOLOv5 v6.0 head\n",
    "head:\n",
    "  [[-1, 1, Conv, [512, 1, 1]],\n",
    "   [-1, 1, nn.Upsample, [None, 2, 'nearest']],\n",
    "   [[-1, 6], 1, Concat, [1]],  # cat backbone P4\n",
    "   [-1, 3, C3, [512, False]],  # 13\n",
    "\n",
    "   [-1, 1, Conv, [256, 1, 1]],\n",
    "   [-1, 1, nn.Upsample, [None, 2, 'nearest']],\n",
    "   [[-1, 4], 1, Concat, [1]],  # cat backbone P3\n",
    "   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)\n",
    "\n",
    "   [-1, 1, Conv, [256, 3, 2]],\n",
    "   [[-1, 14], 1, Concat, [1]],  # cat head P4\n",
    "   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)\n",
    "\n",
    "   [-1, 1, Conv, [512, 3, 2]],\n",
    "   [[-1, 10], 1, Concat, [1]],  # cat head P5\n",
    "   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)\n",
    "\n",
    "   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)\n",
    "  ]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eba8e6fd-ac8a-4021-9827-b189837b9e89",
   "metadata": {},
   "source": [
    "In the **backbone** section we can see that there are 10 layers. So that's the number of layers we want to freeze in our training!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6bcfae3-e88f-42e9-9c69-7502f5f97acf",
   "metadata": {},
   "source": [
    "### Some training parameters (hyperparameters) we are going to use\n",
    "\n",
    "#### The base model\n",
    "weights = yolov5n.pt\n",
    "\n",
    "This is the smallest available model, to save processing time in this example. Adjust for better results.\n",
    "Other models are available here: https://github.com/ultralytics/yolov5#pretrained-checkpoints\n",
    "\n",
    "####  The number of training iterations\n",
    "epochs = 50\n",
    "\n",
    "This is voluntarily low to save processing time for this example. Adjust for better results (>150).\n",
    "\n",
    "#### The number of images analyzed in a single pass\n",
    "batch = 256\n",
    "\n",
    "You may have to adjust this depending on the memory available on your GPU. The higher (in a power of 2) the better, until you run out of memory...\n",
    "\n",
    "#### Number of layers to freeze\n",
    "freeze = 10\n",
    "\n",
    "As per above, we want to freeze 10 layers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2231bb5-081a-4e98-872b-89073a80ef18",
   "metadata": {},
   "source": [
    "### We can now launch the training!\n",
    "\n",
    "NOTE: PyTorch first caches images to speed up the process. If you have enough memory and shared memory that is not an issue. However, this may not be always the case, especially with large datasets. Therefore the cache is forced to disk here. But of course this is a parameter you can change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7463168-e2e1-429d-92e4-86486059afa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python yolov5/train.py --data configuration.yaml --weights yolov5n.pt --epochs 50 --batch 256 --freeze 10 --cache disk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad367ec3-2c18-4532-8ba9-68c2d07b0ba8",
   "metadata": {},
   "source": [
    "The model has been saved as `yolov5n.pt`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6245d6cb-a3d4-4322-a781-0c4196ca3762",
   "metadata": {},
   "source": [
    "### Test the model\n",
    "\n",
    "We can now simply test the model against an image: `test/img_test.jpg`\n",
    "The image will be analyzed, enriched with bounding boxes, and the result will be saved under `yolov5/runs/detect/exp`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a42f9280-ed54-4821-a707-1b78f2eaa75c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python yolov5/detect.py --weights yolov5n.pt --img 640 --conf-thres 0.2 --source test/img_test.jpg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44e1ec57-1189-44c6-a9b5-c55cb8c96bee",
   "metadata": {},
   "source": [
    "![Result](./yolov5/runs/detect/exp/img_test.jpg)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}