Train YOLO V9 on local GPU

Photo by Steve Johnson on Unsplash

The YOLOv9, the latest iteration in the YOLO (You Only Look Once) series, was unveiled on February 21, 2024. Since its introduction in 2015, the YOLO object-detection algorithm has captivated the attention of tech enthusiasts, data scientists, and ML engineers worldwide. Its open-source nature and active community contributions have led to a substantial following. With each new release, YOLO’s architecture has evolved to be more user-friendly and faster, making it accessible to a broader audience.

YOLO, presented in a research paper by J. Redmon, S. Divvala, R. Girshick, and A. Farhadi, marked a significant advancement in real-time object detection, surpassing its predecessor, the Region-based Convolutional Neural Network (R-CNN). Notably, YOLO operates as a single-pass algorithm, employing a lone neural network to predict bounding boxes and class probabilities by analyzing the entire image as input.

Now let's take a look at how to train a YOLO v9 model on your local GPU.

In my case, I am using a NVIDIA Geforce 3060 with a 12GB VRAM and my computer has 32GB RAM.

First things first, let’s clone the Yolo V9 github repository. Before that, cd to the desired the directory you want this to get cloned.

cd <PATH WHERE YOU WANT THIS DIRECTORY TO BE CLONED>

git clone https://github.com/hsivasub/yolov9.git

Note the original repository of YOLO v9 is https://github.com/WongKinYiu/yolov9.git. I was facing some issues with the detection after training so I modified few lines and I have added it to my repository.

After cloning, this will make a local folder.

Then let’s collect the dataset and place it inside the folder. For this, I am using a dataset downloaded from https://universe.roboflow.com/roboflow-100/weed-crop-aerial/dataset/2.

Unzip the dataset and then create a new folder inside the cloned YOLOV9 directory and name it as ‘dataset’

Inside them copy the dataset that you just extracted. The dataset would have the train, test, and valid folder.

Each of the train, test, and valid has the images and the labels arranged as follows,

The images would contain the actual images and the labels would contain the labels for the same images in the YOLO format.

From this dataset, copy the data.yaml file to the cloned Yolo directory.

And inside that change the values to these,

Here we specify the train, test, and val path.

Now let’s download the pre-trained model that would be used for training. Using a pre-trained model would save a lot of time in training and the model would be much more accurate. This is called transfer learning.

These are the pretrained model that can be found on the Github repository https://github.com/WongKinYiu/yolov9

In my case, I used the YOLOv9-E. Download the model and place it inside the cloned repository.

Now go to the cloned Yolo directory. Inside that navigate to models/detect/ and then copy and paste the yolov9.yaml in the same directory and rename them as yolov9_custom.yaml.

Inside that file, change the nc=2 because here the number of classes is basically 2 in our dataset as shown below,

By default, it would be 80 as they were trained on the COCO dataset which has 80 classes.

Now let’s start the training by running this line in the anaconda prompt inside the folder.

python train_dual.py --workers 8 --batch 4  --img 640 --epochs 50 --data data.yaml --weights yolov9-e-converted.pt --device 0 --cfg models/detect/yolov9_custom.yaml --hyp data/hyps/hyp.scratch-h
igh.yaml

Here device 0 indicates it would use the GPU to train.

Once the training is completed, the model artifacts would be stored in this folder automatically,

The following files would be present in the location runs/train/exp2. The exp2 indicates my 2nd experiment. It would automatically create folders like exp3, exp4, and so on as you train multiple times.

Here the weights would contain the best and the last epoch weights stored in the .pt format.

Prediction

Now in the Anaconda Prompt, type in the following for detection,

python detect.py - img 1280 - conf 0.1 - device 0 - weights runs/train/exp2/weights/best.pt - source dataset/test/images/12122_jpg.rf.6cb950f8064f7d2101f5210f77760ca4.jpg

This will do the detection for a test image. The results would be stored in the runs/detect/exp folder.

The above image shows the successful detection using YOLO v9.

That is it!!!

If you like this article, please upvote my article and also follow me in Medium for similar content. Your support is what keeps budding writers like me encouraged!!

Follow me https://medium.com/@sharish2395

If you have any questions, please reach out to me on LinkedIn at https://www.linkedin.com/in/harish-siva-subramanian-320b78b7/

Thank you for reading!!