Detecting Chess Pieces / Pipeline For Training Custom YoloV8 Models
Training custom YOLO models can be quite complex, especially for the first time. There are common steps for training and testing a model, and for every task, these steps are nearly the same. So, creating a pipeline and following it makes your task easier.
- In this article, I am going to explain how to systematically train a custom YOLO model. By following this guide, you can train models for any object detection task.
- Before following these pipeline, you need to decide on a training method: Using GPU, CPU, or servers like Google Colab.
Training Model with GPU
Training a model can be time-consuming, it can even take more than a week. There are 3 main factors that affect training time:
- Size of the Data, GPU power, and Training parameters.
More data requires more time to train, it is as simple as that. For GPU, if you have an Nvidia GPU you are lucky. Simply install PyTorch with CUDA support. You can follow this tutorial for installing PyTorch with CUDA support(link).
f you don’t have a GPU, you can still train models with your CPU but it will be much slower. If you don’t have a GPU or if you prefer not to install PyTorch on your computer, I strongly recommend you to use Google Colab or Kaggle, they provide free GPU for users.
There are seven main steps in my pipeline, and now it is time to discuss them.
- Creating Data
- Create a Folder for Project
- Create a YAML file
- Choose a pretrained YOLOv8 model
- Create a Python file for training and train the model
- Observing Model Metrics
- Testing Model
1. Creating Data
Data is one of the most important things in Deep Learning models. Without proper data, it is impossible to obtain a good model. There are various ways to acquire data, you can either create your own or find it on the internet.
- If you decide to create your own data, it can be a painful and time consuming process. First, you need to gather images( take screenshots , take photos with your camera, etc.), and then annotate images by drawing bounding boxes and labeling objects using a tool like labelImg.
- There are millions of data available on the internet, and a significant portion of it is completely free. Here are some websites that you can find suitable data for your purpose (about 95% of the time ): Kaggle, Roboflow
I always first check the internet, and if I can’t find data that fits my purpose, then I create my data by myself.
→ I am going to use a dataset from Roboflow for chess piece detection (link)
2. Create a Folder for the Project
Create one main folder that contains all the necessary information:
- dataset folder : images and labels
- YOLO model : pretrained YOLO model
- YAML file : information about data
- Python file : for training
You can name the main project folder whatever you want, it doesn’t matter. I named it “chess_model_big_data”. After creating the main project folder, paste your data into that folder (one folder for train and valid folders, I named it chess_big_data), the format should be as following diagram :
I am going to explain all the folders/files in the following steps, I just wanted to show you the main hierarchy of folders/files.
3. Create a YAML File
YAML file contains information about the data path, the number of different classes, and the names of classes. I named the YAML file “chess_big_data.yaml “
It contains the paths to the train and val folders , the number of classes, and the names of these classes.
4. Choose a pretrained YOLOv8 model
There are bunch of pretrained YOLOv8 models and choosing the pretrained model is completely depends on your purpose. If you want higher accuracy you need to choose the more complex model, or if you need higher FPS you need to choose more basic models. There is a trade-off between speed and accuracy.
Choose a pretrained model and download it, then paste it into the main project folder. I downloaded YOLOv8m, you can choose whatever model you want.
5. Create a Python file for training and train the model
Simply a file for training a model, I created a .ipynb file because I am using jupyter notebook, you can create a .py file as well, that is okay. I named it as a Train_chess_model.ipynb
There are numerous parameters for training a custom YOLO model. Some of them are easily understandable from their names, but some need explanation :
- patience : Number of epochs to wait without improvement in validation metrics before early stopping the training. (it helps to prevent overfitting)
- save : Enables saving of training checkpoints and final model weights.
- save_period : Frequency of saving model checkpoints, specified in epochs.
- project : Name of the project directory where training outputs are saved.
- name : Name of the training run.
- verbose : providing detailed logs and progress updates.
- resume : Resumes training from the last saved checkpoint.
- amp : Enables Automatic Mixed Precision (AMP) training, reducing memory usage and possibly speeding up training with minimal impact on accuracy.
- lr0 : Initial learning rate (i.e. SGD=1E-2, Adam=1E-3)
- dropout : Dropout rate for regularization in classification tasks.
- plots : Generates and saves plots of training and validation metrics, as well as prediction examples, providing visual insights into model performance and learning progression.
from ultralytics import YOLO
import torch
# check GPU
print(torch.cuda.get_device_name(0))
# Load the model.
model = YOLO('yolov8m.pt') # pretrained model
# Training.
results = model.train(
data='chess_big_data.yaml', # .yaml file
imgsz=416, # image size
epochs=50, # epoch number
batch=4, # batch size , I normally use 8 or 16 but my GPU gave memory errors, therefore I reduced it to 4 for this time.
name='chess_big_data_model4', # output folder name, it contains model weights and all of the other things.
plots=True, # Plots about metrics (precision, recall,f1 score)
amp=False, # amp=True gives me an error, I don't know why , If it doesn't give you an error, set it to True
- For training, simply run python file. At first glance, you may be surprised because there are going to be many lines about the model, labels, and folders. If everything works well, the training will start.
- It may take a long time depending to your GPU.
- After training is finished, there will be a summary of precision, recall, and mAP values.
Results saved to C:\Users\sirom\runs\detect\chess_big_data_model4
6. Observing Model Metrics
After training is finished, you will see a line indicating: Results saved to “some_folder”. Inside this folder, there are a lot of important informations about the model. By looking this information, you can examine your model’s performance.
And most importantly, this folder contains your trained model(best.pt)
7. Testing Model
Testing a model on images and videos is a straightforward process.
from ultralytics import YOLO
from PIL import Image
model = YOLO("best.pt") # path to .pt model file
results = model("test_image.jpg") # path to test image
# Show the results
for r in results:
im_array = r.plot() # plot a BGR numpy array of predictions
im = Image.fromarray(im_array[..., ::-1]) # RGB PIL image
im.show() # show image
im.save('results.jpg') # save image