So, I am working on a semantic segmentation task using U-Net. The dataset is very unbalanced, with the background being by far the most common, and the last class being very scarce. First I trained it using Categorical Cross Entropy as the loss function, and in the end it simply classified everything as background (I used IoU as a measurement of success, and the confusion matrix had non-null values only on the first column, which can only mean that).

I also tried precomputing balanced weights using sklearn, however Keras does not accept precomputed weights for data that has more than three dimensions. I then tried using Focal Loss instead of CCE, because I saw it is very robust with imbalanced data, but I had the exact same results (to be fair, I only ran it for 25 epochs and the loss value was still very high and far from converging, so I will let it run again with 50 while I sleep, but I'm exploring other options in the meantime). Now I am running it with a DICE-based loss and hoping to see better results.

In case none of those work, is there anything you would recommend to improve the issue with imbalance?

  • $\begingroup$ Did you manage to solve the problem? I'm having the same difficulty, in an imbalanced dataset. $\endgroup$ Aug 29, 2023 at 13:11

2 Answers 2


Are you using some data augmentation with random crops / rotations / zooms ? If you do, you might have some images with only background labels and if so I would suggest you to add a condition to only retain the patches with a ratio of non-background pixels above a certain threshold value.

  • $\begingroup$ Thanks for your answer. I did not do any augmentation, but I do have a few images that are only background. Would it be a good idea to remove them then? $\endgroup$
    – BMC98
    Jul 20, 2021 at 14:48
  • $\begingroup$ Yes, you will not learn anaything from these images so I think so $\endgroup$
    – y-prudent
    Jul 21, 2021 at 13:41

This answer might be a little late, but I believe that what you need is the Focal Tversky Loss (https://arxiv.org/abs/1810.07842).

Neither the vanilla Focal Loss nor the Dice Loss generalize well for this kind of problems. And it gave me good results, some time ago, in identifying the same small object across many pictures.

According to the authors, one must avoid missing any of the pixels of the imbalanced class (ideally, 0% FN) in detriment of misclassifying background class pixels (FPs) to improve recall:

One of the limitations of the Dice loss function is that it equally weighs false positive (FP) and false negative (FN) detections. In practice, this results in segmentation maps with high precision but low recall. With highly imbalanced data and small ROIs such as skin lesions, FN detections need to be weighted higher than FPs to improve recall rate.

With a quick google search, I could find you a Keras implementation of this Focal Tversky Loss function: https://www.kaggle.com/code/bigironsphere/loss-function-library-keras-pytorch#Focal-Tversky-Loss


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.