Object Detection

Train Gold-YOLO object detection on a custom dataset with trainYOLO.

September 27, 2023
5 min

Gold-YOLO is the newest addition to the YOLO family and was developed by researchers from Noah’s Ark Lab (Huawei). They scored new highs on the COCO benchmark and surpassed the popular YOLOv8 algorithm, both in accuracy and speed.

Comparison of state-of-the-art efficient object detectors on Tesla T4 GPU. Both latency and throughput (batch size of 32) are given, using TensorRT8. Picture taken from original paper.

How did they do that? Well, they achieved these results by mainly redesigning the mechanism that fuses the features of multiple resolutions in the decoder, which up to now was mostly done with so-called Feature Pyramid Networks. They developed a new mechanism, called Gather-and-Distribute, which boosts the multi-scale feature fusion capabilities and achieves an ideal balance between latency and accuracy across all model scales. Like most YOLO versions, they also release multiple versions, ranging from gold_yolo-n to gold_yolo-l, trading speed for accuracy.

The Gold-YOLO architecture. Picture taken from original paper.

In this tutorial, we will first take a look at their COCO-pre-trained model and next, we will guide you through the steps on how to train your own Gold-YOLO detector on your own data. For this, we will use the trainYOLO platform. So if you don’t have an account yet, make sure you create a free account here.

Inference with a pre-trained COCO model

Let’s take a look at running Gold-YOLO using a pre-trained COCO model. We start by installing the library, downloading the gold-yolo-n model (the smallest one available), and running the model on some demo images. You can find the Colab Notebook here.

The results are displayed below. We notice that with the default confidence value of 0.4, the horse and football are not detected. When lowering the threshold to 0.2, they are detected, although there also appear some false positives.

Results for gold-yolo-n.pt with a confidence of 0.4
Results for gold-yolo-n.pt with a confidence of 0.2

Training a custom Gold-YOLO model on your own dataset

Next, we will go over the steps on how to train your own Gold-YOLO object detector on your own dataset. We will guide you through the steps of labeling your data and training the model using the trainYOLO platform. As an example, we will train a Padel player detector for game analysis.

1. Create a new project

We start by creating a new project, which will hold our images, clicking on the green ‘plus’ icon. We set the name to padel, keep the annotation type at Bounding Boxes, and set a single category player. If you have multiple categories, you can add more.

2. Upload your images

To upload images to our project, go to the data tab and click on the green ‘plus’ icon. Select the images you would like to upload and click upload. In this example, we will upload some images of a Padel game.

3. Label your images

Once we have uploaded the images, we start labeling. By clicking on an image, you enter the labeling editor. By selecting the bounding box tool in the left menu, we can draw a box around each object, in our case each player. Note that if you have multiple object categories, you need to set the right category using the drop-down list on top. Once finished, you can use the arrow keys on top to go to the next image or select a specific one in the carousel at the bottom.

4. Train Gold-YOLO object detector

Next, we will be going over all the steps on how to train the Gold-YOLO object detector. You can find all the steps in this colab notebook. You only need to fill in your trainYOLO API key (which you find under your avatar) and Project Name and run Runtime -> Run all.

Let’s start by downloading our labeled images. To do so, we need to install the trainYOLO python library and enter our API key (which you can find under your avatar) and Project name. As Gold-YOLO is based on the YOLOv6 codebase, we use the YOLOv6 export format.

Once our dataset is downloaded, we clone the Gold-YOLO repository and download the COCO-pre-trained model (version n)

Now we have everything ready to start training our model. We can set the image size, number of epochs, and batch_size:

Once training is finished, we can infer the model on our validation set and visualize the results. Note that there is currently still a small bug in the codebase, which we fix with the sed command:

You can find the results under ‘runs/inference/’. For example:

Conclusion

In this tutorial, we went through the steps on how to train your own Gold-YOLO object detector. For more details, you can take a look at their official GitHub repository. If you want to compare with YOLOv8, you can check out our tutorial on how to train your own YOLOv8 object detector.

Latest posts

Get in touch.

Thank you! Your submission has been received!
Oops! Something went wrong while submitting the form.
By clicking “Accept All Cookies”, you agree to the storing of cookies on your device to enhance site navigation, analyze site usage, and assist in our marketing efforts. View our Privacy Policy for more information.