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
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:
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.
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:
These things are even harder to identify than our original image! Hopefully, the labelers knew what they were doing.
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.
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.
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.
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!