Title: Training the YOLOv5 Object Detector on a Custom Dataset - PyImageSearch

URL Source: https://pyimagesearch.com/2022/06/20/training-the-yolov5-object-detector-on-a-custom-dataset/

Published Time: 2022-06-20T13:00:00+00:00

Markdown Content:
* * *

**Table of Contents**
---------------------

*   [Training YOLOv5 Object Detector on a Custom Dataset](#h2Title)

*   [Configuring Your Development Environment](#h3ConfigDE)
*   [Having Problems Configuring Your Development Environment?](#h3Problems)

*   [About the Dataset](#h2About)

*   [YOLOv5 Label Format](#h3Label)
*   [Vehicles-OpenImages Dataset](#h3Vehicles)
*   [Udacity Self-Driving Car Dataset](#h3Udacity)
*   [Selecting the Model](#h3Selecting)

*   [YOLOv5 Training](#h2Training)

*   [Downloading the Vehicles-OpenImages Dataset](#h3Downloading)
*   [Configuration Setup](#h3Setup)
*   [YOLOv5 Training Hyperparameters and Model Configuration](#h3Hyperparameters)
*   [Training the YOLOv5s Model](#h3Model)
*   [Visualizing Model Artifacts](#h3Results)
*   [Freeze Initial Layers and Fine-Tune Remaining Layers](#h3Freeze)

*   [Summary](#h2Summary)

*   [Citation Information](#h3Citation)

* * *

With the help of Deep Learning, we all know that the field of Computer Vision has proliferated in the last decade. As a result, so many prevalent computer vision problems like image classification, object detection, and segmentation having real industrial use-case started to achieve accuracy like never before. A new benchmark was set every year from 2012. And today, we will look at object detection from a practical perspective.

A custom, annotated image dataset is vital for training the YOLOv5 object detector. It allows us to train the model on specific objects of interest, leading to a detector tailored to our requirements.

[Roboflow](https://roboflow.com/?ref=pyimagesearch) has free tools for each stage of the computer vision pipeline that will streamline your workflows and supercharge your productivity.

Sign up or Log in to your [Roboflow account](https://roboflow.com/?ref=pyimagesearch) to access state of the art dataset libaries and revolutionize your computer vision pipeline.

You can start by choosing your own datasets or using our [PyimageSearch’s assorted library of useful datasets](https://universe.roboflow.com/pyimagesearch?ref=pyimagesearch).

Bring data in any of 40+ formats to [Roboflow](https://roboflow.com/?ref=pyimagesearch), train using any state-of-the-art model architectures, deploy across multiple platforms (API, NVIDIA, browser, iOS, etc), and connect to applications or 3rd party tools.

![Image 1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20393'%3E%3C/svg%3E)

Object detection has various state-of-the-art architectures that can be used off-the-shelf on real-world datasets to detect objects with reasonable accuracy. The only condition is that the test dataset has the same classes as the pre-trained detector.

However, the application you often build might have objects that differ from the pretrained object detector classes. For example, the dataset distribution is very different from where the dataset detector was trained. In such a scenario, we often use the concept of transfer learning, where we use the pre-trained detector and fine-tune it on the newer dataset. In today’s tutorial, you will learn to train the pretrained YOLOv5 object detector on a custom dataset without writing much code.

We will not go into the theoretical details of the YOLOv5 object detector; however, you can check our [Introduction to the YOLO Family](https://pyimg.co/dgbvi) blog post, where we cover some ground around it.

This lesson is the last in our 7-part series on **YOLO**:

1.  [_Introduction to the YOLO Family_](https://pyimg.co/dgbvi)
2.  [_Understanding a Real-Time Object Detection Network: You Only Look Once (YOLOv1)_](https://pyimg.co/3cpmz)
3.  [_A Better, Faster, and Stronger Object Detector (YOLOv2)_](https://pyimg.co/ehaox)
4.  [_Mean Average Precision (mAP) Using the COCO Evaluator_](https://pyimg.co/nwoka)
5.  [_An Incremental Improvement with Darknet-53 and Multi-Scale Predictions (YOLOv3)_](https://pyimg.co/8xfpg)
6.  [_Achieving Optimal Speed and Accuracy in Object Detection (YOLOv4)_](https://pyimg.co/c6kiu)
7.  **_[Training the YOLOv5 Object Detector on a Custom Dataset](https://pyimg.co/fq0a3)_ (today’s tutorial)**

**To learn how to train a YOLOv5 object detector on a custom dataset****_, just keep reading._**

![Image 2](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20128%20128'%3E%3C/svg%3E)

#### Looking for the source code to this post?

[Jump Right To The Downloads Section](#download-the-code)

* * *

[**Training YOLOv5 Object Detector on a Custom Dataset**](#TOC)
---------------------------------------------------------------

In 2020, Glenn Jocher, the founder and CEO of [**Ultralytics**](https://ultralytics.com/), released its open-source implementation of YOLOv5 on [**GitHub**](https://github.com/ultralytics/yolov5). YOLOv5 offers a family of object detection architectures pre-trained on the MS COCO dataset.

Today, YOLOv5 is one of the official state-of-the-art models with tremendous support and is easier to use in production. The best part is that YOLOv5 is natively implemented in PyTorch, eliminating the Darknet framework’s limitations (based on C programming language).

This massive change of YOLO to the PyTorch framework made it easier for the developers to modify the architecture and export to many deployment environments straightforwardly. And not to forget, YOLOv5 is one of the official state-of-the-art models hosted in the Torch Hub showcase.

**Table 1** shows the performance (mAP) and speed (FPS) benchmarks of five YOLOv5 variants on the MS COCO validation dataset at 640×640 image resolution on Volta 100 GPU. All five models were trained on the MS COCO training dataset. The model benchmarks are shown in ascending order starting with YOLOv5n (i.e., the nano variant having the smallest model footprint to the largest model, YOLOv5x).

<table><tbody><tr><td><span><br></span></td><td><span><br></span></td><td><span><br></span></td><td><span><br></span></td><td><span><strong>Speed</strong><br></span></td><td><span><strong>Speed</strong><br></span></td><td><span><strong>Speed</strong><br></span></td><td><span><br></span></td><td><span><br></span></td></tr><tr><td><span><br></span></td><td><span><strong>size<br></strong></span></td><td><span><strong>mAPval<br></strong></span></td><td><span><strong>mAPval<br></strong></span></td><td><span><strong>CPU b1<br></strong></span></td><td><span><strong>V100 b1<br></strong></span></td><td><span><strong>V100 b32<br></strong></span></td><td><span><strong>params<br></strong></span></td><td><span><strong>FLOPs<br></strong></span></td></tr><tr><td><span><strong>Model</strong><br></span></td><td><span><strong>(pixels)<br></strong></span></td><td><span><strong>0.5:0.95<br></strong></span></td><td><span><strong>0.5<br></strong></span></td><td><span><strong>(ms)<br></strong></span></td><td><span><strong>(ms)<br></strong></span></td><td><span><strong>(ms)<br></strong></span></td><td><span><strong>(M)<br></strong></span></td><td><span><strong>@640 (B)</strong></span></td></tr><tr><td><span>YOLOv5n<br></span></td><td><span>640<br></span></td><td><span>28.4<br></span></td><td><span>46.0<br></span></td><td><span>45<br></span></td><td><span>6.3<br></span></td><td><span>0.6<br></span></td><td><span>1.9<br></span></td><td><span>4.5<br></span></td></tr><tr><td><span>YOLOv5s<br></span></td><td><span>640<br></span></td><td><span>37.2<br></span></td><td><span>56.0<br></span></td><td><span>98<br></span></td><td><span>6.4<br></span></td><td><span>0.9<br></span></td><td><span>7.2<br></span></td><td><span>16.5<br></span></td></tr><tr><td><span>YOLOv5m<br></span></td><td><span>640<br></span></td><td><span>45.2<br></span></td><td><span>63.9<br></span></td><td><span>224<br></span></td><td><span>8.2<br></span></td><td><span>1.7<br></span></td><td><span>21.2<br></span></td><td><span>49<br></span></td></tr><tr><td><span>YOLOv5l<br></span></td><td><span>640<br></span></td><td><span>48.8<br></span></td><td><span>67.2<br></span></td><td><span>430<br></span></td><td><span>10.1<br></span></td><td><span>2.7<br></span></td><td><span>46.5<br></span></td><td><span>109.1<br></span></td></tr><tr><td><span>YOLOv5x<br></span></td><td><span>640<br></span></td><td><span>50.7<br></span></td><td><span>68.9<br></span></td><td><span>766<br></span></td><td><span>12.1<br></span></td><td><span>4.8<br></span></td><td><span>86.7<br></span></td><td><span>205.7<br></span></td></tr><tr><td colspan="9"><strong>Table 1:</strong> Performance and Speed benchmarks of five YOLOv5 variants on the MS COCO dataset.<br></td></tr></tbody></table>

Today, we’ll learn how to harness the power of YOLOv5 in the PyTorch framework by transfer learning it on a custom dataset!

* * *

### [**Configuring Your Development Environment**](#TOC)

To follow this guide, you need to clone the [Ultralytics YOLOv5 repository](https://github.com/ultralytics/yolov5) and `pip` install all the necessary packages from `requirements.txt`.

Luckily, to run the YOLOv5 training, you can just do a pip install on the `requirements.txt` file, which means that all the libraries are pip-installable!

$ git clone https://github.com/ultralytics/yolov5.git #clone repo
$ cd yolov5/
$ pip install -r requirements.txt #install dependencies

* * *

### [**Having Problems Configuring Your Development Environment?**](#TOC)

[![Image 3](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20334'%3E%3C/svg%3E)](https://www.pyimagesearch.com/pyimagesearch-university/)

**Figure 1:** Having trouble configuring your dev environment? Want access to pre-configured Jupyter Notebooks running on Google Colab? Be sure to join [PyImageSearch University](https://www.pyimagesearch.com/pyimagesearch-university/) — you’ll be up and running with this tutorial in a matter of minutes.

All that said, are you:

*   Short on time?
*   Learning on your employer’s administratively locked system?
*   Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
*   **Ready to run the code** **_right now_** **on your Windows, macOS, or Linux system?**

Then join [PyImageSearch University](https://www.pyimagesearch.com/pyimagesearch-university/) today!

**Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides that are** **_pre-configured_** **to run on Google Colab’s ecosystem right in your web browser!** No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

* * *

[**About the Dataset**](#TOC)
-----------------------------

For today’s experiment, we will be training the YOLOv5 model on two different datasets, namely the [Udacity Self-driving Car dataset](https://public.roboflow.com/object-detection/self-driving-car) and the [Vehicles-OpenImages dataset](https://public.roboflow.com/object-detection/vehicles-openimages).

These datasets are public, but we download them from [Roboflow](https://roboflow.com/), which provides a great platform to train your models with various datasets in the Computer Vision domain. Even more interesting is that you can download the datasets in multiple formats like COCO JSON, YOLO Darknet TXT, and YOLOv5 PyTorch. This saves time for writing helper functions for converting the ground-truth annotations to the format required by the model.

* * *

### [**YOLOv5 Label Format**](#TOC)

Since we will train the YOLOv5 PyTorch model, we will download the datasets in YOLOv5 format. The ground-truth annotation format of YOLOv5 is pretty simple (an example is shown in **Figure 2**), so you could write a script on your own that does that for you. There is one `text` file with a single line for each bounding box for each image. For example, if there are four objects in one image, the `text` file would have four rows containing the class label and bounding box coordinates. The format of each row is

`class_id center_x center_y width height`

![Image 4](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20435'%3E%3C/svg%3E)

**Figure 2:** YOLO bounding box format example ([source](https://roboflow.com/formats/yolov5-pytorch-txt)).

where fields are space-delimited, and the coordinates are normalized from `0` to `1`. To convert to normalized `xywh` from pixel values, divide x & box width by the image’s width and divide y & box height by the image’s height.

* * *

### [**Vehicles-OpenImages Dataset**](#TOC)

This dataset contains only 627 images of various vehicle classes for object detection like _Car_, _Bus_, _Ambulance_, _Motorcycle_, and _Truck_. These images are derived from the [Open Images](https://arxiv.org/pdf/1811.00982.pdf) open-source computer vision datasets. The dataset falls under the [Creative Commons License](https://creativecommons.org/licenses/by/4.0/), which allows you to share and adapt the dataset, and even use it commercially.

**Figure 3** shows some sample images from the dataset with ground-truth bounding boxes annotated in green color.

![Image 5](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20645%20500'%3E%3C/svg%3E)

**Figure 3:** Sample images from the Vehicles-OpenImages dataset with ground-truth annotations ([source](https://public.roboflow.com/object-detection/vehicles-openimages)).

There are 1194 regions of interest (objects) in 627 images, meaning there are at least 1.9 objects per image. Based on the heuristic shown in **Figure 4**, the car class contributes to more than 50% of the objects. In contrast, the remaining classes: bus, truck, motorcycle, and ambulance, are under-represented relative to the car class.

![Image 6](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20489'%3E%3C/svg%3E)

**Figure 4:** Class Distribution of Vehicles Open Image Dataset showing that more than half of the objects belong to the `car` class.

* * *

### [**Udacity Self-Driving Car Dataset**](#TOC)

Please note that we will not train the YOLOv5 model on this dataset. Instead, we have identified this great dataset for you all as an exercise so that once you are done learning from this tutorial, you can use this dataset for training the object detector.

This dataset is derived from the [original Udacity Self-Driving Car Dataset](https://github.com/udacity/self-driving-car/tree/master/annotations). Unfortunately, the original dataset was missing labels for thousands of pedestrians, bikers, cars, and traffic lights. Hence, Roboflow managed to re-label the dataset to correct errors and omissions.

**Figure 5** shows some examples from the dataset and labels missing from the original dataset annotated by Roboflow. If you have worked with autonomous driving urban-scene datasets like [Cityscapes](https://www.cityscapes-dataset.com/), [ApolloScape](http://apolloscape.auto/), and [Berkeley DeepDrive](https://bdd-data.berkeley.edu/), you will find this dataset very similar to those.

![Image 7](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20441'%3E%3C/svg%3E)

**Figure 5**: Sample images from the Udacity Self-Driving Car dataset.

The dataset contains 97,942 labels across 11 classes and 15,000 images. The dataset is available on Roboflow in two different fashions: images with `1920x1200` (download size `~3.1` GB) and a downsampled version with `512x512` (download size `~580` MB) suitable for most. We will use the downsampled version since it is smaller in size and fits our requirements in terms of the network.

Like the previous Vehicles-OpenImages dataset, this dataset has the most number of objects belonging to the `car` class (more than 60% of the total objects). **Figure 6** shows the distribution of classes in the Udacity Self-Driving Car dataset:

![Image 8](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20474'%3E%3C/svg%3E)

**Figure 6:** Class distribution of Udacity Self-Driving Car dataset in which 66% of the objects belong to the `car` class.

* * *

### [**Selecting the Model**](#TOC)

**Figure 7** shows five YOLOv5 variants starting with the most miniature YOLOv5 nano model built for running on mobile and embedded devices to YOLOv5 XLarge on the other end of the spectrum. For today’s experiment, we would leverage the base model YOLOv5s, which provides a nice balance between accuracy and speed.

![Image 9](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20244'%3E%3C/svg%3E)

**Figure 7:** YOLOv5 variants starting with YOLOv5 Nano to YOLOv5 XLarge.

* * *

[**YOLOv5 Training**](#TOC)
---------------------------

This section is the heart of today’s tutorial, where we will cover most of the things, starting with

*   Downloading the dataset
*   Creating the data configuration
*   Training the YOLOv5 model
*   Visualizing the YOLOv5 model artifacts
*   Quantitative results
*   Freeze initial layers and fine-tune the remaining layers
*   Results

* * *

### [**Downloading the Vehicles-OpenImages Dataset**](#TOC)

\# Download the vehicles-open image dataset
!mkdir vehicles\_open\_image
%cd vehicles\_open\_image
!curl -L "https://public.roboflow.com/ds/2Tb6yXY8l8?key=Eg82WpxUEr" > vehicles.zip 
!unzip vehicles.zip
!rm vehicles.zip

On **Lines 2 and 3**, we create the `vehicles_open_image` directory and `cd` into the directory where we download the dataset. Then, on **Line 4**, we use the `curl` command and pass the dataset URL we got from [here](https://public.roboflow.com/object-detection/vehicles-openimages/1). Finally, we unzip the dataset and remove the zip file on **Lines 5 and 6**.

Let’s look at the contents of the `vehicles_open_image` folder:

$tree /content/vehicles\_open\_image -L 2
/content/vehicles\_open\_image
├── data.yaml
├── README.dataset.txt
├── README.roboflow.txt
├── test
│   ├── images
│   └── labels
├── train
│   ├── images
│   └── labels
└── valid
    ├── images
    └── labels

9 directories, 3 files

The parent directory has three files, out of which only `data.yaml` is essential, and three sub-directories:

*   `data.yaml`: It has the data-related configurations like the train and valid data directory path, the total number of classes in the dataset, and the name of each class
*   `train`: Training images along with training labels
*   `valid`: Validation images with annotations
*   `test`: Test images and labels. Assessing your model’s performance becomes easy if test data with labels are available.

* * *

### [**Configuration Setup**](#TOC)

Next, we will edit the `data.yaml` file to have the `path` and absolute path for `train` and `valid` images.

\# Create configuration
import yaml
config = {'path': '/content/vehicles\_open\_image',
         'train': '/content/vehicles\_open\_image/train',
         'val': '/content/vehicles\_open\_image/valid',
         'nc': 5,
         'names': \['Ambulance', 'Bus', 'Car', 'Motorcycle', 'Truck'\]}
 
with open("data.yaml", "w") as file:
   yaml.dump(config, file, default\_flow\_style=False)

On **Line 2**, we import the `yaml` module, which would allow us to save the `data.yaml` configuration file in `.yaml` format. Then from **Lines 3-7**, we define the data path, train, validation, number of classes, and class names in a `config` variable. So `config` is a dictionary.

Finally, on **Lines 9 and 10**, we open the existing `data.yaml` file that was downloaded along with the dataset, overwrite it with the contents in `config`, and store it on the disk.

* * *

### [**YOLOv5 Training Hyperparameters and Model Configuration**](#TOC)

YOLOv5 has about 30 hyperparameters used for various training settings. These are defined in `hyp.scratch-low.yaml` for low-augmentation COCO training from scratch, placed in the `/data` directory. The training data hyperparameters are shown below, which are very important for producing good results, so make sure to initialize these values properly before starting the training. For this tutorial, we would simply use the default values, which are optimized for YOLOv5 COCO training from scratch.

As you can see, it has `learning rate`, `weight_decay`, and `iou_t` (IoU training threshold), to name a few, and some data augmentation hyperparameters like `translate`, `scale`, `mosaic`, `mixup`, and `copy_paste`. A `mixup:0.0` means the mixup data augmentation should not be applied.

lr0: 0.01  # initial learning rate (SGD=1E-2, Adam=1E-3)
lrf: 0.01  # final OneCycleLR learning rate (lr0 \* lrf)
momentum: 0.937  # SGD momentum/Adam beta1
weight\_decay: 0.0005  # optimizer weight decay 5e-4
warmup\_epochs: 3.0  # warmup epochs (fractions ok)
warmup\_momentum: 0.8  # warmup initial momentum
warmup\_bias\_lr: 0.1  # warmup initial bias lr
box: 0.05  # box loss gain
cls: 0.5  # cls loss gain
cls\_pw: 1.0  # cls BCELoss positive\_weight
obj: 1.0  # obj loss gain (scale with pixels)
obj\_pw: 1.0  # obj BCELoss positive\_weight
iou\_t: 0.20  # IoU training threshold
anchor\_t: 4.0  # anchor-multiple threshold
# anchors: 3  # anchors per output layer (0 to ignore)
fl\_gamma: 0.0  # focal loss gamma (efficientDet default gamma=1.5)
hsv\_h: 0.015  # image HSV-Hue augmentation (fraction)
hsv\_s: 0.7  # image HSV-Saturation augmentation (fraction)
hsv\_v: 0.4  # image HSV-Value augmentation (fraction)
degrees: 0.0  # image rotation (+/- deg)
translate: 0.1  # image translation (+/- fraction)
scale: 0.5  # image scale (+/- gain)
shear: 0.0  # image shear (+/- deg)
perspective: 0.0  # image perspective (+/- fraction), range 0-0.001
flipud: 0.0  # image flip up-down (probability)
fliplr: 0.5  # image flip left-right (probability)
mosaic: 1.0  # image mosaic (probability)
mixup: 0.0  # image mixup (probability)
copy\_paste: 0.0  # segment copy-paste (probability)

Next, you can briefly look at the structure of the `YOLOv5s` network architecture, though you would hardly modify the model configuration file, unlike the training data hyperparameters. It has `nc` set to `80` for MS COCO dataset, `backbone` for feature extraction, followed by `head` for detection.

\# Parameters
nc: 80  # number of classes
depth\_multiple: 0.33  # model depth multiple
width\_multiple: 0.50  # layer channel multiple
anchors:
  - \[10,13, 16,30, 33,23\]  # P3/8
  - \[30,61, 62,45, 59,119\]  # P4/16
  - \[116,90, 156,198, 373,326\]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # \[from, number, module, args\]
  \[\[-1, 1, Conv, \[64, 6, 2, 2\]\],  # 0-P1/2
   \[-1, 1, Conv, \[128, 3, 2\]\],  # 1-P2/4
   \[-1, 3, C3, \[128\]\],
   \[-1, 1, Conv, \[256, 3, 2\]\],  # 3-P3/8
   \[-1, 6, C3, \[256\]\],
   \[-1, 1, Conv, \[512, 3, 2\]\],  # 5-P4/16
   \[-1, 9, C3, \[512\]\],
   \[-1, 1, Conv, \[1024, 3, 2\]\],  # 7-P5/32
   \[-1, 3, C3, \[1024\]\],
   \[-1, 1, SPPF, \[1024, 5\]\],  # 9
  \]

# YOLOv5 v6.0 head
head:
  \[\[-1, 1, Conv, \[512, 1, 1\]\],
   \[-1, 1, nn.Upsample, \[None, 2, 'nearest'\]\],
   \[\[-1, 6\], 1, Concat, \[1\]\],  # cat backbone P4
   \[-1, 3, C3, \[512, False\]\],  # 13

   \[-1, 1, Conv, \[256, 1, 1\]\],
   \[-1, 1, nn.Upsample, \[None, 2, 'nearest'\]\],
   \[\[-1, 4\], 1, Concat, \[1\]\],  # cat backbone P3
   \[-1, 3, C3, \[256, False\]\],  # 17 (P3/8-small)

   \[-1, 1, Conv, \[256, 3, 2\]\],
   \[\[-1, 14\], 1, Concat, \[1\]\],  # cat head P4
   \[-1, 3, C3, \[512, False\]\],  # 20 (P4/16-medium)

   \[-1, 1, Conv, \[512, 3, 2\]\],
   \[\[-1, 10\], 1, Concat, \[1\]\],  # cat head P5
   \[-1, 3, C3, \[1024, False\]\],  # 23 (P5/32-large)

   \[\[17, 20, 23\], 1, Detect, \[nc, anchors\]\],  # Detect(P3, P4, P5)
  \]

* * *

### [**Training the YOLOv5s Model**](#TOC)

We are almost ready to train the YOLOv5 model, and as discussed above, we will be training the YOLOv5s model. However, before we run the training, let’s define a few parameters:

SIZE = 640
BATCH\_SIZE = 32
EPOCHS = 20
MODEL = "yolov5s"
WORKERS = 1
PROJECT = "vehicles\_open\_image\_pyimagesearch"
RUN\_NAME = f"{MODEL}\_size{SIZE}\_epochs{EPOCHS}\_batch{BATCH\_SIZE}\_small"

We define a few standard model parameters:

*   `SIZE`: Image size or network input while training. The images will be resized to this value before being fed to the network. The preprocessing pipeline would resize them to `640` pixels.
*   `BATCH_SIZE`: Number of images fed as a single batch into the network for a forward pass. You can modify it according to the GPU memory available. We have set it to `32`.
*   `EPOCHS`: Number of times we want to train the model on the full dataset.
*   `MODEL`: The base model we want to use for training. We use the small model `yolov5s` from the YOLOv5 family.
*   `WORKERS`: Maximum Dataloader workers to be used.
*   `PROJECT`: This will create a project directory inside the current directory (`yolov5`).
*   `RUN_NAME`: Each time you run this model, it will create a sub-directory under the project directory, which would have a lot of information on the model like weights, sample input images, a few validation predictions outputs, metrics plot, etc.

!python train.py --img {SIZE}\\
               --batch {BATCH\_SIZE}\\
               --epochs {EPOCHS}\\
               --data ../vehicles\_open\_image/data.yaml\\
               --weights {MODEL}.pt\\
               --workers {WORKERS}\\
               --project {PROJECT}\\
               --name {RUN\_NAME}\\
               --exist-ok

If there are no errors, the training will start as shown below. The logs indicate that the YOLOv5 model would train with Torch version 1.11 on a Tesla T4 GPU; along with that, it also shows the `hyperparameters` initialized.

The `yolov5s.pt` weights are downloaded, which means that the YOLOv5s model is initialized with the parameters trained with the MS COCO dataset. Finally, we can see that two epochs have been completed with a `mAP@0.5=0.237`.

github: up to date with https://github.com/ultralytics/yolov5 ✅
YOLOv5 ? v6.1-236-gdcf8073 Python-3.7.13 torch-1.11.0+cu113 CUDA:0 (Tesla T4, 15110MiB)
hyperparameters: lr0=0.01, lrf=0.01, momentum=0.937, weight\_decay=0.0005, warmup\_epochs=3.0, warmup\_momentum=0.8, warmup\_bias\_lr=0.1, box=0.05, cls=0.5, cls\_pw=1.0, obj=1.0, obj\_pw=1.0, iou\_t=0.2, anchor\_t=4.0, fl\_gamma=0.0, hsv\_h=0.015, hsv\_s=0.7, hsv\_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy\_paste=0.0
Weights & Biases: run 'pip install wandb' to automatically track and visualize YOLOv5 ? runs (RECOMMENDED)
TensorBoard: Start with 'tensorboard --logdir parking\_lot\_pyimagesearch', view at http://localhost:6006/
Downloading https://ultralytics.com/assets/Arial.ttf to /root/.config/Ultralytics/Arial.ttf...
100% 755k/755k \[00:00<00:00, 18.0MB/s\]
YOLOv5 temporarily requires wandb version 0.12.10 or below. Some features may not work as expected.
Downloading https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt to yolov5s.pt...
100% 14.1M/14.1M \[00:00<00:00, 125MB/s\]

Overriding model.yaml nc=80 with nc=5
Logging results to parking\_lot\_pyimagesearch/yolov5s\_size640\_epochs20\_batch32\_simple
Starting training for 20 epochs...

     Epoch   gpu\_mem       box       obj       cls    labels  img\_size
      0/19     7.36G   0.09176   0.03736   0.04355        31       640: 100% 28/28 \[00:28<00:00,  1.03s/it\]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 4/4 \[00:03<00:00,  1.04it/s\]
                 all        250        454      0.352      0.293      0.185      0.089

     Epoch   gpu\_mem       box       obj       cls    labels  img\_size
      1/19     8.98G   0.06672   0.02769   0.03154        45       640: 100% 28/28 \[00:25<00:00,  1.09it/s\]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 4/4 \[00:05<00:00,  1.45s/it\]
                 all        250        454      0.271      0.347      0.237     0.0735

Voila! With this, you have learned to train an object detector on a custom dataset you downloaded from Roboflow. Isn’t that amazing?

Even more exciting is that YOLOv5 logs the model artifacts inside the `runs` directory, which we will look at in the next section.

Once the training is complete, you will see the output similar to the one shown below:

​​    Epoch   gpu\_mem       box       obj       cls    labels  img\_size
     19/19     7.16G   0.02747   0.01736  0.004772        46       640: 100% 28/28 \[01:03<00:00,  2.27s/it\]
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 4/4 \[00:05<00:00,  1.42s/it\]
                 all        250        454      0.713      0.574      0.606      0.416

20 epochs completed in 0.386 hours.
Optimizer stripped from runs/train/exp/weights/last.pt, 14.5MB
Optimizer stripped from runs/train/exp/weights/best.pt, 14.5MB

Validating runs/train/exp/weights/best.pt...
Fusing layers... 
Model summary: 213 layers, 7023610 parameters, 0 gradients, 15.8 GFLOPs
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 4/4 \[00:08<00:00,  2.23s/it\]
                 all        250        454      0.715      0.575      0.606      0.416
           Ambulance        250         64      0.813      0.814      0.853      0.679
                 Bus        250         46      0.771      0.652      0.664       0.44
                 Car        250        238      0.653      0.496      0.518      0.354
          Motorcycle        250         46      0.731      0.478      0.573      0.352
               Truck        250         60      0.608      0.433      0.425      0.256

The above results show that the YOLOv5s model achieved an mAP of `0.606@0.5` IoU and `0.416@0.5:0.95` IoU in all classes. It also indicates class-wise mAP, and the model achieved the best score for the Ambulance class (i.e., `0.853 mAP@0.5` IoU). The model took `23.16` minutes to complete the training for `20` epochs on a Tesla T4 or Tesla K80.

* * *

### [**Visualizing Model Artifacts**](#TOC)

Now that we are done training our model, let’s look at the results generated inside the `yolov5/vehicles_open_pyimagesearch_model` directory.

All training results are logged by default to `yolov5/runs/train` with a new incrementing directory created for each run as `runs/train/exp`, `runs/train/exp1`, etc. However, while training the model, we passed the `PROJECT` and the `RUN_NAME`, so in this case, it does not create the default directory to log the training results. Hence, in this experiment `runs` is `parking_lot_pyimagesearch`.

Next, let’s look at the files created in the experiment.

$tree parking\_lot\_pyimagesearch/yolov5s\_size640\_epochs20\_batch32\_small/
parking\_lot\_pyimagesearch/yolov5s\_size640\_epochs20\_batch32\_small/
├── confusion\_matrix.png
├── events.out.tfevents.1652810418.f70b01be1223.864.0
├── F1\_curve.png
├── hyp.yaml
├── labels\_correlogram.jpg
├── labels.jpg
├── opt.yaml
├── P\_curve.png
├── PR\_curve.png
├── R\_curve.png
├── results.csv
├── results.png
├── train\_batch0.jpg
├── train\_batch1.jpg
├── train\_batch2.jpg
├── val\_batch0\_labels.jpg
├── val\_batch0\_pred.jpg
├── val\_batch1\_labels.jpg
├── val\_batch1\_pred.jpg
├── val\_batch2\_labels.jpg
├── val\_batch2\_pred.jpg
└── weights
    ├── best.pt
    └── last.pt

1 directory, 23 files

On **Line 1**, we use the `tree` command followed by the `PROJECT` and `RUN_NAME`, displaying various evaluation metrics and weights files for the trained object detector. As we can observe, it has a precision curve, recall curve, precision-recall curve, confusion matrix, prediction on validation images, and finally, the weights file in PyTorch format.

Let’s now look at a few images from the `runs` directory.

**Figure 8** shows the training images batch with Mosaic data augmentation. There are 16 images clubbed together; if we pick one image from the 3rd row 1st column, then we can see that the image is a combination of four different images. We explain the concept of Mosaic data augmentation in the YOLOv4 post, so do check that out if you haven’t already.

![Image 10](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20500%20500'%3E%3C/svg%3E)

**Figure 8:** Training images batch with mosaic data augmentation.

Next, we look at the `results.png`, which comprises training and validation loss for bounding box, objectness, and classification. It also has the metrics: precision, recall, `mAP@0.5`, and `mAP@0.5:0.95` for training (**Figure 9**).

![Image 11](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20350'%3E%3C/svg%3E)

**Figure 9:** Loss and Evaluation Metrics Plot.

**Figure 10** shows the ground-truth images and the YOLOv5s model prediction on the Vehicles-OpenImages dataset. From the two images below, it is clear that the model did a great job detecting the objects. Unfortunately, the model failed to detect a bike in the second image and a car in the sixth image. And misclassified a truck as a car in the first image, but this was a tough one to crack as it’s even difficult for the human to predict it correctly. But overall, it did great on these images.

![Image 12](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%201080%20542'%3E%3C/svg%3E)

![Image 13](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%201080%20539'%3E%3C/svg%3E)

**Figure 10:** Ground-truth images (_top_) and YOLOv5s model prediction (_bottom_) fine-tuned with all layers.

* * *

### [**Freeze Initial Layers and Fine-Tune Remaining Layers**](#TOC)

We learned how to train an object detector that was pretrained on the MS COCO dataset, which means we fine-tuned the model parameters (all layers) on the Vehicles-OpenImages dataset. But the question is, do we need to train all the model layers on the new dataset? Maybe not, since the pretrained model has been trained on a large, well-curated MS COCO dataset.

The benefit of freezing layers when often fine-tuning a model on a custom dataset reduces the training time. If the custom dataset is not too complex, then you can expect, if not the same, but comparable accuracies. When we compare the two models’ training times, you will see for yourself.

In this section, we will again train or fine-tune the YOLOv5s model on the Vehicles-OpenImages dataset but freeze the initial `11` layers of the network, unlike before, where we fine-tune all the detector layers. Thanks to the creators of YOLOv5, freezing the model layers is very easy. But, first, we must pass the `--freeze` argument with the layer numbers we would like to freeze in the model.

Let’s now train the model by executing the `train.py` script. First, we change the `--name`, that is, the run name to `freeze_layers`, pass the `--freeze` parameter, and all other parameters are the same.

!python train.py --img {SIZE}\\
               --batch {BATCH\_SIZE}\\
               --epochs {EPOCHS}\\
               --data ../vehicles\_open\_image/data.yaml\\
               --weights {MODEL}.pt\\
               --workers {WORKERS}\\
               --project {PROJECT}\\
               --name freeze\_layers\\
               --exist-ok\\
               --freeze 0 1 2 3 4 5 6 7 8 9 10

The output below is produced when you run the training script; as you can see, the first 11 layers of the network are shown with a `freezing` prefix which means that the parameters (weights and biases) for these layers would remain constant. At the same time, the remaining 15 layers would be fine-tuned on the custom dataset.

That being said, let us look at the results now!

freezing model.0.conv.weight
freezing model.0.bn.weight
freezing model.0.bn.bias
freezing model.1.conv.weight
freezing model.1.bn.weight
freezing model.1.bn.bias
.
.
.
.
freezing model.10.conv.weight
freezing model.10.bn.weight
freezing model.10.bn.bias

We can observe from the below output that `20` epochs were completed in only `0.158` hours (i.e., `9.48` mins), and the model we fine-tuned with none of the layers frozen took `23.16` minutes. Wow! That’s close to a 2.5X reduction in time.

But hold on, let’s look at the mAP of this model shown for all classes and class-wise.

With 11 layers frozen, the model achieved `0.551 mAP@0.5` IoU and `0.336 mAP@0.5:0.95` IoU. There is definitely a difference between the accuracies of the two models but not too significant.

20 epochs completed in 0.158 hours.
Optimizer stripped from parking\_lot\_pyimagesearch/freeze\_layers/weights/last.pt, 14.5MB
Optimizer stripped from parking\_lot\_pyimagesearch/freeze\_layers/weights/best.pt, 14.5MB

Validating parking\_lot\_pyimagesearch/freeze\_layers/weights/best.pt...
Fusing layers... 
Model summary: 213 layers, 7023610 parameters, 0 gradients, 15.8 GFLOPs
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 4/4 \[00:04<00:00,  1.25s/it\]
                 all        250        454      0.579      0.585      0.551      0.336
           Ambulance        250         64      0.587      0.688      0.629      0.476
                 Bus        250         46      0.527      0.696      0.553      0.304
                 Car        250        238      0.522      0.462      0.452      0.275
          Motorcycle        250         46      0.733      0.478      0.587      0.311
               Truck        250         60      0.527        0.6      0.533      0.313

Finally, in **Figure 11**, we can see the detector prediction on the validation images. The results clearly show that they are not as good as the detector trained with all layers. For example, it misses the object in the first image, in the second image, it misclassifies a `motorcycle` with a `car`, in the fourth image fails to detect the `car`, and even in the sixth image, it sees only three `cars`. Maybe training for a bit more epochs or freezing fewer layers could improve the detector’s performance.

![Image 14](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20350'%3E%3C/svg%3E)

**Figure 11:** Ground-truth images (_top_) and YOLOv5s model prediction (_bottom_) fine-tuned with 11 layers frozen.

* * *

### What's next? We recommend [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend).

**Course information:**  
84 total classes • 114+ hours of on-demand code walkthrough videos • Last updated: February 2024  
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

**I strongly believe that if you had the right teacher you could _master_ computer vision and deep learning.**

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s _not_ the case.

All you need to master computer vision and deep learning is for someone to explain things to you in _simple, intuitive_ terms. _And that’s exactly what I do_. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to _successfully_ and _confidently_ apply computer vision to your work, research, and projects. Join me in computer vision mastery.

**Inside PyImageSearch University you'll find:**

*   ✓ **84 courses** on essential computer vision, deep learning, and OpenCV topics
*   ✓ **84 Certificates** of Completion
*   ✓ **114+ hours** of on-demand video
*   ✓ **Brand new courses released _regularly_**, ensuring you can keep up with state-of-the-art techniques
*   ✓ **Pre-configured Jupyter Notebooks in Google Colab**
*   ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
*   ✓ Access to **centralized code repos for _all_ 536+ tutorials** on PyImageSearch
*   ✓ **Easy one-click downloads** for code, datasets, pre-trained models, etc.
*   ✓ **Access** on mobile, laptop, desktop, etc.

[Click here to join PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend)

* * *

[**Summary**](#TOC)
-------------------

Congratulations on making it this far! Let’s quickly summarize what we learned in this tutorial.

1.  We started by giving a brief overview of YOLOv5 and discussed the performance and speed benchmarks of YOLOv5 variants.
2.  Then we discussed the two datasets: Vehicles-OpenImages dataset and Udacity Self-Driving Car dataset. Along with that we also covered the YOLOv5 ground-truth annotation format.
3.  After finalizing the YOLOv5 model variant for training we dived into the hands-on part of the tutorial where we covered aspects like downloading the dataset, creating `configuration.yaml` for the given data, and training and visualizing the YOLOv5 model artifacts.
4.  Finally, we trained the YOLOv5 model for the second time but with the initial 11 layers of the model frozen and compared the results with the YOLOv5 model trained fully.

* * *

### [**Citation Information**](#TOC)

**Sharma, A.** “Training the YOLOv5 Object Detector on a Custom Dataset,” _PyImageSearch_, D. Chakraborty, P. Chugh, A. R. Gosthipaty, S. Huot, K. Kidriavsteva, R. Raha, and A. Thanki, eds., 2022, [https://pyimg.co/fq0a3](https://pyimg.co/fq0a3)

@incollection{Sharma\_2022\_Custom\_Dataset,
  author = {Aditya Sharma},
  title = {Training the {YOLOv5} Object Detector on a Custom Dataset},
  booktitle = {PyImageSearch},
  editor = {Devjyoti Chakraborty and Puneet Chugh and Aritra Roy Gosthipaty and Susan Huot and Kseniia Kidriavsteva and Ritwik Raha and Abhishek Thanki},
  year = {2022},
  note = {https://pyimg.co/fq0a3},
}

* * *

**To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), _simply enter your email address in the form below!_**

![Image 15](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20410%20567'%3E%3C/svg%3E)

#### Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a **FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning.** Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!