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:

```
def weighted_categorical_crossentropy(weights):
""" weighted_categorical_crossentropy
Args:
* weights<ktensor|nparray|list>: crossentropy weights
Returns:
* weighted categorical crossentropy function
"""
if isinstance(weights,list) or isinstance(np.ndarray):
weights=K.variable(weights)
def loss(target,output,from_logits=False):
if not from_logits:
output /= tf.reduce_sum(output,
len(output.get_shape()) - 1,
True)
_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)
else:
raise ValueError('WeightedCategoricalCrossentropy: not valid with logits')
return loss
```

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.