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?


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 at 14:48
  • $\begingroup$ Yes, you will not learn anaything from these images so I think so $\endgroup$
    – y-prudent
    Jul 21 at 13:41

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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