Unet pixel-wise weighted loss function

username: SamT
cooked: <p>Hi, did anyone successfully implemented the pixel-wise weighted loss function that was described in the Unet paper?<br>
Thanks</p>

username: Nikronic
cooked: <p>Hi,<br>
Yes of course.<br>
Here is some link to implementation of Github community. They have trained their model and used, so they have to be correct.<br>
link1: <a href="https://github.com/milesial/Pytorch-UNet" rel="nofollow noopener">https://github.com/milesial/Pytorch-UNet</a><br>
link2: <a href="https://github.com/usuyama/pytorch-unet" rel="nofollow noopener">https://github.com/usuyama/pytorch-unet</a><br>
link3: <a href="https://github.com/jvanvugt/pytorch-unet" rel="nofollow noopener">https://github.com/jvanvugt/pytorch-unet</a><br>
link4: <a href="https://www.kaggle.com/witwitchayakarn/u-net-with-pytorch" rel="nofollow noopener">https://www.kaggle.com/witwitchayakarn/u-net-with-pytorch</a></p>
<p>If you did not understand any part or you have other questions, feel free to ask.</p>

username: SamT
cooked: <p>Hi Nikronic, Thanks for the links!<br>
However, None of these Unet implementation are using the pixel-weighted soft-max cross-entropy loss that is defined in the <a href="https://arxiv.org/pdf/1505.04597.pdf" rel="nofollow noopener">Unet paper</a> (page 5).</p>
<p>I’ve tried to implement it myself using a modified version of <a href="https://stackoverflow.com/questions/50255438/pixel-wise-loss-weight-for-image-segmentation-in-keras" rel="nofollow noopener">this code</a> to compute the weights which I multiply by the CrossEntropyLoss:</p>
<pre><code class="lang-auto">loss = nn.CrossEntropyLoss(reduction='none')(output, target)
loss = torch.mean(loss * weights)
</code></pre>
<p>where the output is the output tensor from the Unet.</p>
<p>When I am using this implementation my model is a complete failure (F1 score = 0).</p>
<p>So, I’ve replace the CrossEntropyLoss by:</p>
<pre><code class="lang-auto">m = nn.Sigmoid()
loss_per_pixel = torch.nn.BCELoss(reduce=False)(m(input), target) * weights
</code></pre>
<p>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.</p>
<p>So, I am wondering if I am doing something wrong.</p>
<p>Thanks for your help!</p>

username: ptrblck
cooked: <p>I’ve implemented a weighted pixelwise NLLLoss a while ago in <a href="https://discuss.pytorch.org/t/weighted-pixelwise-nllloss2d/7766">this topic</a>. The code is quite old (e.g. don’t use <code>Variables</code> anymore), but might be a good starter for your custom loss.</p>

username: nicolo-lucchesi
cooked: <p>Hey,<br>
this is an old thread but I believe Sam’s formulation is right:</p>
<aside class="quote no-group" data-username="SamT" data-post="3" data-topic="46689">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/c57346/48.png" class="avatar"> SamT:</div>
<blockquote>
<pre><code class="lang-auto">loss = nn.CrossEntropyLoss(reduction='none')(output, target)
loss = torch.mean(loss * weights)
</code></pre>
</blockquote>
</aside>
<p>I’ve used this loss myself in the past and it worked out fairly well, assuming the scale of the weights is<br>
“correct”. Make sure to pair it with something like a dice loss so that the CE can be used to incentivize precision around the edges.</p>