Binary Unet Segmentation BCEWithLogitsLoss


Just wondering if anyone has an example of Unet for binary segmentation using BCEWithLogitsLoss ?

I’m segmenting foreground vs background and there are many more 0s than 1s due to this. It looks like i should be using BCEWithLogitsLoss as my loss function, however using fastai this doesnt plug and play super well.

Do i need to define a new Unet to work with this loss function?

Also, this might be out of scope of this forum so i can ask on pytorch too, but it looks like this function takes as an argument the proportion of class imbalance - “pos_weight (Tensor, optional) – a weight of positive examples. Must be a vector with length equal to the number of classes.”

Does that mean i need to know across all my images how many pixels are 1 v 0 :sweat_smile:

Thanks again!

You should be able to use it as loss function, dice worked better in my case though. You just pass the loss function parameter when you create unet, if you want positive weight then you can use partial. I did calculate the ratio of classes of all dataset but then run it with multiple values and others worked better