I’m also doing image segmentation with unbalanced classes but have been working in Keras, so I haven’t used the pytorch BCEWITHLOGITSLOSS.
Keras doesn’t have a weights parameter but I wrote my own (simply by copying the Keras source code for categorical-crossentropy and adding weight parameter).
All that said - I did see improvements in my results. If
BCEWITHLOGITSLOSS isn’t behaving as you expect maybe you could try writing your own pytorch version to ensure its doing what you want it too.
If it helps here’s what I did in Keras:
* weights<ktensor|nparray|list>: crossentropy weights
* weighted categorical crossentropy function
if isinstance(weights,list) or isinstance(np.ndarray):
if not from_logits:
output /= tf.reduce_sum(output,
len(output.get_shape()) - 1,
_epsilon = tf.convert_to_tensor(K.epsilon(), dtype=output.dtype.base_dtype)
output = tf.clip_by_value(output, _epsilon, 1. - _epsilon)
weighted_losses = target * tf.log(output) * weights
return - tf.reduce_sum(weighted_losses,len(output.get_shape()) - 1)
raise ValueError('WeightedCategoricalCrossentropy: not valid with logits')
I didn’t implement the logits part since it wasn’t immediately clear how to do it and I didn’t need it.
The other piece of the puzzle is of course picking the correct weights. I didn’t think too deeply about this but I found a stackoverflow answer where they took the
log of the ratio
(total_for_all_categories/total_for_category) because the different categories were imbalanced by many orders of magnitude. I tried both - with the log and without the log and got better results with the log version.
Hopefully that rambling is a little helpful.