Yolov5: Custom dataset with imbalanced classes.

[weihuangxu]

❔Question

I have a custom dataset that has imbalanced classes: ~10K A, ~10K B, and ~100 C. It is hard to detect C class in the test set. The bounding box works on C object but with wrong lables. Is there anyway to add weight to the C class in the loss function? Or is there any other solution to increase the detection of C class?
Thank you!

Additional context

[github-actions[bot]]

Hello @weihuangxu, thank you for your interest in our work! Please visit our Custom Training Tutorial to get started, and see our Jupyter Notebook , Docker Image, and Google Cloud Quickstart Guide for example environments.

If this is a bug report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom model or data training question, please note that Ultralytics does not provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as:

• Cloud-based AI systems operating on hundreds of HD video streams in realtime.
• Edge AI integrated into custom iOS and Android apps for realtime 30 FPS video inference.
• Custom data training, hyperparameter evolution, and model exportation to any destination.

For more information please visit https://www.ultralytics.com.

[glenn-jocher]

@weihuangxu sure, there's a number of methods that you can research online, but none of them work as well as simply collecting more training data for your underrepresented classes.

[weihuangxu]

@glenn-jocher Thanks for the quick response. Unfortunately, we don't have extra data available for underrepresented classes. I found a solution to partially solve my issue.

[gizemtanriver]

Hi @weihuangxu, I have a similar issue with imbalanced classes in my dataset. How did you solve the issue? Did you change the class weights in the loss fx? Thanks in advance.

[Raian-Rahman]

Hi @gizemtanriver, can you tell me where should i change the class weight?

[gizemtanriver]

Hi @Raian-Rahman, I actually ended up using only one of the classes in my dataset. So I am not sure how to change the class weights. Perhaps have a look at the labels_to_class_weights function in utils-->general.py.