Training a Custom Object Detection Model With Yolo-V5
By: Prabhat Kumar Sahu
Tags: Yolo, Computer Vision, Yolov5, Object Detection

Object Detection is a task in computer vision that focuses on detecting objects in images/videos.

There are various object detection algorithms out there like YOLO (You Only Look Once,) Single Shot Detector (SSD), Faster R-CNN, Histogram of Oriented Gradients (HOG), etc.

In this article, we are going to use Yolo-V5 to train our custom object detection model. YOLO is one of the most famous object detection models.

Version 1.0 of YOLO-V5 got released on 27 May 2020 by Glenn Jocher who is the founder & CEO of Utralytics. It’s written in PyTorch and it’s available in Github.

Prerequisites

It’s good to have a basic knowledge of deep learning computer vision. And how to work in a Google Colab environment.

Steps Covered in this Tutorial

To train our own custom object detector these are the steps to follow

Preparing the dataset

Environment Setup: Install YOLOv5 dependencies

Setup the data and the directories

Setup the YAML files for training

Training the model

Evaluate the model

Visualize the training data

Running inference on test images

Export the weight files for later use

Preparing the dataset

We are going to use the Aquarium Dataset which is available in Roboflow Public Object Detection datasets. You can check out more datasets that are there.

The aquarium dataset consists of 638 images. The images were already labeled by the Roboflow team. It has 7 classes such as fish, jellyfish, penguins, sharks, puffins, stingrays, and starfish, and most images contain multiple bounding boxes.

To download the dataset you need to create a roboflow account first. It’s very simple and easy.

If you want to use your own images you can use annotation tools such as LabelImg, CVAT or you can try any large-scale solutions like Scale or AWS Ground Truth.

When you are annotating yourself make sure to follow the best practices. Check this link out for more details.

After the dataset is prepared then we are all set to set up the environment and train the dataset.

Environment Setup

Here’s the link to my Notebook: Google Colab

You need a google account to use Google Colab. You can either use my notebook to train or you can create your own notebook and follow along.

In Google Colab, you will receive a free GPU for 12 hours. If you use a new notebook in Colab change the runtime session to GPU.

If you are planning to use my notebook then make sure to File → save a copy in your drive. Then you will be able to edit the code.

Installation of the dependencies

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

Somehow the PyTorch version didn’t get compatible with the GPU so I installed another version of PyTorch by

#installing for google colab GPU use

!pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html

.

%cd /content/yolov5

!ls

We can import and take a look at our GPU Specification provided by Google Colab.

import torch from IPython.display import Image # for displaying images from utils.google_utils import gdrive_download # for downloading models/datasets print('Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))



Here’s what I got

Using torch 1.6.0+cu101 _CudaDeviceProperties(name='Tesla T4', major=7, minor=5, total_memory=15079MB, multi_processor_count=40)

Google Colab comes preinstalled with Cuda and Torch and some other dependencies. If you are planning to train locally then you will have to setup Cuda and the dependencies on your own. I will surely make a tutorial about it later on.

Setup the data and the directories

After the environment set up is done. We can import the dataset into colab. As I am using the Roboflow dataset I will be downloading, if you plan to use your own you can import it using Google Drive.

# You need to sign up in roboflow to get the key and then you can use the dataset !curl -L “https://public.roboflow.com/ds/PUT YOUR OWN KEY HERE” > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

This will download the data, unzip and save it inside the yolov5 directory.

Project Folder structure

Setup the YAML files for training

To train a YOLO-V5 model, we need to have two YAML files.

The first YAML to specify:

where our training and validation data is

the number of classes that we want to detect

and the names corresponding to those classes

This YAML of ours looks like this:

train: ./train/images

val: ./valid/images nc: 7

names: ['fish', 'jellyfish', 'penguin', 'puffin', 'shark', 'starfish', 'stingray']

The second YAML is to specify the whole model configuration. You can change the network architecture in this step if you want but we will go with the default one.

The YAML which we term custom_yolov5s.yaml :

We can put the YAML file anywhere we want because we can reference the file path later on. But it’s a good idea to put it inside the YoloV5 directory.

Training the model

After the configuration is done we can begin our training.

There are multiple hyperparameters that we can specify which are:

img: define input image size

define input image size batch: determine batch size

determine batch size epochs: define the number of training epochs.

define the number of training epochs. data: set the path to our YAML file

set the path to our YAML file cfg: specify our model configuration

specify our model configuration weights: specify a custom path to weights

specify a custom path to weights name: result names

result names nosave: only save the final checkpoint

only save the final checkpoint cache: cache images for faster training

We need to specify the path of both YAML files which we created above.