Train yolov5. A quick guide from a model to the actual use case.

Part 1 of the two-part series shows you how to train yolov5 and later deploy it to production in a simple way.

Last update: February 2022

Introduction

This is the first article from the two-part series. In general, we will go from zero to one in building your AI startup. Well, we will not exactly invent anything new, but we will show you how to train (part 1) and deploy YOLOv5 to production (part 2). By the deployment, we mean that you will be able to connect your model to any service and device as a result. Specifically, the problem we will be solving is called Object Detection.

For example, I’ve found a beautiful image with a godly creature in the middle on one of my favorite websites - Unsplash. I would like to know what it is so I could find more images like this on the internet.

Fortunately, there are already a lot of open-source models that can tackle this problem. Therefore, we will use the infamous YOLO architecture that achieves state-of-the-art results on many datasets.

Guide to train yolov5

We will train yolov5 model which is the newest version of the underlying architecture. Important to realize that we will walk you through the whole process below. For your convenience, we’ve also prepared a Google Colab notebook with all the necessary steps to do it.

Prepare a picture

We will first do some imports:

from IPython.display import Image, clear_output
import requests
import torch

Now, we will download the picture from the Introduction section to the notebook:

img_data = requests.get('https://unsplash.com/photos/w2DsS-ZAP4U/download?force=true&w=640').content

with open('undentified.jpg', 'wb') as handler:
    handler.write(img_data)

You can find the image in the ‘Files’ section (the icon looks like a folder) on the left of the screen.

Clone GitHub repo

We will clone the GitHub repo to the notebook:

!git clone "https://github.com/ultralytics/yolov5.git"

For those that don’t know, the ! sign lets you use bash commands inside of the notebook.

Install other packages

We’ve already imported some of the packages, but these are not enough for the model to work. We will install all the required packages using the requiremnets.txt file, but firstly let’s switch the current path to the folder with the model.

%cd yolov5

Now, we can use pip to install the required packages:

%pip install -qr requirements.txt

Note that % is a symbol that lets you use the special notebook commands.

Prepare a coco128 dataset to train yolov5

Before we can train the model, we will need to get an image dataset. Thus, we will use a lightweight dataset coco128 recommended by the authors of the model.

The following commands will download the dataset from an official site, unzip it and remove the zip file:

torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip', 'tmp.zip')
!unzip -q tmp.zip -d ./ && rm tmp.zip

Later, check some images from the dataset:

Image('coco128/images/train2017/000000000009.jpg', width=600)

These things are even harder to identify than our original image! Hopefully, the labelers knew what they were doing.

Train yolov5

Now, we are ready to train yolov5 model. To do this, we will use the python script train.py provided by the authors. It processes the dataset, trains the model, and logs all of the important information. We will also use some default hyperparameters. If you have time to tune them, you can call the following command to find out more about the settings and hyperparameters.

!python train.py --help

The command for training to start is

!python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5s.pt --nosave --cache

The important thing is that we specified the script to save the weights in the yolov5s.pt file. You will need them later when you will be deploying your model.

Detect the picture

FInally, we managed to train yolov5 model. Normally, we would examine the performance of the model on the test dataset before deploying it on production, but it’s a matter for another tutorial. Here, we will check what the algorithm can tell us about the picture from the Introduction section.

We will use the detect.py script provided by the authors. It will not only detect objects on the image but also add bounding boxes with the names of the object and the probability of them being correct.

!python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source ../undentified.jpg
Image(filename='runs/detect/exp/undentified.jpg', width=600)

It turns out that the godly creature in the middle is a cat!

The model gives it 79% confidence. Given the fact that we trained the model for only 3 epochs, it seems like a good result. The model also identified chairs and bowls on the cupboard. Although, the latter ones are not visible well. It’s probably a matter of tweaking the resolution of the picture.

Conclusion

In this tutorial, we’ve learned how to train the YOLOv5 model to detect a picture from the Internet. We left out the details that would be typically taken care of by the machine learning engineer. Nevertheless, we familiarised ourselves with the process. We’re in the middle of the zero-to-one path.

In part 2 of the tutorial, we will find out how to deploy YOLOv5 we’ve just trained on production. This is usually a job of DevOps and requires a lot of insider knowledge to do this properly. Nevertheless, we will leverage the power of the Syndicai platform that will allow you to deploy your AI model on production in a few simple clicks.