Loss function for multi-class semantic segmentation

Vannila: <p>I’m doing a semantic segmentation problem where each pixel may belong to one or more classes. However, I cannot find a suitable loss function to compute binary crossent loss over each pixel in the image. BCELoss requires a single scalar value as the target, while CrossEntropyLoss allows only one class for each pixel.<br>
Is there any built-in loss for this problem (similar to binary_crossentropy in Keras), or I need to write a new loss?<br>
Thank you very much.</p>

ptrblck: <p>If you are dealing with a multi-label classification, <code>nn.BCELoss</code> should work fine, if you pass the target as a “multi-hot” encoded tensor.</p>
<pre><code class="lang-python">target = torch.tensor([[0,1,0,1,0,0]], dtype=torch.float32)
output = torch.randn(1, 6, requires_grad=True)

criterion = nn.BCEWithLogitsLoss()
loss = criterion(output, target)
loss.backward()
</code></pre>
<p>In this small example, I just passed a single data sample with a target for two active classes.<br>
For a semantic segmentation use case, each pixel should contain the corresponding label (0 or 1) in the channel dimension.</p>

Vannila: <p>My tatget and output are of shape [B, C, H, W].  Positive weights is of shape [C].</p>
<pre><code class="lang-auto">target = torch.ones([1, 2, 5, 5], dtype=torch.float32)
output = torch.randn(1, 2, 5, 5, requires_grad=True)
positive_weights = torch.FloatTensor([2, 2])
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=positive_weights)
loss = criterion(output, target)
loss.backward()
</code></pre>
<p>The error is:</p>
<pre><code class="lang-auto">RuntimeError: The size of tensor a (2) must match the size of tensor b (5) at non-singleton dimension 3
</code></pre>
<p>How to fix this?</p>

MariosOreo: <p>It seems that the problem occurs in <code>pos_weight = positive_weights</code>.<br>
From the docs</p>
<ul>
<li>
<strong>pos_weight</strong>  (<a href="https://pytorch.org/docs/master/tensors.html#torch.Tensor" rel="nofollow noopener"> <em>Tensor</em> </a> <em>,</em>  <em>optional</em> ) – a weight of positive examples. Must be a vector with length equal to the number of classes.</li>
</ul>
<p>But why does it consider 5 as the number of classes?<img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=7" title=":thinking:" class="emoji" alt=":thinking:"></p>

ptrblck: <p>As <a class="mention" href="/u/mariosoreo">@MariosOreo</a> said, it seems the <code>pos_weight</code> argument throws this error.<br>
A quick fix might be to permute and view the output and target such that the two classes are in dim1:</p>
<pre><code class="lang-python">loss = criterion(output.permute(0, 2, 3, 1).view(-1, 2), target.permute(0, 2, 3, 1).view(-1, 2))
</code></pre>
<p>or to expand the <code>pos_weight</code> tensor manually:</p>
<pre><code class="lang-python">positive_weights = positive_weights[:, None, None].expand(-1, 5, 5)
criterion = torch.nn.BCEWithLogitsLoss(pos_weight=positive_weights)
</code></pre>
<p>However, it seems like unclear behavior to me, so feel free to post a <a href="https://github.com/pytorch/pytorch/issues" rel="nofollow noopener">Github issue</a> to further discuss this use case.</p>

velikodniy: <p>Actually it’s a bug. When <code>pos_weight</code> was added to BCEWithLogits loss it wasn’t supposed to be used with per-pixel classifiers. So, broadcasting doesn’t work well in this case.</p>
<p>You can reshape <code>pos_weight</code> by adding two dummy dimensions for W and H to to work around this bug.</p>
<pre><code class="lang-auto">positive_weights = torch.FloatTensor([2, 2]).reshape(1, 2, 1, 1)
</code></pre>
<p>Unfortunately my code with implementation of pos_weight was rewritten from Python to C++ and I’m not sure if I can fix it right away.</p>

ADONIS_HAMMOUTENE: <p>Hey, frist of all, if each pixel may belong to one or more classes, it means that you are in a multilabel segmentation task and not multi-class because multiclass, a pixel belongs to one of the classes.<br>
I don’t know if you’ve thought of using the segmentation-models-pytorch (smp) library that also has a package of losses including DiceLoss, Focal Loss and Tversky Loss and you can choose the mode of loss : 'binary, multi-label, multi-class"  : <a href="https://smp.readthedocs.io/en/latest/losses.html" rel="noopener nofollow ugc">SMP’s defined losses</a></p>