Use CrossEntropyLoss() in multiclass semantic segmentation

SimCan: <p>I would like to know how to properly use <code>CrossEntropyLoss()</code> for the multiclass semantic segmentation task.</p>
<p>This is an implementation of <a href="https://paperswithcode.com/method/focal-loss" rel="noopener nofollow ugc">FocalLoss</a>, where I substituted the original <code>binary_cross_entropy()</code> with <code>CrossEntropyLoss()</code></p>
<pre><code class="lang-auto">
class FocalLoss(nn.Module):
    def __init__(self, weight=None, size_average=True):
        super(FocalLoss, self).__init__()

    def forward(self, inputs, targets, alpha=ALPHA, gamma=GAMMA, smooth=1):
        #inputs = torch.squeeze(inputs, 1)
        inputs = torch.as_tensor(inputs, dtype=torch.float32, device=torch.device('cuda'))
        targets = torch.as_tensor(targets, dtype=torch.float32, device=torch.device('cuda'))
        
        #comment out if your model contains a sigmoid or equivalent activation layer
        #inputs = torch.sigmoid(inputs)       
        #flatten label and prediction tensors
        inputs = inputs.view(-1)
        targets = targets.view(-1)
        
        #first compute binary cross-entropy 
        #BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
        BCE = nn.CrossEntropyLoss(inputs, targets, reduction='mean')
        BCE_EXP = torch.exp(-BCE)
        focal_loss = alpha * (1- BCE_EXP)**gamma * BCE
                       
        return focal_loss
</code></pre>
<p>When I run the above script I get this error:</p>
<pre><code class="lang-auto">/usr/local/lib/python3.7/dist-packages/torch/nn/_reduction.py in legacy_get_string(size_average, reduce, emit_warning)
     33         reduce = True
     34 
---&gt; 35     if size_average and reduce:
     36         ret = 'mean'
     37     elif reduce:

RuntimeError: Boolean value of Tensor with more than one value is ambiguous
</code></pre>
<ul>
<li>Should I change the target’s shape from (BxCxWxH) to (BxWxH)?</li>
<li>Is it necessary for the target to be one-hot-encoded?</li>
</ul>

KFrank: <p>Hi Simone!</p>
<aside class="quote no-group" data-username="SimCan" data-post="1" data-topic="158141" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/simcan/48/48222_2.png" class="avatar"> SimCan:</div>
<blockquote>
<pre><code class="lang-auto">        #BCE = F.binary_cross_entropy(inputs, targets, reduction='mean')
        BCE = nn.CrossEntropyLoss(inputs, targets, reduction='mean')
</code></pre>
</blockquote>
</aside>
<p>Your problem is that you are <em>instantiating</em> a <code>CrossEntropyLoss</code> object<br>
with garbled constructor arguments.  I would recommend using the<br>
functional form (as you had been doing with <code>binary_cross_entropy()</code>):</p>
<pre data-code-wrap="python"><code class="lang-python">        BCE = F.cross_entropy (inputs, targets, reduction='mean')
</code></pre>
<p>You could instantiate <code>CrossEntropyLoss</code> on the fly and then call it:</p>
<pre data-code-wrap="python"><code class="lang-python">BCE = nn.CrossEntropyLoss (reduction = 'mean') (inputs, targets)
</code></pre>
<p>but, stylistically, I prefer the functional form.</p>
<aside class="quote no-group">
<blockquote>
<ul>
<li>Should I change the masks’ shape from (BxCxWxH) to (BxWxH)?</li>
</ul>
</blockquote>
</aside>
<p>If your <code>inputs</code> (your predictions, i.e., the output of your model) has shape<br>
<code>[B, C, W, H]</code>, then your <code>targets</code> (which I assume you mean by “masks”)<br>
should have shape <code>[B, W, H]</code> (with no class dimension).  The values of<br>
<code>targets</code> should be <code>long</code> class labels that run from <code>0</code> to <code>C - 1</code>, where <code>C</code><br>
is the number of classes.</p>
<p>Also, your <code>inputs</code> should be raw-score <em>logits</em>, that is, the output of your<br>
model should be the output of its final <code>Linear</code> layer – <em>not</em> followed by<br>
<code>softmax()</code> (or <code>sigmoid()</code>).</p>
<p>Last, as an aside, unless you have a <em>very</em> large number of classes,<br>
consider using regular cross entropy as your loss criterion, using class<br>
weights if you have a significant class imbalance in your data.  There are<br>
also claims that you are likely to get better results using a focal-loss term<br>
as an add-on to cross-entropy compared to using focal loss alone.</p>
<p>Best.</p>
<p>K. Frank</p>

SimCan: <p><a class="mention" href="/u/kfrank">@KFrank</a> Thanks for the advice!<br>
Would you suggest using cross-entropy with weights instead of focal-loss?</p>

KFrank: <p>Hi Simone!</p>
<aside class="quote no-group" data-username="SimCan" data-post="3" data-topic="158141" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/simcan/48/48222_2.png" class="avatar"> SimCan:</div>
<blockquote>
<p>Would you suggest using cross-entropy with weights instead of focal-loss?</p>
</blockquote>
</aside>
<p>I would certainly suggest at least trying cross entropy.  If you have a<br>
significant class imbalance, then try using <code>CrossEntropyLoss</code>’s <code>weight</code><br>
constructor argument.</p>
<p>Focal loss is a perfectly reasonable loss criterion, and could well work for<br>
you, especially if you very large number of classes.  (How many classes do<br>
you have?)  But I would only use it if you can show that it gives you better<br>
results than cross entropy for your specific use case.</p>
<p>Best.</p>
<p>K. Frank</p>

SimCan: <p>I have only 3 classes, therefore I will try <code>CrossEntropyLoss</code> as suggested.</p>
<p>Furthermore, I would try  <code>dice loss</code>, but I am not sure how it could work with a significant class imbalance as in my case.</p>