Train yolov5. A quick guide from a model to the actual use case.⋅ Michał Zmysłowski
Part 1 of the two-part series that show you how to train yolov5 and later deploy it to production in a simple way.
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.
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.
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.
Now, we can use
pip to install the required packages:
%pip install -qr requirements.txt
% is a symbol that lets you use the special notebook commands.
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:
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 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)
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.