Changing the Loss Function in pretrained models

Hi There! I’m a current beginner and wanted to ask a general question regarding the YOLOv8 model (and v11 as well) to see if it was possible to override the loss function (More specifically, BBoxLoss) when training on my dataset while still maintaining the other pretrained weights. If not, would I have to implement my model from scratch?

Thank you!

This:

plus:

in most cases means you shouldn’t be attempting this, but that really depends on your goal (which wasn’t stated).

If you’re looking to learn where to make changes to the loss for learning purposes, you can check here in the source code for the v8DetectionLoss class. If however your goal is to improve your model’s performance when training, it is very likely that you should not attempt to change something related to the loss as a first or early step, as there are several other more important things to attempt as a priority.

My apologies. I realized my last post was extremely vague. Let me provide some clarification:

I am currently working in a project where we are detecting small defects (cracks, missing bolts). The subset of the dataset I’m training on is ~3,000 images. My previous task was to look to modify the loss function (more specifically, BBoxLoss) in order for it to focus more on smaller objects. Please let me know what steps I should take first / what suggestions you have to improve the model’s performance, as well as any general recommendations (such as what model I should use). Please also let me know if this post is still vague.

Thank you again for your time.

The additional context is helpful. A few more details could help, and not that you have to reply to them, but they are likely going to be useful to consider:

  • Are you using a detection model or segmentation model?
    • I would suggest using a segmentation model for cracks especially.
  • What are the pixel dimensions of the images?
  • What’s considered “small” for your object size? Do you mean something like 20 x 20 pixels for an image that’s 1600 x 1200 or do you mean 4 x 6 pixels for an image that’s 1200 x 900?
  • What imgsz are you training at?
  • What imgsz are you using for prediction?
  • Have you explored tried image slicing?
  • Have you adjusted any of the image augmentation hyperparameters?
  • What’s the current performance of the model?
    • Anywhere in particular it’s not meeting performance expectations (low mAP, high FP, misclassification, etc.)
  • What is the performance target, or how do you know when you’ve achieved the model performance you need?

If your end goal is a better model, there’s probably many ways to improve model performance before attempting to make custom modifications the model loss. Something else to consider is that the quality and consistency of your ground-truth annotations will be critical in how well the model is able to detect and classify correctly, so I would recommend verifying these as a primary task. After that, you could work thru some of the ideas in the bullets above, as those are the best places to start before attempting to use a custom loss (unless you have some sort of evidence showing that’s what’s needed, but if you have such justification, it’s likely you’d know how to modify the loss already).