Hi Nikronic, Thanks for the links!
However, None of these Unet implementation are using the pixel-weighted soft-max cross-entropy loss that is defined in the Unet paper 70 (page 5).
I’ve tried to implement it myself using a modified version of this code 93 to compute the weights which I multiply by the CrossEntropyLoss:
loss = nn.CrossEntropyLoss(reduction='none')(output, target)
loss = torch.mean(loss * weights)

where the output is the output tensor from the Unet.
When I am using this implementation my model is a complete failure (F1 score = 0).
So, I’ve replace the CrossEntropyLoss by:
m = nn.Sigmoid()
loss_per_pixel = torch.nn.BCELoss(reduce=False)(m(input), target) * weights

In this case my model converge to a relatively good result but this result is less accurate than if I’ve used a non-weighted loss function.
So, I am wondering if I am doing something wrong.
Thanks for your help!