trainers.segmentation JaccardLoss receiving num_classes, should be a List[int]?

[robmarkcole]

Description

From the trainer:

        elif loss == "jaccard":
            self.criterion = smp.losses.JaccardLoss(
                mode="multiclass", classes=self.hparams["num_classes"]
            )

num_classes is an int, Yet the docs for JaccardLoss state this is expected to be a list of int, presumably the indices of the output channels to use in the calc. Checking if this is a bug

Steps to reproduce

NA

Version

main

[calebrob6]

I verified the behavior of JaccardLoss here https://gist.github.com/calebrob6/658edaa59c68f0c0a510f8d9d7a41458.

Currently our code will always use the average jaccard over all classes regardless of what self.hparams["num_classes"] is.

We definitely should fix this.

@isaaccorley ignore_index support might be a cool feature for the torchseg fork.

[robmarkcole]

I suggest this requires a new arg of classes_in_loss which is a List[int]. This should pop off the ignore_index

[calebrob6]

Have you run into scenarios in which you would want to ignore multiple classes? If not, I think we can just pop ignore_index from the entire class list if it is given.

[robmarkcole]

I've not really used JaccardLoss, and commonly use nn.CrossEntropyLoss where the weights parameter has been very effective for imbalanced datasets. If I understand the use of classes in JaccardLoss it has a simiar purpose, but is binary on/off per class, rather than a weighting

Coming back to your question, for multiclass datasets where I might only care about 1 or 2 classes, I can imagine ignoring multiple classes

[calebrob6]

+1 for fiddling with weights in CrossEntropyLoss.

Yep, I can imagine a few scenarios too (e.g. dropping loss over noisy classes). I would usually do this by simply remapping all classes I didn't care about to a nodata mask value, then setting ignore_index to that value, but this is more elegant. My concern now would be implementing this functionality for the other loss functions.

[isaaccorley]

TorchSeg now supports both classes and ignore_index in DiceLoss, JaccardLoss, and FocalLoss.

[robmarkcole]

Thanks @isaaccorley

For anyone looking, they are on branch https://github.com/isaaccorley/torchseg/blob/losses/ignore-index/torchseg/losses/

[isaaccorley]

@robmarkcole it's merged into main now. I'll make another pre-release on pypi today.