Advanced loss manipulation in YOLO

[ShaharJamshy]

Hello,

I’m a computer vision researcher at Scopio Labs (https://scopiolabs.com/), we are looking into YoloV8 to detect blood cells in images.

We encountered some challenges that we would like to solve using more advanced logic on the loss function, and were wondering what would be the best way to implement these solutions within the Ultralytics framework.

Examples for scenarios we want to address:

1. We want efficiently train a detector when the classes for the detected object are highly skewed, and we want to include enough examples of the low prevalence class without investing a lot effort labeling the high prevalence class

   • For example, red blood cells are much 40x more abundant in the blood than white blood cells, see in the example below image red blood cell detection boxes in red, white blood cell in purple, platelets in yellow
   

   • If we want a good coverage of the different white blood cells types in the training data we would need to sample a large number of images, labeling all the red blood cells for all these images would be impractical.

   • This means we have a lot of images where only white blood cells (the purple box) are labeled, and none of the red blood cells (red boxes) are.

   • One idea we had was, that for the images where we know only white blood cells were labeled we would ignore the loss term for every false detection that isn’t classified as a white blood cell.

2. Another slightly different challenge is how can we train a detector where not all the class labels are well defined

   • There are many cases where we have cells which are hard to label, our labelers might have difficulty deciding between a subset of the classes or even be very uncertain about the class at all.

   • For example, in the images below the cell in the left and right images are clear examples of two stages in the cell development process (Band and Segmented Eosinophils) while the cell in middle image is somewhat in between and hard to label.

   

   • Nevertheless, we do want instances similar to the cell in the middle image to be detected, but we want to be able to avoid providing any class information which might be inaccurate.

   • We thought one way to implement it is to introduce a “don’t care” class, and ignore the detection class loss for all the examples which are labeled with this class.

3. Given our labeled data does not accurately reflect the prevalence of the different classes in the images, we were wondering if using label smoothing to add a prior with the correct distribution would help improve our performance.

We'll greatly appreciate any insight you might have on how to implement these in Yolo, or any other ideas or feedback you might have to address these challenges.

[glenn-jocher]

@ShaharJamshy thank you for reaching out. For advanced loss manipulation in YOLOv8, you can customize the loss function by modifying the training script. To handle class imbalance, consider using focal loss or class weighting. For uncertain labels, you can implement a "don't care" class by adjusting the loss computation to ignore certain classes. Label smoothing can indeed help with class distribution issues. For detailed guidance, please refer to the Ultralytics documentation. If the issue persists, ensure you are using the latest version of the package.

[ShaharJamshy]

@glenn-jocher thank you so much for the prompt answer, looking at the codebase and trying to understand what steps I need to take to customize the loss function, I'm thinking of:

• Copy the v8DetectionLoss code and create a customized loss class in my code, say CustomDetectionLoss
• Amend the __call__ function of CustomDetectionLoss to include the logic we want to introduce
• Subclass DetectionModel
• Override the init_criterion function to return CustomDetectionLoss

Would this be the correct approach to writing our own custom loss logic?

Thanks again
Shahar

[glenn-jocher]

@ShaharJamshy yes, your approach to customizing the loss function by creating a CustomDetectionLoss class and overriding the init_criterion function in a subclass of DetectionModel is correct. This will allow you to implement your custom logic effectively.