What is the difference between classes and ignore_index in some losses?

[DimitrisMantas]

Some losses, such as DiceLoss, accept both a classes and an ignore_index argument.

I've noticed that I'm getting different results for seemingly equivalent configurations (e.g., classes = [1, 2], assuming a three-class problem, and ignore_index = 0), but I can't wrap my head around why.

[DimitrisMantas]

Here's a gist showcasing my concern: https://gist.github.com/DimitrisMantas/658e16bf9c651588474550a1ad5d76ca

[isaaccorley]

I haven't had a chance to look at your gist yet. If you have a (C, H, W) prediction, classes allows you to ignore the loss for one or more of the these. ignore_index is different that it ignores the loss for pixels in the ground truth across all classes. This allows you to ignore pixels of a class present in the ground truth or that may be invalid (-1000, -1, etc.). So setting classes=[1, 2] and ignore_index=0 aren't going to be equivalent.

[DimitrisMantas]

It makes more sense in my head to picture this after the argmax step, so please bare with me; so, classes basically allows you to ignore predictions for a certain class wherever they end up the prediction mask where as ignore_index specific ground truth pixels, like you said.

So, if you wanted to ignore the background in the ground truth such that no matter what the model predicted there it wouldn't matter, you would use ignore_index, since with classes it could happen that the model predicted background for an otherwise valid pixel, which, incorrectly, wouldn't be penalized.

The implication if this is, if I have it correct, is that the model will actually learn the background with classes and will be "free" to output it.

[isaaccorley]

Yes that's correct. Ignore index operates the same way as in torch.nn.CrossEntropyLoss in that it's to identify pixels that we don't want to penalize the model for. The classes code basically computes the loss for all classes and then removes the desired class index entirely before performing the reduction e.g. mean. So the model is still able to output a prediction for the class but it won't be used in the loss function.

A use case for this that I've used is when I use a CE + Dice/Jaccard combination loss. So for CE I may still want the model to classify the background per pixel, but for Dice/Jaccard I prefer not to penalize the model for the background overlap with the ground truth.

[notprime]

Just to summarize what classes and ignore_index do at the moment. We are currently improving all the losses in torchseg, so let's use this discussion to share ideas.

At the moment, with the old code from smp:

• ignore_index: basically converts the pixels in the ground truth with the specified value to 0. For instance, you may have a mask with values ranging from 0 to 3 (so, 4 classes), you may set ignore_index = 2, and all pixels in the ground truth with value == 2 will be converted to 0. So those pixels will actually contribute to the loss, as they will be considered as background.
• classes: the loss is computed across all classes as usual. Before reducing the loss, only the chosen classes will be selected from the loss and will be returned. For instance, let's say we have 3 classes (so indices 0, 1 and 2), and we pass class = [0, 1]: the loss will be computed across all classes (0, 1 and 2), but only classes [0, 1] will be returned.

The current implementation of ignore_index is obviously wrong as we consider ignore_index as a specified class index that should not contribute to the loss, in any way. So we definitely need to change that. A simple idea is to create a specific channel for the index to be ignored, so that we can immediately drop it from all computations. For example, let's say you have a mask with values 0, 1 and 2, and a nodata value in the mask of -100. By passing ignore_index = -100 we should generate a one_hot representation of 4 channels (one for 0, for 1, for 2 and for -100 respectively), drop the channel of -100 and compute the loss. In that way nodata value is actually eliminated from the loss.

Regarding classes, the idea is that you may want to obtain the loss only for some specific classes, even when you have no nodata values or no index to ignore. It may be easier to select reduction = 'none' (we are currently implementing reduction techniques) and just select the desired values, but we may aswell leave classes for readibility. Obviously, being the classes extracted after the loss is computed, even excluded classes will contribute to the backpropagation, being already in the computational graph.

Finally, regarding ignoring the background, we should follow a similar approach to ignore_index, because, again, if we just use classes the predticions over the background will contribute to the backpropagation. So, a better idea would be to add a specific flag ignore_background, similar to ignore_index. If set to True, we can simply drop the first channel (the background) and proceed with the loss computation. As an example, let's say we have a mask and a prediction both of shape [B, C, H , W], we can simply do the following:

mask = mask[:, 1:]
prediction = prediction[:, 1:]

and then compute the loss.

Let me know your thoughts

[DimitrisMantas]

This is the ignore_background implementation that was recently introduced in TorchMetrics for mIoU:

def _ignore_background(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]:
    """Ignore the background class in the computation assuming it is the first, index 0."""
    preds = preds[:, 1:] if preds.shape[1] > 1 else preds
    target = target[:, 1:] if target.shape[1] > 1 else target
    return preds, target

It appears to be somewhat equivalent to classes?

By the way, it does this filtering and then proceeds to compute the metric, so not exactly what we're doing...

[isaaccorley]

Interesting. I'm not sold on the way we are filtering classes (we are just using the same implementation from smp). I'd like to keep the ignore_index implementation the same as PyTorch's.

For your case I would recommend transforming all classes in your mask that you want to ignore to -1 and then setting ignore_index=1. I typically don't ever really use the classes implementation tbh except for JaccardLoss to ignore the background.

[notprime]

That ignore_background can be easily implemented, I've seen it in multiple libraries, it basically drops the first channel immediately from both prediction and ground truth before computing the loss. Btw it's similar to classes, but different. In ignore_background the background channel is immediately excluded, with classes the loss is computed across all channels, only before reduction the desired classes are selected.

I believe the issue is just that the ignore_index "function contained in smp isn't implemented correctly, sticking to pytorch's implementation os definitely what I was looking for (but first i need to find exactly the pytorch file where it is implemented lol 💀)

[isaaccorley]

Yeah I'm wondering if the better choice is to just leave the classes and ignore_index implementation as is for backwards compatibility and implement ignore_background

[DimitrisMantas]

I actually ended up with MONAI which implements include_background.

In my case, the results at the end of training in terms of both validation loss and performance metrics were more or less equivalent to our ignore_index.