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

Introduction

In this two-part series, 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.

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 favorites websites - Unsplash. I would like to know what it is so I could find more images like this on the internet.

train yolov5 model
Cat picture by Paul Hanaoka (https://unsplash.com/@plhnk)

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

Development

We will train the YOLOv5 model which is the newest version of the underlying architecture. 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.

Preparing 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.

Cloning 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.

Installing 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.

Image dataset

Before we can train the model, we will need to get an image dataset. 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

Let’s check some image from the dataset:

Image('coco128/images/train2017/000000000009.jpg', width=600)
coco dataset for model training
Sample image from the Coco dataset

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

Training YOLOv5

Now, we are ready to train the YOLOv5 model. To do this, we will use the python script 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.

Detecting the picture

Okay, so now the model is trained. 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)
train yolov5 model
YoloV5 prediction after training a model with Coco dataset

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.

* * *

If you found that material helpful, have some comments, or want to share some ideas for the next one - don't hesitate to drop us a line via slack or mail. We would love to hear your feedback!

You might like these

How to deploy a PyTorch model in minutes

Tutorial
June 20, 2021
by
Marcin Laskowski

Deploy ML with FastAPI, Docker, and Cloud Run

Tutorial
June 20, 2021
by
Marcin Laskowski