3
$\begingroup$

I'm training a Tensorflow model that receives an image and segments the image into foreground and background. That is, if the input image is w x h x 3, then the neural network outputs a w x h x 1 image of 0's and 1's, where 0 represents background and 1 represents foreground.

I've computed that about 75% of the true mask is background, so the neural network simply trains a model that outputs all 0's and gets a 75% accuracy.

To solve this, I'm thinking of implementing a custom loss function that checks if there are more than a certain percentage of 0's, and if so, to add a very large number to the loss to disincentivize the all 0's strategy.

The issue is that this loss function becomes non-differentiable.

Where should I go from here?

Improve this question
$\endgroup$

1 Answer 1

Reset to default
2
$\begingroup$

The background being an unbalance class is a well known problem in image segmentation. Before digging into custom losses you should take a look to existing ones that address this specific issue like the Dice Loss or Focal Loss, the latter being more tunable having a extra hyper parameter that can be optimized. You can easily find on github tensorflow implementations of both.
For a more detailed comparison and reference to other similar losses you can also check this paper.

Improve this answer
$\endgroup$
1
  • $\begingroup$ Can you briefly summarize Dice Loss and Focal Loss to make your answer more self-contained? $\endgroup$
    – Josiah Yoder
    Commented Dec 27, 2021 at 21:16

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .