Repository: microsoft/torchgeo

Issue:
Add `ignore_index` support for Jaccard Loss
### Summary

Currently, the `SemanticSegmentationTask` recognises the `ignore_index` parameter only when cross entropy and focal loss is used. However, `smp.losses.JaccardLoss` implicitly supports this option via its `classes` parameter, which is currently set to `self.hparams["num_classes"]`. The FR is to adapt what is passed as argument to this parameter so that `ignore_index` can work with all currently supported losses.

### Rationale

The Jaccard index is a common semantic segmentation metric which also makes for a decent loss function. Because of the way it is defined, it is important to ignore overly dominant classes (e.g., the background when classifying building rooftops); otherwise performance can be hindered significantly.

### Implementation

Change the `classes` argument of `smp.losses.JaccardLoss` in `SemanticSegmentationTask.configure_losses` from `self.hparams["num_classes"]` to `list(set(list(range(self.hparams["num_classes"]))).difference(set([ignore_index])))`, assuming that `ignore_index` is not `None`.

### Alternatives

_No response_

### Additional information

_No response_


Discussion:

Resolving patch:
diff --git a/torchgeo/trainers/segmentation.py b/torchgeo/trainers/segmentation.py
--- a/torchgeo/trainers/segmentation.py
+++ b/torchgeo/trainers/segmentation.py
@@ -4,7 +4,6 @@
 """Trainers for semantic segmentation."""
 
 import os
-import warnings
 from typing import Any, Optional, Union
 
 import matplotlib.pyplot as plt
@@ -70,9 +69,6 @@ class and used with 'ce' loss.
             freeze_decoder: Freeze the decoder network to linear probe
                 the segmentation head.
 
-        Warns:
-            UserWarning: When loss='jaccard' and ignore_index is specified.
-
         .. versionchanged:: 0.3
            *ignore_zeros* was renamed to *ignore_index*.
 
@@ -87,13 +83,10 @@ class and used with 'ce' loss.
            The *weights* parameter now supports WeightEnums and checkpoint paths.
            *learning_rate* and *learning_rate_schedule_patience* were renamed to
            *lr* and *patience*.
-        """
-        if ignore_index is not None and loss == "jaccard":
-            warnings.warn(
-                "ignore_index has no effect on training when loss='jaccard'",
-                UserWarning,
-            )
 
+        .. versionchanged:: 0.6
+            The *ignore_index* parameter now works for jaccard loss.
+        """
         self.weights = weights
         super().__init__(ignore="weights")
 
@@ -111,9 +104,13 @@ def configure_losses(self) -> None:
                 ignore_index=ignore_value, weight=self.hparams["class_weights"]
             )
         elif loss == "jaccard":
-            self.criterion = smp.losses.JaccardLoss(
-                mode="multiclass", classes=self.hparams["num_classes"]
-            )
+            # JaccardLoss requires a list of classes to use instead of a class
+            # index to ignore.
+            classes = [
+                i for i in range(self.hparams["num_classes"]) if i != ignore_index
+            ]
+
+            self.criterion = smp.losses.JaccardLoss(mode="multiclass", classes=classes)
         elif loss == "focal":
             self.criterion = smp.losses.FocalLoss(
                 "multiclass", ignore_index=ignore_index, normalized=True


Test patch:
diff --git a/tests/conf/chesapeake_cvpr_5.yaml b/tests/conf/chesapeake_cvpr_5.yaml
--- a/tests/conf/chesapeake_cvpr_5.yaml
+++ b/tests/conf/chesapeake_cvpr_5.yaml
@@ -1,13 +1,13 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: "ce"
+    loss: "jaccard"
     model: "unet"
     backbone: "resnet50"
     in_channels: 4
     num_classes: 5
     num_filters: 1
-    ignore_index: null
+    ignore_index: 0
 data:
   class_path: ChesapeakeCVPRDataModule
   init_args:
diff --git a/tests/conf/chesapeake_cvpr_7.yaml b/tests/conf/chesapeake_cvpr_7.yaml
--- a/tests/conf/chesapeake_cvpr_7.yaml
+++ b/tests/conf/chesapeake_cvpr_7.yaml
@@ -7,7 +7,7 @@ model:
     in_channels: 4
     num_classes: 7
     num_filters: 1
-    ignore_index: null
+    ignore_index: 0
 data:
   class_path: ChesapeakeCVPRDataModule
   init_args:
diff --git a/tests/trainers/test_segmentation.py b/tests/trainers/test_segmentation.py
--- a/tests/trainers/test_segmentation.py
+++ b/tests/trainers/test_segmentation.py
@@ -179,11 +179,6 @@ def test_invalid_loss(self) -> None:
         with pytest.raises(ValueError, match=match):
             SemanticSegmentationTask(loss="invalid_loss")
 
-    def test_ignoreindex_with_jaccard(self) -> None:
-        match = "ignore_index has no effect on training when loss='jaccard'"
-        with pytest.warns(UserWarning, match=match):
-            SemanticSegmentationTask(loss="jaccard", ignore_index=0)
-
     def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
         monkeypatch.setattr(SEN12MSDataModule, "plot", plot)
         datamodule = SEN12MSDataModule(