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.

[Mym1121]

Hello, I have recently encountered a similar issue. For some images, the number of certain samples is too large to be fully annotated, which causes the model to hesitate significantly on such "incompletely annotated" targets during training. As a result, the detection confidence for these targets is low, and the miss rate is high. I came across your post and was wondering if you have managed to resolve this problem.

[glenn-jocher]

@Mym1121 no confirmed resolution was posted here. Standard YOLO training treats unlabeled objects as background, so incomplete annotations can reduce confidence and recall; either complete the annotations or implement an ignore mask that excludes unverified classes/regions from the classification loss using the custom trainer workflow.

[pick7]

For your examples, I would separate three problems that need different treatment.

1. Class imbalance: focal loss, class-balanced sampling, and per-class weighting can help, but they do not solve missing labels by themselves.
2. Partially labeled images: this is the dangerous case. If red blood cells are present but unlabeled, the detector can learn that valid RBC boxes are background. You need an ignore mechanism for unlabeled classes/regions, or train a detector whose label space matches what is actually annotated for that image.
3. Uncertain class labels: this is better modeled as “detect the object, but do not fully supervise the fine class” rather than as a normal class.

For the blood-cell setup, a practical pipeline may be more stable than one heavily customized YOLO loss:

• Train detection for “cell” or for the classes that are consistently annotated.
• Crop detections and train a separate classifier for WBC subtype / uncertain subtype handling.
• For uncertain cells, use a superclass such as eosinophil_uncertain or omit fine-class loss while still keeping box/object supervision.

If you do want to implement it inside Ultralytics, you will likely need a custom trainer/loss rather than a config-only change:

• add metadata per image or per target saying which classes are fully annotated,
• mask classification loss for “don't care” labels,
• prevent unlabeled but present classes from becoming hard negatives in the assigner/objectness path,
• optionally use soft targets for ambiguous class sets instead of one-hot labels.

The key part is the assigner/negative handling. Simply ignoring the class loss for one label is not enough if unmatched predictions over unlabeled RBCs are still penalized as false positives.

Label smoothing can help calibration, but it is not a good way to encode real blood-cell prevalence. If the training set prevalence is intentionally different from deployment prevalence, handle that with sampling, calibration/thresholding on a validation set with realistic prevalence, or post-training prior correction rather than expecting label smoothing to fix missing labels.

[glenn-jocher]

This is the right distinction: partially labeled images require masking both target assignment and classification negatives, not only the class-loss term. Implement per-image annotation metadata through a custom DetectionTrainer and override DetectionModel.init_criterion() with a modified v8DetectionLoss; the custom trainer guide provides the supported extension pattern. Validate against a fully annotated set with deployment-realistic prevalence, since label smoothing will not correct missing supervision.