Correct behavior of ignore_index for JaccardLoss

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:

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:

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(

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