How to Train a Custom Dataset with YOLOv5?

Introduction

We have seen some fancy terms for AI and deep learning, such as pre-trained models, transfer learning, etc. Let me educate you with a widely used technology and one of the most important and effective: Transfer learning with YOLOv5.

You Only Look Once, or YOLO is one of the most extensively used deep learning-based object identification methods. Using a custom dataset, this article will show you how to train one of its most recent variations, YOLOv5.

Learning Objectives 

1. This article will focus mainly on training the YOLOv5 model on a custom dataset implementation.
2. We will see what pre-trained models are and see what transfer learningis.
3. We will understand what YOLOv5 is and why we are using version 5 of YOLO.

So, without wasting time, let’s get started with the process

Table of Content

1. Pre-Trained Models
2. Transfer Learning
3. What and Why YOLOv5?
4. Steps Involved In Transfer Learning
5. Implementation
6. Some Challenges That You Can Face
7. Conclusion

Pre-trained Models

You might have heard data scientists use the term “pre-trained model”  widely. After explaining what a deep learning model/network does, I will explain the term. A deep learning model is a model containing various layers stacked together so as to serve a solitary purpose, such as classification, detection, etc. Deep learning networks learn by discovering complicated structures in the data fed to them and saving the weights in a file which are later used to perform similar tasks. Pretrained models are already trained Deep Learning models. What it means is that they are already trained on a huge dataset containing millions of images.

Here is how the TensorFlow website defines pre-trained models: A pre-trained model is a saved network that was previously trained on a large dataset, typically on a large-scale image-classification task.

Some highly optimized and extraordinarily efficient pre-trained models are available on the internet. Different models are used to perform different tasks. Some of the pre-trained models are VGG-16, VGG-19, YOLOv5, YOLOv3, and ResNet 50.

Which model to use depends on the task you want to perform. For example, if I want to perform an object detection task, I will use the YOLOv5 model.

Transfer Learning

Transfer Learning is the most important technique that eases the task of a data scientist. Training a model is a hefty and time-consuming task; if a model is trained from scratch, it usually does not give very good results. Even if we train a model similar to a pre-trained model, it will not perform as effectively, and it can take weeks for a model to train. Instead, we can use the pre-trained models and use the already learned weights by training them on a custom dataset to perform a similar task. These models are highly efficient and refined in terms of architecture and performance, and they have made their way to the top by performing better in different contests. These models are trained on very large amounts of data, making them more diverse in knowledge.

So transfer learning basically means transferring knowledge gained by training the model on previous data to help the model learn better and faster to perform a different but similar task.

For example, using a YOLOv5 for object detection, but the object is something other than the object’s previous data used.

What and Why YOLOv5?

YOLOv5 is a pre-trained model which stands for you only look once version 5 is used for real-time object detection and has proven to be highly efficient in terms of accuracy and inference time. There are other versions of YOLO, but as one would predict, YOLOv5 performs better than other versions. YOLOv5 is fast and easy to use. It is based on the PyTorch framework, which has a larger community than Yolo v4 Darknet.

We will now look at the architecture of YOLOv5.

The structure may look confusing, but it does not matter as we do not have to look at the architecture instead directly use the model and weights.

In transfer learning, we use the custom dataset i.e., the data the model has never seen before OR the data on which the model is not trained. Since the model is already trained on a large dataset, we already have the weights. We can now train the model for a number of epochs on the data we want to work on. Training is required as the model has seen the data for the first time and will require some knowledge in order to perform the task.

Steps Involved in Transfer Learning

Transfer learning is a simple process, and we can do it in a few simple steps:

1. Data preparation
2. The right format for the annotations
3. Change a few layers if you want to
4. Retrain the model for a few iterations
5. Validate/Test

Data Preparation

Data preparation can be time-consuming if your chosen data is a bit large. Data preparation means annotating the images, which is a process where you label the images by making a box around the object in the image. By doing this, the coordinates of the object marked will be saved in a file which will then be fed to the model for training. There are a few websites, such as makesense.ai and roboflow.com, which can help you label the data. 

Here is how you can annotate the data for the YOLOv5 model on makesense.ai.

1. Visit https://www.makesense.ai/. 

2. Click on get started at the bottom right of the screen.

3. Select the images you want to label by clicking on the box highlighted in the center.

Load the images you want to annotate and click on object detection.

4. After loading the images, you will be asked to create labels for your dataset’s different classes.

I am detecting license plates on a vehicle, so the only label I will be using is “License Plate.” You can make more labels by simply hitting enter by clicking the ‘+’ button on the left side of the dialogue box.

After you have created all the labels, click on start project.

If you have missed any labels, you can edit them later by clicking on actions and then edit labels.

5. Start creating a bounding box around the object in the image. This exercise may be a bit fun initially, but with very large data, it can be tiring.

6. After annotating all the images you need to save the file which will contain the coordinates of bounding boxes along with the class.

So you need to head to the actions button and click on export annotations don’t forget to check the ‘A zip package containing files in YOLO format’ option, as this will save the files in the correct format as required in the YOLO model.

7. This is a significant step, so follow it carefully.

After you have all the files and the images, make a folder with any name. Click on the folder and make two more folders with the name images and labels inside the folder. Don’t forget to name the folder the same as above, as the model automatically searches for labels after you feed the training path in the command.

To give you an idea of the folder, I have created a folder named ‘CarsData’ and in that folder made two folders – ‘images’ and ‘labels.’

Inside the two folders, you have to make two more folders named ‘train’ and ‘val’. In the images folder, you can split the images according to your will, but you have to be careful while splitting the label, as the labels should match the images you have split

8. Now make a zip file of the folder and upload it to the drive so that we can use it in colab.

Implementation

We will now come to the implementation part, which is very simple but tricky. If you don’t know what files to change exactly, you won’t be able to train the model on the custom dataset. 

So here are the codes that you should follow to train the YOLOv5 model on a custom dataset

I recommend you use google colab for this tutorial as it also provides GPU which provides faster computations.

1. !git clone https://github.com/ultralytics/yolov5
This will make a copy of the YOLOv5 repository which is a GitHub repository created by ultralytics.

2. cd yolov5
This is a command-line shell command used to change the current working directory to the YOLOv5 directory.

3. !pip install -r requirements.txt
This command will install all the packages and libraries used in training the model.

4. !unzip ‘/content/drive/MyDrive/CarsData.zip’
Unzipping the folder that contains images and labels in google colab

Here comes the most important step…

You have now performed almost all the steps and need to write one more line of code that will train the model, but, before that, you need to perform a few more steps and change some directories in order to give the path of your custom dataset and train your model on that data.

Here is what you need to do.

After performing the 4 steps above, you will have the yolov5 folder in your google colab. Go to the yolov5 folder, and click on the ‘data’ folder. Now you will see a folder named ‘coco128.yaml’. 

Go ahead and download this folder.

After the folder is downloaded, you need to make a few changes to it and upload it back to the same folder you downloaded it from.

Let’s now look at the content of the file we have downloaded, and it will look something like this.

We are going to customize this file according to our dataset and annotations.

We have already unzipped the dataset on colab, so we are going to copy the path of our train and validation images. After copying the path of the train images, which will be in the dataset folder and looks something like this ‘/content/yolov5/CarsData/images/train’, paste it in the coco128.yaml file, which we just downloaded.

Do the same with the test and validation images.

Now after we are done with this, we will mention the number of classes like ‘nc: 1’. The number of classes, in this case, is only 1. We will then mention the name as shown in the image below. Remove all the other classes and the commented part, which is not needed, after which our file should look something like this.

Save this file with any name you want. I have saved the file with the name customPath.yaml and now upload this file back to the colab at the same place where coco128.yaml was. 

Now we are done with the editing part and ready to train the model.

Run the following command to train your model for a few interactions on your custom dataset.

Do not forget to change the name of the file you have uploaded(‘customPath.yaml). You can also change the number of epochs you want to train the model. In this case, I am going to train the model only for 3 epochs.

5. !python train.py –img 640 –batch 16 –epochs 10 –data /content/yolov5/customPath.yaml –weights yolov5s.pt

Keep in mind the path where you upload the folder. If the path is changed, then the commands will not work at all.

After you run this command, your model should start training and you will see something like this on your screen.

After all the epochs are completed, your model can be tested on any image.

You can do some more customization in the detect.py file on what you want to save and what you don’t like, the detections where the license plates are detected, etc.

6. !python detect.py –weight /content/yolov5/runs/train/exp/weights/best.pt –source path_of_the_image

You can use this command to test the prediction of the model on some of the images.

Some Challenges That You can Face

Although the steps explained above are correct, there are some problems you can face if you don’t follow them exactly. 

1. Wrong path: This can be a headache or a problem. If you have entered the wrong path somewhere in training the image, it can be not easy to identify, and you will not be able to train the model.
2. Wrong format of labels: This is a widespread problem faced by people while training a YOLOv5. The model only accepts a format in which every image has its own text file with the desired format inside. Often, an XLS format file or a single CSV file is fed to the network, resulting in an error. If you are downloading the data from somewhere, instead of annotating each and every image, there can be a different file format in which the labels are saved. Here is an article to convert the XLS format to YOLO format. (link after the completion of the article).
3. Not naming the files correctly: Not naming the file correctly will again lead to an error. Pay attention to the steps while naming the folders and avoid this error.

Conclusion

In this article, we learned what transfer learning is and the pre-trained model. We learned when and why to use the YOLOv5 model and how to train the model on a custom dataset. We went through each and every step, from preparing the dataset to changing the paths and finally feeding them to the network in the implementation of the technique, and thoroughly understood the steps. We also looked at common problems faced while training a YOLOv5 and their solution. I hope this article helped you train your first YOLOv5 on a custom dataset and that you like the article.