How to get logits as neural network output
Simple and short question. I have a network (Unet) which performs image segmentation. I want the logits as the output to feed into the cross entropy loss (using pytorch). Currently my final layer looks as so:
class Logits(nn.Sequential):
    def __init__(self, 
                 in_channels, 
                 n_class
                 ):
      super(Logits, self).__init__()

      # fully connected layer outputting the prediction layers for each of my classes
      self.conv = self.add_module('conv_out', 
                        nn.Conv2d(in_channels, 
                                  n_class, 
                                  kernel_size = 1
                                  )
                       )
      self.activ = self.add_module('sigmoid_out', 
                        nn.Sigmoid()
                       )

Is it correct to use the sigmoid activation function here? Does this give me logits?

When people talk about &quot;logits&quot; they usually refer to the &quot;raw&quot; n_class-dimensional output vector. For multi-class classification (n_class &gt; 2) you want to convert the n_class-dimensional vector of raw &quot;logits&quot; into a n_class-dim probability vector.
That is, you want prob = f(logits) with prob_i &gt;= 0 for all n_class entries, and that sum(prob)=1.
The most straight forward way of doing that in a differentiable way is to use the Softmax function:
prob_i = softmax(logits) = exp(logits_i) / sum_j exp(logits_j)

It is easy to see that the output of softmax is indeed a n_class-dim probability vector (I leave it to you as a short exercise).
BTW, this is why the raw predictions are called &quot;logits&quot; because they are kind of &quot;log&quot; of the output predicted probabilities.
Now, it is customary not to explicitly compute the softmax on top of a classification network and defer its computation to the loss function, e.g. nn.CrossEntropyLoss that internally computes the softmax and requires the raw logits as inputs, rather than the normalized probabilities. This is done mainly for numerical stability.
Therefore, if you are training a multi-class classification network with nn.CrossEntropyLoss you do not need to worry at all about the final activation and simply output the raw logits from your final conv/linear layer.
Most importantly, do not use nn.Sigmoid() activation as it tends to have saturated gradients and will mess up your training.
