Multi-class semantic segmentation using U-Net Error with Binary Cross Entropy with Logits

sarmientoj24: <p>My current implementation for the loss function of Pytorch Multi-class (class = 5) on my U-Net with Pre-trained ResNet is the ff:</p>
<pre><code class="lang-auto">def dice_loss(pred, target, smooth = 1.):
    pred = pred.contiguous()
    target = target.contiguous()    

    intersection = (pred * target).sum(dim=2).sum(dim=2)
    
    loss = (1 - ((2. * intersection + smooth) / (pred.sum(dim=2).sum(dim=2) + target.sum(dim=2).sum(dim=2) + smooth)))
    
    return loss.mean()

def calc_loss(pred, target, metrics, bce_weight=0.5):
    bce = F.binary_cross_entropy_with_logits(pred, target)

    pred = F.sigmoid(pred)
    dice = dice_loss(pred, target)

    loss = bce * bce_weight + dice * (1 - bce_weight)

    metrics['bce'] += bce.data.cpu().numpy() * target.size(0)
    metrics['dice'] += dice.data.cpu().numpy() * target.size(0)
    metrics['loss'] += loss.data.cpu().numpy() * target.size(0)

    return loss
</code></pre>
<p><strong>Error is the ff:</strong></p>
<pre><code class="lang-auto">   2432     if not (target.size() == input.size()):
-&gt; 2433         raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
   2434 
   2435     return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)

ValueError: Target size (torch.Size([4, 1, 320, 480, 3])) must be the same as input size (torch.Size([4, 5, 320, 480]))
</code></pre>

ptrblck: <p><code>F.binary_cross_entropy_with_logits</code> expects the model output and target to have the same shape and can be used for a multi-label segmentation (each pixel can belong to zero, one, or multiple classes).</p>
<p>Since you’ve mentioned a multi-class segmentation (each pixel belongs to one class only), you should use <code>nn.CrossEntropyLoss</code> instead (or <code>nn.NLLLoss</code> with <code>F.log_softmax</code> as the last non-linearity).</p>
<p><code>nn.CrossEntropyLoss</code> expects the model output to have the shape <code>[batch_size, nb_classes, height, width]</code> and the target <code>[batch_size, height, width]</code> containing the class indices in the range <code>[0, nb_classes-1]</code>.</p>
<p>Your target seems to have an additional dimension in dim1 and might be a color image in the channels-last format?<br>
If so, you would have to map the colors to class indices.</p>

sarmientoj24: <p>I already made a possible function to convert the mask image into n_channels of [n_classes, height, width] as a binary mask where background is 0 and 1. Also, each pixel only belongs to one class.</p>
<pre><code class="lang-auto">mask_palette = [1, 2, 3, 4, 5] #basically the values of each pixel converted to one channel. for each classes.

for palette in mask_palette:
  new_mask = np.zeros((mask_array.shape[0], mask_array.shape[1]), dtype=np.float32)
  new_mask[mask == palette] = 1
  final_mask.append(new_mask)

final_mask = np.asarray(final_mask)
</code></pre>
<p>final mask is <code>[5, 600, 900]</code></p>
<p>So this method shall create<br>
[batch_size, nb_classes, height, width] as a mask dataloader</p>
<p>and the same size as generated by my model.</p>
<p>What loss and activation should that be?</p>
<p>The final_mask is like the following</p>
<pre><code class="lang-auto">array([[[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 1., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]],

       [[0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        ...,
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.],
        [0., 0., 0., ..., 0., 0., 0.]]], dtype=float32)
</code></pre>

sarmientoj24: <p>Hi, if I am going to use CrossEntropyLoss, what should be the setup of the output and target?</p>
<p>Currently, I have 5 classes. If I make the mask a 1 channel image with class numbers, should I put 0 as background? So I shall have 0, 1, 2, 3, 4, 5 so 6 classes?</p>

ptrblck: <p>If you are using <code>nn.CrossEntropyLoss</code>, you should have to use this approach:</p>
<aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="85207">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><code>nn.CrossEntropyLoss</code> expects the model output to have the shape <code>[batch_size, nb_classes, height, width]</code> and the target <code>[batch_size, height, width]</code> containing the class indices in the range <code>[0, nb_classes-1]</code> .</p>
</blockquote>
</aside>
<p>Your target mask should not contain a channels dimension and for 5 classes should contain the class indices in the range <code>[0, 4]</code>. If the background class is not included in these 5 classes, then you should use 6. It doesn’t matter, if the background is class index 0 or any other.</p>

pyx2020: <aside class="quote no-group" data-username="sarmientoj24" data-post="1" data-topic="85207">
<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/sarmientoj24/48/30474_2.png" class="avatar"> sarmientoj24:</div>
<blockquote>
<p><code>* target.size(0)</code></p>
</blockquote>
</aside>
<p>Can you explain for me, why we should multiply target.size(0) in loss function?</p>

Mahdi_Farhadi: <p>Hello ,<br>
I want to use CrossEntropyLoss for my segmentation code but I confused!<br>
I calculated the loss in two ways and I found two <strong>difference results</strong>!!! But both of them recommended  in the the pytorch forums !<br>
<em>.e.g my code like it</em> :</p>
<pre><code class="lang-auto">y = torch.randint(low=0,high=4,size=( 2,1,3,3)).type(torch.LongTensor)
y_hat = torch.randn(2,4,3,3)
print('\ntarget shape like it  [batch_size,N_class,h,w] =',y.shape)
print('\npredict shape like it [batch_size,N_class,h,w] =',y_hat.shape)

print('\nloss without flatten = ', criterion(y_hat,y.squeeze(1)))
print('\nloss with flatten = ',criterion(y_hat.view(-1,4), y.view([-1, ])))
</code></pre>
<p><em><strong>and the output :</strong></em></p>
<pre><code class="lang-auto">target shape like it  [batch_size, 1 ,h,w] = torch.Size([2, 1, 3, 3])

predict shape like it [batch_size,N_class,h,w] = torch.Size([2, 4, 3, 3])

loss without flatten =  tensor(2.0514)

loss with flatten =  tensor(1.6351)
</code></pre>
<p>Which one is true ?</p>

ptrblck: <p>The <code>view</code> operation is wrong as it will interleave the tensors and you would need to <code>permute</code> it beforehand:</p>
<pre><code class="lang-python">criterion = nn.CrossEntropyLoss()
y = torch.randint(low=0,high=4,size=( 2,1,3,3)).type(torch.LongTensor)
y_hat = torch.randn(2,4,3,3)
print('\ntarget shape like it  [batch_size,N_class,h,w] =',y.shape)
print('\npredict shape like it [batch_size,N_class,h,w] =',y_hat.shape)

print('\nloss without flatten = ', criterion(y_hat,y.squeeze(1)))
print('\nloss with flatten = ',criterion(y_hat.permute(0, 2, 3, 1).contiguous().view(-1,4), y.permute(0, 2, 3, 1).contiguous().view([-1, ])))
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">target shape like it  [batch_size,N_class,h,w] = torch.Size([2, 1, 3, 3])

predict shape like it [batch_size,N_class,h,w] = torch.Size([2, 4, 3, 3])

loss without flatten =  tensor(1.6820)

loss with flatten =  tensor(1.6820)
</code></pre>