UNet Multiclass Segmentation - Cross Entropy Softmax

Niranjan_Rao: <p>Following is my UNet model for Multi Class Segmentation for 4 classes.</p>
<pre><code class="lang-python">class Unet(nn.Module):
    
    def __init__(self):
        super(Unet, self).__init__()    
        
        #(Double) Convolution 1        
        self.double_conv1 = double_conv(1, start_fm, 3, 1, 1)
        #Max Pooling 1
        self.maxpool1 = nn.MaxPool2d(kernel_size=2)
        
        #Convolution 2
        self.double_conv2 = double_conv(start_fm, start_fm * 2, 3, 1, 1)
        #Max Pooling 2
        self.maxpool2 = nn.MaxPool2d(kernel_size=2)
        
        #Convolution 3
        self.double_conv3 = double_conv(start_fm * 2, start_fm * 4, 3, 1, 1)
        #Max Pooling 3
        self.maxpool3 = nn.MaxPool2d(kernel_size=2)
        
        #Convolution 4
        self.double_conv4 = double_conv(start_fm * 4, start_fm * 8, 3, 1, 1)
        #Max Pooling 4
        self.maxpool4 = nn.MaxPool2d(kernel_size=2)
        
        #Convolution 5
        self.double_conv5 = double_conv(start_fm * 8, start_fm * 16, 3, 1, 1)
        
        #Transposed Convolution 4
        self.t_conv4 = nn.ConvTranspose2d(start_fm * 16, start_fm * 8, 2, 2)
        # Expanding Path Convolution 4 
        self.ex_double_conv4 = double_conv(start_fm * 16, start_fm * 8, 3, 1, 1)
        
        #Transposed Convolution 3
        self.t_conv3 = nn.ConvTranspose2d(start_fm * 8, start_fm * 4, 2, 2)
        #Convolution 3
        self.ex_double_conv3 = double_conv(start_fm * 8, start_fm * 4, 3, 1, 1)
        
        #Transposed Convolution 2
        self.t_conv2 = nn.ConvTranspose2d(start_fm * 4, start_fm * 2, 2, 2)
        #Convolution 2
        self.ex_double_conv2 = double_conv(start_fm * 4, start_fm * 2, 3, 1, 1)
        
        #Transposed Convolution 1
        self.t_conv1 = nn.ConvTranspose2d(start_fm * 2, start_fm, 2, 2)
        #Convolution 1
        self.ex_double_conv1 = double_conv(start_fm * 2, start_fm, 3, 1, 1)
        
        # One by One Conv
        self.one_by_one = nn.Conv2d(start_fm, 4, 1, 1, 0)
        #self.final_act = nn.Sigmoid()
        
        
    def forward(self, inputs):
        # Contracting Path
        conv1 = self.double_conv1(inputs)
        maxpool1 = self.maxpool1(conv1)

        conv2 = self.double_conv2(maxpool1)
        maxpool2 = self.maxpool2(conv2)

        conv3 = self.double_conv3(maxpool2)
        maxpool3 = self.maxpool3(conv3)

        conv4 = self.double_conv4(maxpool3)
        maxpool4 = self.maxpool4(conv4)
            
        # Bottom
        conv5 = self.double_conv5(maxpool4)
        
        # Expanding Path
        t_conv4 = self.t_conv4(conv5)
        cat4 = torch.cat([conv4 ,t_conv4], 1)
        ex_conv4 = self.ex_double_conv4(cat4)
        
        t_conv3 = self.t_conv3(ex_conv4)
        cat3 = torch.cat([conv3 ,t_conv3], 1)
        ex_conv3 = self.ex_double_conv3(cat3)

        t_conv2 = self.t_conv2(ex_conv3)
        cat2 = torch.cat([conv2 ,t_conv2], 1)
        ex_conv2 = self.ex_double_conv2(cat2)
        
        t_conv1 = self.t_conv1(ex_conv2)
        cat1 = torch.cat([conv1 ,t_conv1], 1)
        ex_conv1 = self.ex_double_conv1(cat1)
        
        one_by_one = self.one_by_one(ex_conv1)
        
        return one_by_one

</code></pre>
<p>Should I be using a softmax layer for getting class probabilities while using Cross-Entropy Loss.<br>
How do I convert Logits to Probabilities.</p>

KFrank: <p>Hi Niranjan!</p>
<p>My answer is not specific to UNet, and I haven’t looked at your code<br>
or your particular use case.</p>
<aside class="quote no-group" data-username="Niranjan_Rao" data-post="1" data-topic="113219" 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/niranjan_rao/48/29758_2.png" class="avatar"> Niranjan_Rao:</div>
<blockquote>
<p>Should I be using a softmax layer for getting class probabilities while using Cross-Entropy Loss.</p>
</blockquote>
</aside>
<p>No.  <code>CrossEntropyLoss</code> has, in effect, <code>softmax()</code> built in.  So you<br>
want to feed into it the raw-score logits output by your model.</p>
<aside class="quote no-group">
<blockquote>
<p>How do I convert Logits to Probabilities.</p>
</blockquote>
</aside>
<p>You usually don’t actually need the probabilities.  But if you do, you<br>
convert logits to probabilities by passing them through <code>softmax()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

Niranjan_Rao: <p>Thank you for that, its clear now, however my model output is [batch_size, nb_classes=4, height, width].<br>
When making predictions on a new image using model(image), I get the raw logits score. How do I choose the most “probable” class for each pixel.</p>
<p>Do I have to take some kind of argmax on the nb_classes dimension?</p>
<p><strong>Reminder</strong> - Multi Class Segmentation(4 classes)</p>

KFrank: <p>Hi Niranjan!</p>
<aside class="quote no-group" data-username="Niranjan_Rao" data-post="3" data-topic="113219" 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/niranjan_rao/48/29758_2.png" class="avatar"> Niranjan_Rao:</div>
<blockquote>
<p>Do I have to take some kind of argmax on the nb_classes dimension?</p>
</blockquote>
</aside>
<p>Yes, this is exactly right.  The algebraically largest logit corresponds<br>
to the most probable class.</p>
<p>If you have:</p>
<pre data-code-wrap="python"><code class="lang-python">logits.shape == torch.Size ([batch_size, nb_classes=4, height, width])
</code></pre>
<p>then <code>logits.argmax (dim = 1)</code> will be a <code>LongTensor</code> of shape<br>
<code>[batch_size, height, width]</code>, with the <code>nb_classes</code> dimension<br>
“reduced” away.</p>
<p>Best.</p>
<p>K. Frank</p>