YOLOv8 Weight Penalty

[NerminMostafa]

Search before asking

• I have searched the YOLOv8 issues and discussions and found no similar questions.

Hello! First of all, thank you very much for your team's outstanding contributions!

Question: Is there a way to apply a weight penalty to train specific classes more effectively?

I'm training a YOLOv8n model with a custom dataset (PlantVillage dataset) to detect leaf diseases. The dataset is balanced and split 60:40. I trained the model for 150 epochs with the following settings: imgsz = 256, batch_size = 64, and using a GPU.

Despite trying different versions of the model (YOLOv8s, YOLOv8m) and experimenting with various training parameters (learning rate adjustments, increasing the number of epochs, changing weight decay, modifying patience number, and increasing close_mosaic to 25), the results remain almost the same.

The model performs very well on some classes but struggles with others, specifically classes 4, 5, 7, 8, 9, and 11 (zero-based indexing). I need to improve the model's learning for these specific classes.

More info:

Here are the losses from the last epoch:
Train box_loss = 1.4227
Train cls_loss = 0.94637
Train dfl_loss = 1.2384

Valid box_loss = 1.8663
Valid cls_loss = 1.3516
Valid dfl_loss = 1.5154

Here's the confusion matrix

Please let me know if there is a way to apply a weight penalty to improve my model and how to apply it, Thank you.

Additional

No response

[github-actions]

👋 Hello @NerminMostafa, thank you for your interest in Ultralytics YOLOv8 🚀! We recommend a visit to the Docs for new users where you can find many Python and CLI usage examples and where many of the most common questions may already be answered.

If this is a 🐛 Bug Report, please provide a minimum reproducible example to help us debug it.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset image examples and training logs, and verify you are following our Tips for Best Training Results.

Join the vibrant Ultralytics Discord 🎧 community for real-time conversations and collaborations. This platform offers a perfect space to inquire, showcase your work, and connect with fellow Ultralytics users.

Install

Pip install the ultralytics package including all requirements in a Python>=3.8 environment with PyTorch>=1.8.

pip install ultralytics

Environments

YOLOv8 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all Ultralytics CI tests are currently passing. CI tests verify correct operation of all YOLOv8 Modes and Tasks on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[glenn-jocher]

Hello! Thanks for reaching out and for your detailed explanation of the issue.

To apply a weight penalty specifically to certain classes in YOLOv8, you can adjust the cls weights in the loss function for those classes. This can be done by modifying the cls parameter in your training configuration file. You can increase the weight for classes 4, 5, 7, 8, 9, and 11 to give more importance to these classes during the loss calculation.

Here's a quick example of how you might adjust these settings in your YAML configuration file:

# Assuming 'cls' weights are not initially set, you can define them like this:
cls_weights: [1, 1, 1, 1, 2, 2, 1, 2, 2, 2, 1, 2, 1]  # Increasing weights for classes 4, 5, 7, 8, 9, 11

Make sure the length of cls_weights matches the number of classes, and increase the values for the classes you want to penalize less (or focus more on correcting). This should help the model pay more attention to these specific classes during training.

Let me know if this helps or if you need further assistance! 🚀

[NerminMostafa]

@glenn-jocher Thank you so much for your quick response and helpful instructions! I implemented the class weights as you described, but unfortunately, the improvement was not as significant as I expected. Could you please provide further assistance or suggest another solution to enhance the model's performance for these specific classes? Your support is greatly appreciated!

[glenn-jocher]

@NerminMostafa hello! I'm glad you tried the class weights adjustment. If the improvement isn't as expected, you might consider augmenting your dataset specifically for the underperforming classes. This can involve adding more varied images of these classes or applying more aggressive data augmentation techniques (like rotations, scaling, color variations) to these specific classes to help the model generalize better.

Here's a quick example of how you might adjust your data augmentation settings in your YAML configuration file:

# Custom augmentations for specific classes
augmentations:
  rotation: 20
  scale: 1.2
  hue: 0.05
  saturation: 1.5
  brightness: 0.2

Adjusting these parameters can sometimes lead to better feature learning for challenging classes. Let's give it a try and see how it goes! 🌟

[NerminMostafa]

@glenn-jocher Thank you so much for the suggestion! I will definitely give it a try and adjust the data augmentation settings specifically for the underperforming classes. Your example for the YAML configuration file is very helpful. I appreciate your guidance and will implement these changes to see if it leads to better results. Thanks a lot!

[glenn-jocher]

@NerminMostafa You're very welcome! I'm glad you found the suggestions helpful. 🚀 Adjusting those data augmentation settings can indeed make a significant difference. Keep us posted on your progress, and don't hesitate to reach out if you have more questions or need further assistance. Best of luck with your training! 🌿