Train YOLOv8 on Custom dataset in Windows GPU

In my last post, I guided you how to train YOLOv3. In this post, we will explore how to train YOLOv8 on custom dataset in Windows GPU. We will cover all required steps like preparing the dataset, configuring and training the model, evaluating performance of the model. Finally use custom trained model to inference on new images.

Must Read: Recommend you to read below tutorials before this.

• YOLO object detection using deep learning OpenCV | Real-time
• Train YOLOv3 Custom object detection model in Windows GPU
• Object Detection and More using YOLOv8

Configure YOLOv8 for Windows GPU

One thing I also liked about this YOLO v8 is that you don’t have to clone the repository or do any manual things. In my YOLO v3 setup post you can see that we used to clone the repository, set up requirements, and configure the model file manually. But with YOLO v8, you don’t have to do any of these hard manual things. You only need to install Ultra Analytics, That’s it.

There are two ways to use UltraAnalytics YOLO v8:

1. Using CLI or command line
2. Python script

So let’s set up UltraAnalytics YOLO v8 to train a custom model for our own dataset in Windows 10 (You can follow same steps for Google collab also). I will break entire configuration into some steps:

Step1: Create Virtual Environment

You must have Python 3.7 or above to use UltraAnalytics YOLO v8. It will be a good idea to create a fresh virtual environment with Python 3.7. If you are using Anaconda, you can create an isolated virtual environment using below command.

conda create -n ENV_NAME python=3.7
activate ENV_NAME

Step2: Install Ultalytics

This is the fun part. You no need to do anything manually. You just need to run the below command and YOLO v8 will be installed in your windows system.

pip install ultralytics

Step3: Setup Jupyter notebook

I think we all love Jupyter Notebook to write Python scripts. To use Jupyter Notebook inside our virtual environment, run below three commands inside your Virtual env.

conda install jupyter
conda install nb_conda
conda install ipykernel

Step4: Validate Installation

Now we need to check whether we installed YOLO v8 correctly or not. To do that open Jupyter Notebook (inside virtual environment) and run below Python script.

from IPython import display
display.clear_output()

import ultralytics
ultralytics.checks()

Output

As you can see we have successfully installed Ultralytics YOLOv8, but with CPU. Let’s now try to install YOLOv8 with GPU.

Step5: YOLOv8 GPU Configuration

In the above output, you can see that YOLOv8 is configured with CPU. This is because your torch is CPU version. You can check your torch backend by running below Python code.

import torch
torch.cuda.is_available()

The output is False because torch is not using CUDA or GPU as its backend. Its using CPU.

To use your GPU for training custom YOLOv8 for your own data, you need to install torch with GPU corresponding to your CUDA version.

So let’s check our CUDA version by running below Python script in Jupyter Notebook.

!nvidia-smi

As you can see I am using CUDA version: 11.7. Now go to this link and select your system configuration to get the command to install torch with GPU.

Since I am using CUDA 11.7 and Windows system, I have selected them along with the Stable version. I find difficulty using pip command. So I will recommend you to use conda command.

conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch -c nvidia

Along with that we also need to install conda version of cuda toolkit inside our virtual environment. Below command is to install that.

conda install cudatoolkit

Now let’s check our torch backend again.

import torch
torch.cuda.is_available()

Okay, now we are getting output as True. That means torch is now using GPU as its backend. Let’s now check YOLOv8 is also configured with GPU or not again.

from IPython import display
display.clear_output()

import ultralytics
ultralytics.checks()

Yes! YOLOv8 is now configured with GPU as its backend. You can see it is also showing my NVIDIA graphics card model number which is GeForce GTX 1050 Ti (4 GB).

Train YOLOv8 on Custom dataset

So we successfully installed YOLOv8 from Ultralytics library with GPU on our Windows computer. Now let’s try to train yolov8 on our own custom dataset.

I will break the entire process of training a custom YOLOv8 model into some steps. Note that I am using Windows 10 with GPU but the steps will be same for Linux or Google Colab.

Okay now let’s go into each step to train our custom yolov8 model for our own dataset.

Step1: Data Preparation

For data preparation, you need to follow below steps:

• Collect images (either download or capture)
• Select proper set of images
• Annotate images

In this post, I will not guide you to prepare data to train a YOLO model. The training data format is same as YOLOv3. You can read this detailed post to learn more about data preparation to train custom YOLOv8 or v3 models.

I am going to use Goats detection dataset to train our custom YOLOv8 model. This dataset contains 1.5k training images (with annotation) of goats. You can find different formats of that data set. Since we need to train a YOLOv8 model, I am going to download that version.

Click on YOLOv8 button > select download zip to computer option in the pop-up window > Continue. It will download the dataset as zip file. Unzip that file.

This dataset has three folders: train, test, and valid. Each folder contains images (contains original images) and labels (contains coordinate information for the corresponding image) folder.

Along with that, you can find another important file named data.yaml, which contains class names and each folder location. So our working directory should look like below:

Woking directory
    ├── Downloaded dataset folder
    	   ├── train
    	   |     ├── images
    	   |     └── labels
    	   ├── test
    	   |     ├── images
    	   |     └── labels
    	   ├── valid
    	   |     ├── images
    	   |     └── labels
    	   └── data.yaml

If you find difficulty understanding this dataset, please read this tutorial where I explain data format to train YOLO model.

Step2: Rename Folder

So you unzipped the folder and placed it into your working directory. Now you need to rename the folder to “datasets“. With a little bit of research, I did not find the reason behind it, but we need to name the dataset folder like that. If you find any solution, please let me know in the comment section.

Now copy the data.yaml file from datasets (newly renamed) folder to your root directory.

So now our folder structure should look like below:

Woking directory
    ├── datasets
    |	   ├── train
    |	   |     ├── images
    |	   |     └── labels
    |	   ├── test
    |	   |     ├── images
    |	   |     └── labels
    |	   ├── valid
    |	   |     ├── images
    |	   |     └── labels
    |	   └── data.yaml
    └── data.yaml

That’s it we are all set to train the YOLOv8 model on our custom dataset.

Step3: Train YOLOv8 model for custom data

There are mainly two ways to train custom YOLOv8 model:

1. Using CLI (or command line)
2. Using Python Script

1. CLI Method

Using CLI method, you just need to run below command in the command line.

yolo task=detect mode=train model=yolov8n.pt data= data.yaml epochs=30 imgsz=832 plots=True device=0

In the above command:

• task: Here you need to mention the task.
  • detect – for object detection
  • segment – for image segmentation
  • classify – for image classification
• model: Here you can select any one of the 5 available pre-trained YOLOv8 models for custom object detection (you can also train from scratch with those models). List of models can be found here.
• imgsz: This is the image dimension of your training images.
• device = 0 is to utilize your GPU while training the model. You can skip this if you want to train using CPU
• plot Ture is to save model statistics in your working directiory

In my system (4 GB NVIDIA 1050 Ti graphics) it took around 1 hour to complete training for 30 epochs.

After the successful completion of training yolov8 on your custom dataset, a folder called “runs” will be created inside your working directory.

You can find your final custom YOLOv8 model inside runs > detect > train> weights. So now our folder structure should looks like bellow:

Woking directory
	├── runs
	|    └── detect
	|         └── train
	|              ├── Some statistical images
	|              └── weights
	|                   ├── best.pt
        |	            └── last.pt
	|
        ├── datasets
        |	 ├── train
        |	 |     ├── images
        |	 |     └── labels
        |	 ├── test
        |	 |     ├── images
        |	 |     └── labels
        |	 ├── valid
        |	 |     ├── images
        |	 |     └── labels
        |	 └── data.yaml
	└── data.yaml

2. Python Method

Along with Command ine, you can train custom YOLO v8 model through Python. Below Python code is to train yolov8 on custom dataset:

from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)

# Train the model
model.train(data='data.yaml', epochs=2, imgsz=832, device=0)

Note: Theoretically you can train using above Python script but when I was trying to do this using Jupyter. The Notebook was getting died after some time. But if you do it using CLI (command line) there is no issue.

Step4: Inference Trained YOLOv8 model

So now we successfully trained YOLOv8 model for our custom dataset. Now let’s see how our model is performing for images. Below Python code is to detect Goats from an input image.

from ultralytics import YOLO

# Load our custom goat model
model = YOLO("runs/detect/train/weights/best.pt")

# Use the model to detect object - goat
model.predict(source="goat_image.jpeg", save=True, show=True)

That looks good from a model trained on only 30 epochs. It is successfully detecting goat from the input image.

Step5: Real-time object detection using YOLOv8

Now let’s see how we can detect object real-time using our custom-trained YOLOv8 model. There are two ways: 1. Using Webcam, 2. Using external video. Let’s see both of them.

Webcam Code

Below Python code is for real time yolov8 object detection using Webcam.

# For webcam
from ultralytics import YOLO

# Load custom trained YOLOv8 model
model = YOLO("runs/detect/train/weights/best.pt")

# Use the model to detect object
model.predict(source="0", show=True)

In this code source, 0 means webcam.

Video Code

To do the same thing for external video, you just need to mention the video path instead of 0 in the source.

# For video
from ultralytics import YOLO

# Load custom trained YOLOv8 model
model = YOLO("runs/detect/train/weights/best.pt")

# Use the model to detect object
model.predict(source="goat_video.mp4", show=True)

End Note

In this tutorial, I showed how you can train your own YOLOv8 model on custom dataset. For this tutorial I used windows 10 system with GPU. But you can follow same steps for other systems like Linux or google collab etc using CPU also.

I have not discussed much about data preparation for YOLO model in this post. If you want to understand more about data preparation, read this detailed tutorial.

That’s it for this tutorial. If you have any questions or suggestions regarding this post, please let me know in the comment section below.

Anindya

Hi there, I’m Anindya Naskar, Data Science Engineer. I created this website to show you what I believe is the best possible way to get your start in the field of Data Science.