2

I'm training YOLO on custom dataset(using Alexey AB implementation of Darknet). It has 3 classes of images where class 1 has 45k images, and remaining two have around 1k images.

After training it for 6k iterations, the loss is in between 1.5 and 2. However, when i tried running it on a video, it is only detecting class 1.

I would like to know what is the reason for this, is it because of the imbalance in the no of images in the dataset? Is there a way to solve this problem?

user8128965
  • 41
  • 1
  • 7

1 Answers1

3

Yes, to begin with you've an unbalanced dataset. The recommended number of images per class is >2000 (according to the directions on the repository).

I'd recommend you to have a test set and track your mean average precision along with loss while you're training. If the average precision is low, you can assume that not all classes are being detected equally well.

There are a few ways to solve this but I'm assuming you already might have figured out.

  • Collect more data for other classes.
  • Try data augmentation for other classes with rotation, noise, etc.
  • You could try adding synthetic images for the other classes.
  • Even though it's a bad idea to leave out data for the 1st class, at least try reducing it so there's a balanced dataset.

I'm calling other classes for the classes with less data.

rajput
  • 166
  • 7
  • To add to this, you could also define a custom loss function that weighs an error for a class it has less examples more than one for a class with more examples. You can read more about that [here](https://www.tensorflow.org/tutorials/structured_data/imbalanced_data#class_weights) – Victor Sonck Mar 10 '20 at 08:39