We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchgeo/trainers/segmentation.py`
Content:
```
1 # Copyright (c) Microsoft Corporation. All rights reserved.
2 # Licensed under the MIT License.
3 
4 """Trainers for semantic segmentation."""
5 
6 import os
7 import warnings
8 from typing import Any, Optional, Union
9 
10 import matplotlib.pyplot as plt
11 import segmentation_models_pytorch as smp
12 import torch.nn as nn
13 from matplotlib.figure import Figure
14 from torch import Tensor
15 from torchmetrics import MetricCollection
16 from torchmetrics.classification import MulticlassAccuracy, MulticlassJaccardIndex
17 from torchvision.models._api import WeightsEnum
18 
19 from ..datasets import RGBBandsMissingError, unbind_samples
20 from ..models import FCN, get_weight
21 from . import utils
22 from .base import BaseTask
23 
24 
25 class SemanticSegmentationTask(BaseTask):
26     """Semantic Segmentation."""
27 
28     def __init__(
29         self,
30         model: str = "unet",
31         backbone: str = "resnet50",
32         weights: Optional[Union[WeightsEnum, str, bool]] = None,
33         in_channels: int = 3,
34         num_classes: int = 1000,
35         num_filters: int = 3,
36         loss: str = "ce",
37         class_weights: Optional[Tensor] = None,
38         ignore_index: Optional[int] = None,
39         lr: float = 1e-3,
40         patience: int = 10,
41         freeze_backbone: bool = False,
42         freeze_decoder: bool = False,
43     ) -> None:
44         """Inititalize a new SemanticSegmentationTask instance.
45 
46         Args:
47             model: Name of the
48                 `smp <https://smp.readthedocs.io/en/latest/models.html>`__ model to use.
49             backbone: Name of the `timm
50                 <https://smp.readthedocs.io/en/latest/encoders_timm.html>`__ or `smp
51                 <https://smp.readthedocs.io/en/latest/encoders.html>`__ backbone to use.
52             weights: Initial model weights. Either a weight enum, the string
53                 representation of a weight enum, True for ImageNet weights, False or
54                 None for random weights, or the path to a saved model state dict. FCN
55                 model does not support pretrained weights. Pretrained ViT weight enums
56                 are not supported yet.
57             in_channels: Number of input channels to model.
58             num_classes: Number of prediction classes.
59             num_filters: Number of filters. Only applicable when model='fcn'.
60             loss: Name of the loss function, currently supports
61                 'ce', 'jaccard' or 'focal' loss.
62             class_weights: Optional rescaling weight given to each
63                 class and used with 'ce' loss.
64             ignore_index: Optional integer class index to ignore in the loss and
65                 metrics.
66             lr: Learning rate for optimizer.
67             patience: Patience for learning rate scheduler.
68             freeze_backbone: Freeze the backbone network to fine-tune the
69                 decoder and segmentation head.
70             freeze_decoder: Freeze the decoder network to linear probe
71                 the segmentation head.
72 
73         Warns:
74             UserWarning: When loss='jaccard' and ignore_index is specified.
75 
76         .. versionchanged:: 0.3
77            *ignore_zeros* was renamed to *ignore_index*.
78 
79         .. versionchanged:: 0.4
80            *segmentation_model*, *encoder_name*, and *encoder_weights*
81            were renamed to *model*, *backbone*, and *weights*.
82 
83         .. versionadded: 0.5
84             The *class_weights*, *freeze_backbone*, and *freeze_decoder* parameters.
85 
86         .. versionchanged:: 0.5
87            The *weights* parameter now supports WeightEnums and checkpoint paths.
88            *learning_rate* and *learning_rate_schedule_patience* were renamed to
89            *lr* and *patience*.
90         """
91         if ignore_index is not None and loss == "jaccard":
92             warnings.warn(
93                 "ignore_index has no effect on training when loss='jaccard'",
94                 UserWarning,
95             )
96 
97         self.weights = weights
98         super().__init__(ignore="weights")
99 
100     def configure_losses(self) -> None:
101         """Initialize the loss criterion.
102 
103         Raises:
104             ValueError: If *loss* is invalid.
105         """
106         loss: str = self.hparams["loss"]
107         ignore_index = self.hparams["ignore_index"]
108         if loss == "ce":
109             ignore_value = -1000 if ignore_index is None else ignore_index
110             self.criterion = nn.CrossEntropyLoss(
111                 ignore_index=ignore_value, weight=self.hparams["class_weights"]
112             )
113         elif loss == "jaccard":
114             self.criterion = smp.losses.JaccardLoss(
115                 mode="multiclass", classes=self.hparams["num_classes"]
116             )
117         elif loss == "focal":
118             self.criterion = smp.losses.FocalLoss(
119                 "multiclass", ignore_index=ignore_index, normalized=True
120             )
121         else:
122             raise ValueError(
123                 f"Loss type '{loss}' is not valid. "
124                 "Currently, supports 'ce', 'jaccard' or 'focal' loss."
125             )
126 
127     def configure_metrics(self) -> None:
128         """Initialize the performance metrics."""
129         num_classes: int = self.hparams["num_classes"]
130         ignore_index: Optional[int] = self.hparams["ignore_index"]
131         metrics = MetricCollection(
132             [
133                 MulticlassAccuracy(
134                     num_classes=num_classes,
135                     ignore_index=ignore_index,
136                     multidim_average="global",
137                     average="micro",
138                 ),
139                 MulticlassJaccardIndex(
140                     num_classes=num_classes, ignore_index=ignore_index, average="micro"
141                 ),
142             ]
143         )
144         self.train_metrics = metrics.clone(prefix="train_")
145         self.val_metrics = metrics.clone(prefix="val_")
146         self.test_metrics = metrics.clone(prefix="test_")
147 
148     def configure_models(self) -> None:
149         """Initialize the model.
150 
151         Raises:
152             ValueError: If *model* is invalid.
153         """
154         model: str = self.hparams["model"]
155         backbone: str = self.hparams["backbone"]
156         weights = self.weights
157         in_channels: int = self.hparams["in_channels"]
158         num_classes: int = self.hparams["num_classes"]
159         num_filters: int = self.hparams["num_filters"]
160 
161         if model == "unet":
162             self.model = smp.Unet(
163                 encoder_name=backbone,
164                 encoder_weights="imagenet" if weights is True else None,
165                 in_channels=in_channels,
166                 classes=num_classes,
167             )
168         elif model == "deeplabv3+":
169             self.model = smp.DeepLabV3Plus(
170                 encoder_name=backbone,
171                 encoder_weights="imagenet" if weights is True else None,
172                 in_channels=in_channels,
173                 classes=num_classes,
174             )
175         elif model == "fcn":
176             self.model = FCN(
177                 in_channels=in_channels, classes=num_classes, num_filters=num_filters
178             )
179         else:
180             raise ValueError(
181                 f"Model type '{model}' is not valid. "
182                 "Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
183             )
184 
185         if model != "fcn":
186             if weights and weights is not True:
187                 if isinstance(weights, WeightsEnum):
188                     state_dict = weights.get_state_dict(progress=True)
189                 elif os.path.exists(weights):
190                     _, state_dict = utils.extract_backbone(weights)
191                 else:
192                     state_dict = get_weight(weights).get_state_dict(progress=True)
193                 self.model.encoder.load_state_dict(state_dict)
194 
195         # Freeze backbone
196         if self.hparams["freeze_backbone"] and model in ["unet", "deeplabv3+"]:
197             for param in self.model.encoder.parameters():
198                 param.requires_grad = False
199 
200         # Freeze decoder
201         if self.hparams["freeze_decoder"] and model in ["unet", "deeplabv3+"]:
202             for param in self.model.decoder.parameters():
203                 param.requires_grad = False
204 
205     def training_step(
206         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
207     ) -> Tensor:
208         """Compute the training loss and additional metrics.
209 
210         Args:
211             batch: The output of your DataLoader.
212             batch_idx: Integer displaying index of this batch.
213             dataloader_idx: Index of the current dataloader.
214 
215         Returns:
216             The loss tensor.
217         """
218         x = batch["image"]
219         y = batch["mask"]
220         y_hat = self(x)
221         loss: Tensor = self.criterion(y_hat, y)
222         self.log("train_loss", loss)
223         self.train_metrics(y_hat, y)
224         self.log_dict(self.train_metrics)
225         return loss
226 
227     def validation_step(
228         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
229     ) -> None:
230         """Compute the validation loss and additional metrics.
231 
232         Args:
233             batch: The output of your DataLoader.
234             batch_idx: Integer displaying index of this batch.
235             dataloader_idx: Index of the current dataloader.
236         """
237         x = batch["image"]
238         y = batch["mask"]
239         y_hat = self(x)
240         loss = self.criterion(y_hat, y)
241         self.log("val_loss", loss)
242         self.val_metrics(y_hat, y)
243         self.log_dict(self.val_metrics)
244 
245         if (
246             batch_idx < 10
247             and hasattr(self.trainer, "datamodule")
248             and hasattr(self.trainer.datamodule, "plot")
249             and self.logger
250             and hasattr(self.logger, "experiment")
251             and hasattr(self.logger.experiment, "add_figure")
252         ):
253             datamodule = self.trainer.datamodule
254             batch["prediction"] = y_hat.argmax(dim=1)
255             for key in ["image", "mask", "prediction"]:
256                 batch[key] = batch[key].cpu()
257             sample = unbind_samples(batch)[0]
258 
259             fig: Optional[Figure] = None
260             try:
261                 fig = datamodule.plot(sample)
262             except RGBBandsMissingError:
263                 pass
264 
265             if fig:
266                 summary_writer = self.logger.experiment
267                 summary_writer.add_figure(
268                     f"image/{batch_idx}", fig, global_step=self.global_step
269                 )
270                 plt.close()
271 
272     def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
273         """Compute the test loss and additional metrics.
274 
275         Args:
276             batch: The output of your DataLoader.
277             batch_idx: Integer displaying index of this batch.
278             dataloader_idx: Index of the current dataloader.
279         """
280         x = batch["image"]
281         y = batch["mask"]
282         y_hat = self(x)
283         loss = self.criterion(y_hat, y)
284         self.log("test_loss", loss)
285         self.test_metrics(y_hat, y)
286         self.log_dict(self.test_metrics)
287 
288     def predict_step(
289         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
290     ) -> Tensor:
291         """Compute the predicted class probabilities.
292 
293         Args:
294             batch: The output of your DataLoader.
295             batch_idx: Integer displaying index of this batch.
296             dataloader_idx: Index of the current dataloader.
297 
298         Returns:
299             Output predicted probabilities.
300         """
301         x = batch["image"]
302         y_hat: Tensor = self(x).softmax(dim=1)
303         return y_hat
304
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
             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 @@
            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 @@
                 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

{"golden_diff": "diff --git a/torchgeo/trainers/segmentation.py b/torchgeo/trainers/segmentation.py\n--- a/torchgeo/trainers/segmentation.py\n+++ b/torchgeo/trainers/segmentation.py\n@@ -4,7 +4,6 @@\n \"\"\"Trainers for semantic segmentation.\"\"\"\n \n import os\n-import warnings\n from typing import Any, Optional, Union\n \n import matplotlib.pyplot as plt\n@@ -70,9 +69,6 @@\n             freeze_decoder: Freeze the decoder network to linear probe\n                 the segmentation head.\n \n-        Warns:\n-            UserWarning: When loss='jaccard' and ignore_index is specified.\n-\n         .. versionchanged:: 0.3\n            *ignore_zeros* was renamed to *ignore_index*.\n \n@@ -87,13 +83,10 @@\n            The *weights* parameter now supports WeightEnums and checkpoint paths.\n            *learning_rate* and *learning_rate_schedule_patience* were renamed to\n            *lr* and *patience*.\n-        \"\"\"\n-        if ignore_index is not None and loss == \"jaccard\":\n-            warnings.warn(\n-                \"ignore_index has no effect on training when loss='jaccard'\",\n-                UserWarning,\n-            )\n \n+        .. versionchanged:: 0.6\n+            The *ignore_index* parameter now works for jaccard loss.\n+        \"\"\"\n         self.weights = weights\n         super().__init__(ignore=\"weights\")\n \n@@ -111,9 +104,13 @@\n                 ignore_index=ignore_value, weight=self.hparams[\"class_weights\"]\n             )\n         elif loss == \"jaccard\":\n-            self.criterion = smp.losses.JaccardLoss(\n-                mode=\"multiclass\", classes=self.hparams[\"num_classes\"]\n-            )\n+            # JaccardLoss requires a list of classes to use instead of a class\n+            # index to ignore.\n+            classes = [\n+                i for i in range(self.hparams[\"num_classes\"]) if i != ignore_index\n+            ]\n+\n+            self.criterion = smp.losses.JaccardLoss(mode=\"multiclass\", classes=classes)\n         elif loss == \"focal\":\n             self.criterion = smp.losses.FocalLoss(\n                 \"multiclass\", ignore_index=ignore_index, normalized=True\n", "issue": "Add `ignore_index` support for Jaccard Loss\n### Summary\r\n\r\nCurrently, 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.\r\n\r\n### Rationale\r\n\r\nThe 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.\r\n\r\n### Implementation\r\n\r\nChange 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`.\r\n\r\n### Alternatives\r\n\r\n_No response_\r\n\r\n### Additional information\r\n\r\n_No response_\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"Trainers for semantic segmentation.\"\"\"\n\nimport os\nimport warnings\nfrom typing import Any, Optional, Union\n\nimport matplotlib.pyplot as plt\nimport segmentation_models_pytorch as smp\nimport torch.nn as nn\nfrom matplotlib.figure import Figure\nfrom torch import Tensor\nfrom torchmetrics import MetricCollection\nfrom torchmetrics.classification import MulticlassAccuracy, MulticlassJaccardIndex\nfrom torchvision.models._api import WeightsEnum\n\nfrom ..datasets import RGBBandsMissingError, unbind_samples\nfrom ..models import FCN, get_weight\nfrom . import utils\nfrom .base import BaseTask\n\n\nclass SemanticSegmentationTask(BaseTask):\n    \"\"\"Semantic Segmentation.\"\"\"\n\n    def __init__(\n        self,\n        model: str = \"unet\",\n        backbone: str = \"resnet50\",\n        weights: Optional[Union[WeightsEnum, str, bool]] = None,\n        in_channels: int = 3,\n        num_classes: int = 1000,\n        num_filters: int = 3,\n        loss: str = \"ce\",\n        class_weights: Optional[Tensor] = None,\n        ignore_index: Optional[int] = None,\n        lr: float = 1e-3,\n        patience: int = 10,\n        freeze_backbone: bool = False,\n        freeze_decoder: bool = False,\n    ) -> None:\n        \"\"\"Inititalize a new SemanticSegmentationTask instance.\n\n        Args:\n            model: Name of the\n                `smp <https://smp.readthedocs.io/en/latest/models.html>`__ model to use.\n            backbone: Name of the `timm\n                <https://smp.readthedocs.io/en/latest/encoders_timm.html>`__ or `smp\n                <https://smp.readthedocs.io/en/latest/encoders.html>`__ backbone to use.\n            weights: Initial model weights. Either a weight enum, the string\n                representation of a weight enum, True for ImageNet weights, False or\n                None for random weights, or the path to a saved model state dict. FCN\n                model does not support pretrained weights. Pretrained ViT weight enums\n                are not supported yet.\n            in_channels: Number of input channels to model.\n            num_classes: Number of prediction classes.\n            num_filters: Number of filters. Only applicable when model='fcn'.\n            loss: Name of the loss function, currently supports\n                'ce', 'jaccard' or 'focal' loss.\n            class_weights: Optional rescaling weight given to each\n                class and used with 'ce' loss.\n            ignore_index: Optional integer class index to ignore in the loss and\n                metrics.\n            lr: Learning rate for optimizer.\n            patience: Patience for learning rate scheduler.\n            freeze_backbone: Freeze the backbone network to fine-tune the\n                decoder and segmentation head.\n            freeze_decoder: Freeze the decoder network to linear probe\n                the segmentation head.\n\n        Warns:\n            UserWarning: When loss='jaccard' and ignore_index is specified.\n\n        .. versionchanged:: 0.3\n           *ignore_zeros* was renamed to *ignore_index*.\n\n        .. versionchanged:: 0.4\n           *segmentation_model*, *encoder_name*, and *encoder_weights*\n           were renamed to *model*, *backbone*, and *weights*.\n\n        .. versionadded: 0.5\n            The *class_weights*, *freeze_backbone*, and *freeze_decoder* parameters.\n\n        .. versionchanged:: 0.5\n           The *weights* parameter now supports WeightEnums and checkpoint paths.\n           *learning_rate* and *learning_rate_schedule_patience* were renamed to\n           *lr* and *patience*.\n        \"\"\"\n        if ignore_index is not None and loss == \"jaccard\":\n            warnings.warn(\n                \"ignore_index has no effect on training when loss='jaccard'\",\n                UserWarning,\n            )\n\n        self.weights = weights\n        super().__init__(ignore=\"weights\")\n\n    def configure_losses(self) -> None:\n        \"\"\"Initialize the loss criterion.\n\n        Raises:\n            ValueError: If *loss* is invalid.\n        \"\"\"\n        loss: str = self.hparams[\"loss\"]\n        ignore_index = self.hparams[\"ignore_index\"]\n        if loss == \"ce\":\n            ignore_value = -1000 if ignore_index is None else ignore_index\n            self.criterion = nn.CrossEntropyLoss(\n                ignore_index=ignore_value, weight=self.hparams[\"class_weights\"]\n            )\n        elif loss == \"jaccard\":\n            self.criterion = smp.losses.JaccardLoss(\n                mode=\"multiclass\", classes=self.hparams[\"num_classes\"]\n            )\n        elif loss == \"focal\":\n            self.criterion = smp.losses.FocalLoss(\n                \"multiclass\", ignore_index=ignore_index, normalized=True\n            )\n        else:\n            raise ValueError(\n                f\"Loss type '{loss}' is not valid. \"\n                \"Currently, supports 'ce', 'jaccard' or 'focal' loss.\"\n            )\n\n    def configure_metrics(self) -> None:\n        \"\"\"Initialize the performance metrics.\"\"\"\n        num_classes: int = self.hparams[\"num_classes\"]\n        ignore_index: Optional[int] = self.hparams[\"ignore_index\"]\n        metrics = MetricCollection(\n            [\n                MulticlassAccuracy(\n                    num_classes=num_classes,\n                    ignore_index=ignore_index,\n                    multidim_average=\"global\",\n                    average=\"micro\",\n                ),\n                MulticlassJaccardIndex(\n                    num_classes=num_classes, ignore_index=ignore_index, average=\"micro\"\n                ),\n            ]\n        )\n        self.train_metrics = metrics.clone(prefix=\"train_\")\n        self.val_metrics = metrics.clone(prefix=\"val_\")\n        self.test_metrics = metrics.clone(prefix=\"test_\")\n\n    def configure_models(self) -> None:\n        \"\"\"Initialize the model.\n\n        Raises:\n            ValueError: If *model* is invalid.\n        \"\"\"\n        model: str = self.hparams[\"model\"]\n        backbone: str = self.hparams[\"backbone\"]\n        weights = self.weights\n        in_channels: int = self.hparams[\"in_channels\"]\n        num_classes: int = self.hparams[\"num_classes\"]\n        num_filters: int = self.hparams[\"num_filters\"]\n\n        if model == \"unet\":\n            self.model = smp.Unet(\n                encoder_name=backbone,\n                encoder_weights=\"imagenet\" if weights is True else None,\n                in_channels=in_channels,\n                classes=num_classes,\n            )\n        elif model == \"deeplabv3+\":\n            self.model = smp.DeepLabV3Plus(\n                encoder_name=backbone,\n                encoder_weights=\"imagenet\" if weights is True else None,\n                in_channels=in_channels,\n                classes=num_classes,\n            )\n        elif model == \"fcn\":\n            self.model = FCN(\n                in_channels=in_channels, classes=num_classes, num_filters=num_filters\n            )\n        else:\n            raise ValueError(\n                f\"Model type '{model}' is not valid. \"\n                \"Currently, only supports 'unet', 'deeplabv3+' and 'fcn'.\"\n            )\n\n        if model != \"fcn\":\n            if weights and weights is not True:\n                if isinstance(weights, WeightsEnum):\n                    state_dict = weights.get_state_dict(progress=True)\n                elif os.path.exists(weights):\n                    _, state_dict = utils.extract_backbone(weights)\n                else:\n                    state_dict = get_weight(weights).get_state_dict(progress=True)\n                self.model.encoder.load_state_dict(state_dict)\n\n        # Freeze backbone\n        if self.hparams[\"freeze_backbone\"] and model in [\"unet\", \"deeplabv3+\"]:\n            for param in self.model.encoder.parameters():\n                param.requires_grad = False\n\n        # Freeze decoder\n        if self.hparams[\"freeze_decoder\"] and model in [\"unet\", \"deeplabv3+\"]:\n            for param in self.model.decoder.parameters():\n                param.requires_grad = False\n\n    def training_step(\n        self, batch: Any, batch_idx: int, dataloader_idx: int = 0\n    ) -> Tensor:\n        \"\"\"Compute the training loss and additional metrics.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n\n        Returns:\n            The loss tensor.\n        \"\"\"\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self(x)\n        loss: Tensor = self.criterion(y_hat, y)\n        self.log(\"train_loss\", loss)\n        self.train_metrics(y_hat, y)\n        self.log_dict(self.train_metrics)\n        return loss\n\n    def validation_step(\n        self, batch: Any, batch_idx: int, dataloader_idx: int = 0\n    ) -> None:\n        \"\"\"Compute the validation loss and additional metrics.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n        \"\"\"\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self(x)\n        loss = self.criterion(y_hat, y)\n        self.log(\"val_loss\", loss)\n        self.val_metrics(y_hat, y)\n        self.log_dict(self.val_metrics)\n\n        if (\n            batch_idx < 10\n            and hasattr(self.trainer, \"datamodule\")\n            and hasattr(self.trainer.datamodule, \"plot\")\n            and self.logger\n            and hasattr(self.logger, \"experiment\")\n            and hasattr(self.logger.experiment, \"add_figure\")\n        ):\n            datamodule = self.trainer.datamodule\n            batch[\"prediction\"] = y_hat.argmax(dim=1)\n            for key in [\"image\", \"mask\", \"prediction\"]:\n                batch[key] = batch[key].cpu()\n            sample = unbind_samples(batch)[0]\n\n            fig: Optional[Figure] = None\n            try:\n                fig = datamodule.plot(sample)\n            except RGBBandsMissingError:\n                pass\n\n            if fig:\n                summary_writer = self.logger.experiment\n                summary_writer.add_figure(\n                    f\"image/{batch_idx}\", fig, global_step=self.global_step\n                )\n                plt.close()\n\n    def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:\n        \"\"\"Compute the test loss and additional metrics.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n        \"\"\"\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self(x)\n        loss = self.criterion(y_hat, y)\n        self.log(\"test_loss\", loss)\n        self.test_metrics(y_hat, y)\n        self.log_dict(self.test_metrics)\n\n    def predict_step(\n        self, batch: Any, batch_idx: int, dataloader_idx: int = 0\n    ) -> Tensor:\n        \"\"\"Compute the predicted class probabilities.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n\n        Returns:\n            Output predicted probabilities.\n        \"\"\"\n        x = batch[\"image\"]\n        y_hat: Tensor = self(x).softmax(dim=1)\n        return y_hat\n", "path": "torchgeo/trainers/segmentation.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"Trainers for semantic segmentation.\"\"\"\n\nimport os\nfrom typing import Any, Optional, Union\n\nimport matplotlib.pyplot as plt\nimport segmentation_models_pytorch as smp\nimport torch.nn as nn\nfrom matplotlib.figure import Figure\nfrom torch import Tensor\nfrom torchmetrics import MetricCollection\nfrom torchmetrics.classification import MulticlassAccuracy, MulticlassJaccardIndex\nfrom torchvision.models._api import WeightsEnum\n\nfrom ..datasets import RGBBandsMissingError, unbind_samples\nfrom ..models import FCN, get_weight\nfrom . import utils\nfrom .base import BaseTask\n\n\nclass SemanticSegmentationTask(BaseTask):\n    \"\"\"Semantic Segmentation.\"\"\"\n\n    def __init__(\n        self,\n        model: str = \"unet\",\n        backbone: str = \"resnet50\",\n        weights: Optional[Union[WeightsEnum, str, bool]] = None,\n        in_channels: int = 3,\n        num_classes: int = 1000,\n        num_filters: int = 3,\n        loss: str = \"ce\",\n        class_weights: Optional[Tensor] = None,\n        ignore_index: Optional[int] = None,\n        lr: float = 1e-3,\n        patience: int = 10,\n        freeze_backbone: bool = False,\n        freeze_decoder: bool = False,\n    ) -> None:\n        \"\"\"Inititalize a new SemanticSegmentationTask instance.\n\n        Args:\n            model: Name of the\n                `smp <https://smp.readthedocs.io/en/latest/models.html>`__ model to use.\n            backbone: Name of the `timm\n                <https://smp.readthedocs.io/en/latest/encoders_timm.html>`__ or `smp\n                <https://smp.readthedocs.io/en/latest/encoders.html>`__ backbone to use.\n            weights: Initial model weights. Either a weight enum, the string\n                representation of a weight enum, True for ImageNet weights, False or\n                None for random weights, or the path to a saved model state dict. FCN\n                model does not support pretrained weights. Pretrained ViT weight enums\n                are not supported yet.\n            in_channels: Number of input channels to model.\n            num_classes: Number of prediction classes.\n            num_filters: Number of filters. Only applicable when model='fcn'.\n            loss: Name of the loss function, currently supports\n                'ce', 'jaccard' or 'focal' loss.\n            class_weights: Optional rescaling weight given to each\n                class and used with 'ce' loss.\n            ignore_index: Optional integer class index to ignore in the loss and\n                metrics.\n            lr: Learning rate for optimizer.\n            patience: Patience for learning rate scheduler.\n            freeze_backbone: Freeze the backbone network to fine-tune the\n                decoder and segmentation head.\n            freeze_decoder: Freeze the decoder network to linear probe\n                the segmentation head.\n\n        .. versionchanged:: 0.3\n           *ignore_zeros* was renamed to *ignore_index*.\n\n        .. versionchanged:: 0.4\n           *segmentation_model*, *encoder_name*, and *encoder_weights*\n           were renamed to *model*, *backbone*, and *weights*.\n\n        .. versionadded: 0.5\n            The *class_weights*, *freeze_backbone*, and *freeze_decoder* parameters.\n\n        .. versionchanged:: 0.5\n           The *weights* parameter now supports WeightEnums and checkpoint paths.\n           *learning_rate* and *learning_rate_schedule_patience* were renamed to\n           *lr* and *patience*.\n\n        .. versionchanged:: 0.6\n            The *ignore_index* parameter now works for jaccard loss.\n        \"\"\"\n        self.weights = weights\n        super().__init__(ignore=\"weights\")\n\n    def configure_losses(self) -> None:\n        \"\"\"Initialize the loss criterion.\n\n        Raises:\n            ValueError: If *loss* is invalid.\n        \"\"\"\n        loss: str = self.hparams[\"loss\"]\n        ignore_index = self.hparams[\"ignore_index\"]\n        if loss == \"ce\":\n            ignore_value = -1000 if ignore_index is None else ignore_index\n            self.criterion = nn.CrossEntropyLoss(\n                ignore_index=ignore_value, weight=self.hparams[\"class_weights\"]\n            )\n        elif loss == \"jaccard\":\n            # JaccardLoss requires a list of classes to use instead of a class\n            # index to ignore.\n            classes = [\n                i for i in range(self.hparams[\"num_classes\"]) if i != ignore_index\n            ]\n\n            self.criterion = smp.losses.JaccardLoss(mode=\"multiclass\", classes=classes)\n        elif loss == \"focal\":\n            self.criterion = smp.losses.FocalLoss(\n                \"multiclass\", ignore_index=ignore_index, normalized=True\n            )\n        else:\n            raise ValueError(\n                f\"Loss type '{loss}' is not valid. \"\n                \"Currently, supports 'ce', 'jaccard' or 'focal' loss.\"\n            )\n\n    def configure_metrics(self) -> None:\n        \"\"\"Initialize the performance metrics.\"\"\"\n        num_classes: int = self.hparams[\"num_classes\"]\n        ignore_index: Optional[int] = self.hparams[\"ignore_index\"]\n        metrics = MetricCollection(\n            [\n                MulticlassAccuracy(\n                    num_classes=num_classes,\n                    ignore_index=ignore_index,\n                    multidim_average=\"global\",\n                    average=\"micro\",\n                ),\n                MulticlassJaccardIndex(\n                    num_classes=num_classes, ignore_index=ignore_index, average=\"micro\"\n                ),\n            ]\n        )\n        self.train_metrics = metrics.clone(prefix=\"train_\")\n        self.val_metrics = metrics.clone(prefix=\"val_\")\n        self.test_metrics = metrics.clone(prefix=\"test_\")\n\n    def configure_models(self) -> None:\n        \"\"\"Initialize the model.\n\n        Raises:\n            ValueError: If *model* is invalid.\n        \"\"\"\n        model: str = self.hparams[\"model\"]\n        backbone: str = self.hparams[\"backbone\"]\n        weights = self.weights\n        in_channels: int = self.hparams[\"in_channels\"]\n        num_classes: int = self.hparams[\"num_classes\"]\n        num_filters: int = self.hparams[\"num_filters\"]\n\n        if model == \"unet\":\n            self.model = smp.Unet(\n                encoder_name=backbone,\n                encoder_weights=\"imagenet\" if weights is True else None,\n                in_channels=in_channels,\n                classes=num_classes,\n            )\n        elif model == \"deeplabv3+\":\n            self.model = smp.DeepLabV3Plus(\n                encoder_name=backbone,\n                encoder_weights=\"imagenet\" if weights is True else None,\n                in_channels=in_channels,\n                classes=num_classes,\n            )\n        elif model == \"fcn\":\n            self.model = FCN(\n                in_channels=in_channels, classes=num_classes, num_filters=num_filters\n            )\n        else:\n            raise ValueError(\n                f\"Model type '{model}' is not valid. \"\n                \"Currently, only supports 'unet', 'deeplabv3+' and 'fcn'.\"\n            )\n\n        if model != \"fcn\":\n            if weights and weights is not True:\n                if isinstance(weights, WeightsEnum):\n                    state_dict = weights.get_state_dict(progress=True)\n                elif os.path.exists(weights):\n                    _, state_dict = utils.extract_backbone(weights)\n                else:\n                    state_dict = get_weight(weights).get_state_dict(progress=True)\n                self.model.encoder.load_state_dict(state_dict)\n\n        # Freeze backbone\n        if self.hparams[\"freeze_backbone\"] and model in [\"unet\", \"deeplabv3+\"]:\n            for param in self.model.encoder.parameters():\n                param.requires_grad = False\n\n        # Freeze decoder\n        if self.hparams[\"freeze_decoder\"] and model in [\"unet\", \"deeplabv3+\"]:\n            for param in self.model.decoder.parameters():\n                param.requires_grad = False\n\n    def training_step(\n        self, batch: Any, batch_idx: int, dataloader_idx: int = 0\n    ) -> Tensor:\n        \"\"\"Compute the training loss and additional metrics.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n\n        Returns:\n            The loss tensor.\n        \"\"\"\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self(x)\n        loss: Tensor = self.criterion(y_hat, y)\n        self.log(\"train_loss\", loss)\n        self.train_metrics(y_hat, y)\n        self.log_dict(self.train_metrics)\n        return loss\n\n    def validation_step(\n        self, batch: Any, batch_idx: int, dataloader_idx: int = 0\n    ) -> None:\n        \"\"\"Compute the validation loss and additional metrics.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n        \"\"\"\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self(x)\n        loss = self.criterion(y_hat, y)\n        self.log(\"val_loss\", loss)\n        self.val_metrics(y_hat, y)\n        self.log_dict(self.val_metrics)\n\n        if (\n            batch_idx < 10\n            and hasattr(self.trainer, \"datamodule\")\n            and hasattr(self.trainer.datamodule, \"plot\")\n            and self.logger\n            and hasattr(self.logger, \"experiment\")\n            and hasattr(self.logger.experiment, \"add_figure\")\n        ):\n            datamodule = self.trainer.datamodule\n            batch[\"prediction\"] = y_hat.argmax(dim=1)\n            for key in [\"image\", \"mask\", \"prediction\"]:\n                batch[key] = batch[key].cpu()\n            sample = unbind_samples(batch)[0]\n\n            fig: Optional[Figure] = None\n            try:\n                fig = datamodule.plot(sample)\n            except RGBBandsMissingError:\n                pass\n\n            if fig:\n                summary_writer = self.logger.experiment\n                summary_writer.add_figure(\n                    f\"image/{batch_idx}\", fig, global_step=self.global_step\n                )\n                plt.close()\n\n    def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:\n        \"\"\"Compute the test loss and additional metrics.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n        \"\"\"\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self(x)\n        loss = self.criterion(y_hat, y)\n        self.log(\"test_loss\", loss)\n        self.test_metrics(y_hat, y)\n        self.log_dict(self.test_metrics)\n\n    def predict_step(\n        self, batch: Any, batch_idx: int, dataloader_idx: int = 0\n    ) -> Tensor:\n        \"\"\"Compute the predicted class probabilities.\n\n        Args:\n            batch: The output of your DataLoader.\n            batch_idx: Integer displaying index of this batch.\n            dataloader_idx: Index of the current dataloader.\n\n        Returns:\n            Output predicted probabilities.\n        \"\"\"\n        x = batch[\"image\"]\n        y_hat: Tensor = self(x).softmax(dim=1)\n        return y_hat\n", "path": "torchgeo/trainers/segmentation.py"}]}