PyTorch: Multi-class segmentation loss value != 0 when using target image as the prediction

user_10371862 0
Viewed 354

I was performing semantic segmentation using PyTorch. There are a total of 103 different classes in the dataset and the targets are RGB images with only the Red channel containing the labels. I was using nn.CrossEntropyLoss as my loss function. For sanity, I wanted to check if using nn.CrossEntropyLoss is correct for this problem and whether it has the expected behaviour

I pick a random mask from my dataset and create a categorical version of it using this custom transform

class ToCategorical:
    def __init__(self, n_classes: int) -> None:
        self.n_classes = n_classes

    def __call__(self, sample: torch.Tensor):
        mask = sample.permute(1, 2, 0)
        categories = torch.unique(mask).tolist()[1:]  # get all categories other than 0
        # build a tensor with `n_classes` channels
        one_hot_image = torch.zeros(self.n_classes, *mask.shape[:-1])
        for category in categories:
            # get spacial locs where the categ is present
            rows, cols, _ = torch.where(mask == category)
            # in same spacial loc but in `categ` channel fill 1
            one_hot_image[category, rows, cols] = 1  
        return one_hot_image

And then I send this image as the output (prediction) and use the ground truth mask as the target to the loss function.

import torch.nn as nn

mask = T.PILToTensor()(Image.open("path_to_image").convert("RGB"))
categorical_mask = ToCategorical(103)(mask).unsqueeze(0)
mask = mask[0].unsqueeze(0)  # get only the red channel, add fake batch_dim

loss_fn = nn.CrossEntropyLoss()

target = mask
output = categorical_mask

print(output.shape, target.shape)
print(loss_fn(output, target.to(torch.long)))

I expected the loss to be zero but to my surprise, the output is as follows

torch.Size([1, 103, 600, 800]) torch.Size([1, 600, 800])
tensor(4.2836)

I verified with other samples in the dataset and I obtained similar values for other masks as well. Am I doing something wrong? I expect the loss to be = 0 when the output is the same as the target.

PS. I also know that nn.CrossEntropyLoss is the same as using log_softmax followed by nn.NLLLoss() but even I obtained the same value by using nllloss as well

For Reference

Dataset used: UECFoodPixComplete

image-segmentation pytorch deep-learning computer-vision
1 Answers

I would like to adress this:

I expect the loss to be = 0 when the output is the same as the target.

If the prediction matches the target, i.e. the prediction corresponds to a one-hot-encoding of the labels contained in the dense target tensor, but the loss itself is not supposed to equal to zero. Actually, it can never be equal to zero because the nn.CrossEntropyLoss function is always positive by definition.

Let us take a minimal example with number of #C classes and a target y_pred and a prediction y_pred consisting of prefect predictions:

As a quick reminder:

  • The softmax is applied on the logits (q_i) as p_i = log(exp(q_i)/sum_j(exp(q_j)):

    >>> p = F.softmax(y_pred, 1)

    Similarly if you are using the log-softmax, defined as logp_i = log(p_i):

    >>> logp = F.log_softmax(y_pred, 1)

  • Then comes the negative likelihood function computed between x the input and y the target: -y*x. In association with the softmax, it comes down to -y*p, or -y*logp respectively. In any case, whether you apply the log or not, only the predictions corresponding to the true classes will remain since the others ones are zeroed-out.

That being said, applying the NLLLoss on y_pred would indeed result with a 0 as you expected in your question. However, here we apply it on the probability distribution or log-probability: p, or logp respectively!

In our specific case, p_i = 1 for the true class and p_i = 0 for all other classes (there are #C - 1 of those). This means the softmax of the logit associated with the true class will equal to exp(1)/sum_i(p_i). And since sum_i(p_i) = (#C-1)*exp(0) + exp(1). We therefore have:

softmax(p) = e / (#C - 1 + e)

Similarly for log-softmax:

log-softmax(p) = log(e / (#C-1 + e)) = 1 - log(#C - 1 + e)

If we proceed by applying the negative likelihood function we simply get cross-entropy(y_pred, y_true) = (nllloss o log-softmax)(y_pred, y_true). This results in:

loss = - (1 - log(#C - 1 + e)) = log(#C - 1 + e) - 1

This effectively corresponds to the minimum of the nn.CrossEntropyLoss function.

Regarding your specific case where #C = 103, you may have an issue in your code... since the average loss should equal to log(102 + e) - 1 i.e. around 3.65.

>>> y_true = torch.randint(0,103,(1,1,2,5))
>>> y_pred = torch.zeros(1,103,2,5).scatter(1, y_true, value=1)

You can see for yourself with one of the provided methods:

  • the builtin function nn.functional.cross_entropy:

    >>> F.cross_entropy(y_pred, y_true[:,0])
tensor(3.6513)

  • manually computing the quantity:

    >>> logp = F.log_softmax(y_pred, 1)
>>> -logp.gather(1, y_true).mean()
tensor(3.6513)

  • analytical result:

    >>> log(102 + e) - 1
3.6513
user_63313690
Related