We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Bug: missing "not" operator to set ignore_zeros in segmentation trainer
I believe there's a bug in the translation from ignore_zeros to ignore_index in the [\_\_init__ of SemanticSegmentationTask](https://github.com/microsoft/torchgeo/blob/c933694ddc9a121af9415389949aefcc84760472/torchgeo/trainers/segmentation.py#L89).
`self.ignore_zeros = None if kwargs["ignore_zeros"] else 0` should be:
`self.ignore_zeros = None if kwargs["ignore_zeros"] is None else 0` as is done [when define loss function](https://github.com/microsoft/torchgeo/blob/c933694ddc9a121af9415389949aefcc84760472/torchgeo/trainers/segmentation.py#L57)

`self.ignore_zeros` is later used to define the "self.ignore_index" of the self.train_metrics. For sake of clarity, I'd suggest renaming it to  `self.ignore_index` directly in the \_\_init__. 

Therefore, this
```
...
        self.ignore_zeros = None if kwargs["ignore_zeros"] else 0

        self.config_task()

        self.train_metrics = MetricCollection(
            [
                Accuracy(
                    num_classes=self.hparams["num_classes"],
                    ignore_index=self.ignore_zeros,
...
```
would become
```
...
        self.ignore_index = None if kwargs["ignore_zeros"] is None else 0

        self.config_task()

        self.train_metrics = MetricCollection(
            [
                Accuracy(
                    num_classes=self.hparams["num_classes"],
                    ignore_index=self.ignore_index,
...
```
If you'd like, I can definitely push a PR for to fix this bug if the solution seems ok to all. Let me know.

--- 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 """Segmentation tasks."""
5 
6 from typing import Any, Dict, cast
7 
8 import segmentation_models_pytorch as smp
9 import torch
10 import torch.nn as nn
11 from pytorch_lightning.core.lightning import LightningModule
12 from torch import Tensor
13 from torch.optim.lr_scheduler import ReduceLROnPlateau
14 from torch.utils.data import DataLoader
15 from torchmetrics import Accuracy, JaccardIndex, MetricCollection
16 
17 from ..datasets.utils import unbind_samples
18 from ..models import FCN
19 
20 # https://github.com/pytorch/pytorch/issues/60979
21 # https://github.com/pytorch/pytorch/pull/61045
22 DataLoader.__module__ = "torch.utils.data"
23 
24 
25 class SemanticSegmentationTask(LightningModule):
26     """LightningModule for semantic segmentation of images."""
27 
28     def config_task(self) -> None:
29         """Configures the task based on kwargs parameters passed to the constructor."""
30         if self.hyperparams["segmentation_model"] == "unet":
31             self.model = smp.Unet(
32                 encoder_name=self.hyperparams["encoder_name"],
33                 encoder_weights=self.hyperparams["encoder_weights"],
34                 in_channels=self.hyperparams["in_channels"],
35                 classes=self.hyperparams["num_classes"],
36             )
37         elif self.hyperparams["segmentation_model"] == "deeplabv3+":
38             self.model = smp.DeepLabV3Plus(
39                 encoder_name=self.hyperparams["encoder_name"],
40                 encoder_weights=self.hyperparams["encoder_weights"],
41                 in_channels=self.hyperparams["in_channels"],
42                 classes=self.hyperparams["num_classes"],
43             )
44         elif self.hyperparams["segmentation_model"] == "fcn":
45             self.model = FCN(
46                 in_channels=self.hyperparams["in_channels"],
47                 classes=self.hyperparams["num_classes"],
48                 num_filters=self.hyperparams["num_filters"],
49             )
50         else:
51             raise ValueError(
52                 f"Model type '{self.hyperparams['segmentation_model']}' is not valid."
53             )
54 
55         if self.hyperparams["loss"] == "ce":
56             self.loss = nn.CrossEntropyLoss(
57                 ignore_index=-1000 if self.ignore_zeros is None else 0
58             )
59         elif self.hyperparams["loss"] == "jaccard":
60             self.loss = smp.losses.JaccardLoss(
61                 mode="multiclass", classes=self.hyperparams["num_classes"]
62             )
63         elif self.hyperparams["loss"] == "focal":
64             self.loss = smp.losses.FocalLoss(
65                 "multiclass", ignore_index=self.ignore_zeros, normalized=True
66             )
67         else:
68             raise ValueError(f"Loss type '{self.hyperparams['loss']}' is not valid.")
69 
70     def __init__(self, **kwargs: Any) -> None:
71         """Initialize the LightningModule with a model and loss function.
72 
73         Keyword Args:
74             segmentation_model: Name of the segmentation model type to use
75             encoder_name: Name of the encoder model backbone to use
76             encoder_weights: None or "imagenet" to use imagenet pretrained weights in
77                 the encoder model
78             in_channels: Number of channels in input image
79             num_classes: Number of semantic classes to predict
80             loss: Name of the loss function
81             ignore_zeros: Whether to ignore the "0" class value in the loss and metrics
82 
83         Raises:
84             ValueError: if kwargs arguments are invalid
85         """
86         super().__init__()
87 
88         # Creates `self.hparams` from kwargs
89         self.save_hyperparameters()  # type: ignore[operator]
90         self.hyperparams = cast(Dict[str, Any], self.hparams)
91 
92         self.ignore_zeros = None if kwargs["ignore_zeros"] else 0
93 
94         self.config_task()
95 
96         self.train_metrics = MetricCollection(
97             [
98                 Accuracy(
99                     num_classes=self.hyperparams["num_classes"],
100                     ignore_index=self.ignore_zeros,
101                     mdmc_average="global",
102                 ),
103                 JaccardIndex(
104                     num_classes=self.hyperparams["num_classes"],
105                     ignore_index=self.ignore_zeros,
106                 ),
107             ],
108             prefix="train_",
109         )
110         self.val_metrics = self.train_metrics.clone(prefix="val_")
111         self.test_metrics = self.train_metrics.clone(prefix="test_")
112 
113     def forward(self, *args: Any, **kwargs: Any) -> Any:
114         """Forward pass of the model.
115 
116         Args:
117             x: tensor of data to run through the model
118 
119         Returns:
120             output from the model
121         """
122         return self.model(*args, **kwargs)
123 
124     def training_step(self, *args: Any, **kwargs: Any) -> Tensor:
125         """Compute and return the training loss.
126 
127         Args:
128             batch: the output of your DataLoader
129 
130         Returns:
131             training loss
132         """
133         batch = args[0]
134         x = batch["image"]
135         y = batch["mask"]
136         y_hat = self.forward(x)
137         y_hat_hard = y_hat.argmax(dim=1)
138 
139         loss = self.loss(y_hat, y)
140 
141         # by default, the train step logs every `log_every_n_steps` steps where
142         # `log_every_n_steps` is a parameter to the `Trainer` object
143         self.log("train_loss", loss, on_step=True, on_epoch=False)
144         self.train_metrics(y_hat_hard, y)
145 
146         return cast(Tensor, loss)
147 
148     def training_epoch_end(self, outputs: Any) -> None:
149         """Logs epoch level training metrics.
150 
151         Args:
152             outputs: list of items returned by training_step
153         """
154         self.log_dict(self.train_metrics.compute())
155         self.train_metrics.reset()
156 
157     def validation_step(self, *args: Any, **kwargs: Any) -> None:
158         """Compute validation loss and log example predictions.
159 
160         Args:
161             batch: the output of your DataLoader
162             batch_idx: the index of this batch
163         """
164         batch = args[0]
165         batch_idx = args[1]
166         x = batch["image"]
167         y = batch["mask"]
168         y_hat = self.forward(x)
169         y_hat_hard = y_hat.argmax(dim=1)
170 
171         loss = self.loss(y_hat, y)
172 
173         self.log("val_loss", loss, on_step=False, on_epoch=True)
174         self.val_metrics(y_hat_hard, y)
175 
176         if batch_idx < 10:
177             try:
178                 datamodule = self.trainer.datamodule  # type: ignore[union-attr]
179                 batch["prediction"] = y_hat_hard
180                 for key in ["image", "mask", "prediction"]:
181                     batch[key] = batch[key].cpu()
182                 sample = unbind_samples(batch)[0]
183                 fig = datamodule.plot(sample)
184                 summary_writer = self.logger.experiment  # type: ignore[union-attr]
185                 summary_writer.add_figure(
186                     f"image/{batch_idx}", fig, global_step=self.global_step
187                 )
188             except AttributeError:
189                 pass
190 
191     def validation_epoch_end(self, outputs: Any) -> None:
192         """Logs epoch level validation metrics.
193 
194         Args:
195             outputs: list of items returned by validation_step
196         """
197         self.log_dict(self.val_metrics.compute())
198         self.val_metrics.reset()
199 
200     def test_step(self, *args: Any, **kwargs: Any) -> None:
201         """Compute test loss.
202 
203         Args:
204             batch: the output of your DataLoader
205         """
206         batch = args[0]
207         x = batch["image"]
208         y = batch["mask"]
209         y_hat = self.forward(x)
210         y_hat_hard = y_hat.argmax(dim=1)
211 
212         loss = self.loss(y_hat, y)
213 
214         # by default, the test and validation steps only log per *epoch*
215         self.log("test_loss", loss, on_step=False, on_epoch=True)
216         self.test_metrics(y_hat_hard, y)
217 
218     def test_epoch_end(self, outputs: Any) -> None:
219         """Logs epoch level test metrics.
220 
221         Args:
222             outputs: list of items returned by test_step
223         """
224         self.log_dict(self.test_metrics.compute())
225         self.test_metrics.reset()
226 
227     def configure_optimizers(self) -> Dict[str, Any]:
228         """Initialize the optimizer and learning rate scheduler.
229 
230         Returns:
231             a "lr dict" according to the pytorch lightning documentation --
232             https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers
233         """
234         optimizer = torch.optim.Adam(
235             self.model.parameters(), lr=self.hyperparams["learning_rate"]
236         )
237         return {
238             "optimizer": optimizer,
239             "lr_scheduler": {
240                 "scheduler": ReduceLROnPlateau(
241                     optimizer,
242                     patience=self.hyperparams["learning_rate_schedule_patience"],
243                 ),
244                 "monitor": "val_loss",
245             },
246         }
247
```
--- 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
@@ -3,6 +3,7 @@
 
 """Segmentation tasks."""
 
+import warnings
 from typing import Any, Dict, cast
 
 import segmentation_models_pytorch as smp
@@ -53,16 +54,15 @@
             )
 
         if self.hyperparams["loss"] == "ce":
-            self.loss = nn.CrossEntropyLoss(
-                ignore_index=-1000 if self.ignore_zeros is None else 0
-            )
+            ignore_value = -1000 if self.ignore_index is None else self.ignore_index
+            self.loss = nn.CrossEntropyLoss(ignore_index=ignore_value)
         elif self.hyperparams["loss"] == "jaccard":
             self.loss = smp.losses.JaccardLoss(
                 mode="multiclass", classes=self.hyperparams["num_classes"]
             )
         elif self.hyperparams["loss"] == "focal":
             self.loss = smp.losses.FocalLoss(
-                "multiclass", ignore_index=self.ignore_zeros, normalized=True
+                "multiclass", ignore_index=self.ignore_index, normalized=True
             )
         else:
             raise ValueError(f"Loss type '{self.hyperparams['loss']}' is not valid.")
@@ -78,7 +78,7 @@
             in_channels: Number of channels in input image
             num_classes: Number of semantic classes to predict
             loss: Name of the loss function
-            ignore_zeros: Whether to ignore the "0" class value in the loss and metrics
+            ignore_index: Optional integer class index to ignore in the loss and metrics
 
         Raises:
             ValueError: if kwargs arguments are invalid
@@ -89,20 +89,26 @@
         self.save_hyperparameters()  # type: ignore[operator]
         self.hyperparams = cast(Dict[str, Any], self.hparams)
 
-        self.ignore_zeros = None if kwargs["ignore_zeros"] else 0
-
+        if not isinstance(kwargs["ignore_index"], (int, type(None))):
+            raise ValueError("ignore_index must be an int or None")
+        if (kwargs["ignore_index"] is not None) and (kwargs["loss"] == "jaccard"):
+            warnings.warn(
+                "ignore_index has no effect on training when loss='jaccard'",
+                UserWarning,
+            )
+        self.ignore_index = kwargs["ignore_index"]
         self.config_task()
 
         self.train_metrics = MetricCollection(
             [
                 Accuracy(
                     num_classes=self.hyperparams["num_classes"],
-                    ignore_index=self.ignore_zeros,
+                    ignore_index=self.ignore_index,
                     mdmc_average="global",
                 ),
                 JaccardIndex(
                     num_classes=self.hyperparams["num_classes"],
-                    ignore_index=self.ignore_zeros,
+                    ignore_index=self.ignore_index,
                 ),
             ],
             prefix="train_",

{"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@@ -3,6 +3,7 @@\n \n \"\"\"Segmentation tasks.\"\"\"\n \n+import warnings\n from typing import Any, Dict, cast\n \n import segmentation_models_pytorch as smp\n@@ -53,16 +54,15 @@\n             )\n \n         if self.hyperparams[\"loss\"] == \"ce\":\n-            self.loss = nn.CrossEntropyLoss(\n-                ignore_index=-1000 if self.ignore_zeros is None else 0\n-            )\n+            ignore_value = -1000 if self.ignore_index is None else self.ignore_index\n+            self.loss = nn.CrossEntropyLoss(ignore_index=ignore_value)\n         elif self.hyperparams[\"loss\"] == \"jaccard\":\n             self.loss = smp.losses.JaccardLoss(\n                 mode=\"multiclass\", classes=self.hyperparams[\"num_classes\"]\n             )\n         elif self.hyperparams[\"loss\"] == \"focal\":\n             self.loss = smp.losses.FocalLoss(\n-                \"multiclass\", ignore_index=self.ignore_zeros, normalized=True\n+                \"multiclass\", ignore_index=self.ignore_index, normalized=True\n             )\n         else:\n             raise ValueError(f\"Loss type '{self.hyperparams['loss']}' is not valid.\")\n@@ -78,7 +78,7 @@\n             in_channels: Number of channels in input image\n             num_classes: Number of semantic classes to predict\n             loss: Name of the loss function\n-            ignore_zeros: Whether to ignore the \"0\" class value in the loss and metrics\n+            ignore_index: Optional integer class index to ignore in the loss and metrics\n \n         Raises:\n             ValueError: if kwargs arguments are invalid\n@@ -89,20 +89,26 @@\n         self.save_hyperparameters()  # type: ignore[operator]\n         self.hyperparams = cast(Dict[str, Any], self.hparams)\n \n-        self.ignore_zeros = None if kwargs[\"ignore_zeros\"] else 0\n-\n+        if not isinstance(kwargs[\"ignore_index\"], (int, type(None))):\n+            raise ValueError(\"ignore_index must be an int or None\")\n+        if (kwargs[\"ignore_index\"] is not None) and (kwargs[\"loss\"] == \"jaccard\"):\n+            warnings.warn(\n+                \"ignore_index has no effect on training when loss='jaccard'\",\n+                UserWarning,\n+            )\n+        self.ignore_index = kwargs[\"ignore_index\"]\n         self.config_task()\n \n         self.train_metrics = MetricCollection(\n             [\n                 Accuracy(\n                     num_classes=self.hyperparams[\"num_classes\"],\n-                    ignore_index=self.ignore_zeros,\n+                    ignore_index=self.ignore_index,\n                     mdmc_average=\"global\",\n                 ),\n                 JaccardIndex(\n                     num_classes=self.hyperparams[\"num_classes\"],\n-                    ignore_index=self.ignore_zeros,\n+                    ignore_index=self.ignore_index,\n                 ),\n             ],\n             prefix=\"train_\",\n", "issue": "Bug: missing \"not\" operator to set ignore_zeros in segmentation trainer\nI believe there's a bug in the translation from ignore_zeros to ignore_index in the [\\_\\_init__ of SemanticSegmentationTask](https://github.com/microsoft/torchgeo/blob/c933694ddc9a121af9415389949aefcc84760472/torchgeo/trainers/segmentation.py#L89).\r\n`self.ignore_zeros = None if kwargs[\"ignore_zeros\"] else 0` should be:\r\n`self.ignore_zeros = None if kwargs[\"ignore_zeros\"] is None else 0` as is done [when define loss function](https://github.com/microsoft/torchgeo/blob/c933694ddc9a121af9415389949aefcc84760472/torchgeo/trainers/segmentation.py#L57)\r\n\r\n`self.ignore_zeros` is later used to define the \"self.ignore_index\" of the self.train_metrics. For sake of clarity, I'd suggest renaming it to  `self.ignore_index` directly in the \\_\\_init__. \r\n\r\nTherefore, this\r\n```\r\n...\r\n        self.ignore_zeros = None if kwargs[\"ignore_zeros\"] else 0\r\n\r\n        self.config_task()\r\n\r\n        self.train_metrics = MetricCollection(\r\n            [\r\n                Accuracy(\r\n                    num_classes=self.hparams[\"num_classes\"],\r\n                    ignore_index=self.ignore_zeros,\r\n...\r\n```\r\nwould become\r\n```\r\n...\r\n        self.ignore_index = None if kwargs[\"ignore_zeros\"] is None else 0\r\n\r\n        self.config_task()\r\n\r\n        self.train_metrics = MetricCollection(\r\n            [\r\n                Accuracy(\r\n                    num_classes=self.hparams[\"num_classes\"],\r\n                    ignore_index=self.ignore_index,\r\n...\r\n```\r\nIf you'd like, I can definitely push a PR for to fix this bug if the solution seems ok to all. Let me know.\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"Segmentation tasks.\"\"\"\n\nfrom typing import Any, Dict, cast\n\nimport segmentation_models_pytorch as smp\nimport torch\nimport torch.nn as nn\nfrom pytorch_lightning.core.lightning import LightningModule\nfrom torch import Tensor\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torch.utils.data import DataLoader\nfrom torchmetrics import Accuracy, JaccardIndex, MetricCollection\n\nfrom ..datasets.utils import unbind_samples\nfrom ..models import FCN\n\n# https://github.com/pytorch/pytorch/issues/60979\n# https://github.com/pytorch/pytorch/pull/61045\nDataLoader.__module__ = \"torch.utils.data\"\n\n\nclass SemanticSegmentationTask(LightningModule):\n    \"\"\"LightningModule for semantic segmentation of images.\"\"\"\n\n    def config_task(self) -> None:\n        \"\"\"Configures the task based on kwargs parameters passed to the constructor.\"\"\"\n        if self.hyperparams[\"segmentation_model\"] == \"unet\":\n            self.model = smp.Unet(\n                encoder_name=self.hyperparams[\"encoder_name\"],\n                encoder_weights=self.hyperparams[\"encoder_weights\"],\n                in_channels=self.hyperparams[\"in_channels\"],\n                classes=self.hyperparams[\"num_classes\"],\n            )\n        elif self.hyperparams[\"segmentation_model\"] == \"deeplabv3+\":\n            self.model = smp.DeepLabV3Plus(\n                encoder_name=self.hyperparams[\"encoder_name\"],\n                encoder_weights=self.hyperparams[\"encoder_weights\"],\n                in_channels=self.hyperparams[\"in_channels\"],\n                classes=self.hyperparams[\"num_classes\"],\n            )\n        elif self.hyperparams[\"segmentation_model\"] == \"fcn\":\n            self.model = FCN(\n                in_channels=self.hyperparams[\"in_channels\"],\n                classes=self.hyperparams[\"num_classes\"],\n                num_filters=self.hyperparams[\"num_filters\"],\n            )\n        else:\n            raise ValueError(\n                f\"Model type '{self.hyperparams['segmentation_model']}' is not valid.\"\n            )\n\n        if self.hyperparams[\"loss\"] == \"ce\":\n            self.loss = nn.CrossEntropyLoss(\n                ignore_index=-1000 if self.ignore_zeros is None else 0\n            )\n        elif self.hyperparams[\"loss\"] == \"jaccard\":\n            self.loss = smp.losses.JaccardLoss(\n                mode=\"multiclass\", classes=self.hyperparams[\"num_classes\"]\n            )\n        elif self.hyperparams[\"loss\"] == \"focal\":\n            self.loss = smp.losses.FocalLoss(\n                \"multiclass\", ignore_index=self.ignore_zeros, normalized=True\n            )\n        else:\n            raise ValueError(f\"Loss type '{self.hyperparams['loss']}' is not valid.\")\n\n    def __init__(self, **kwargs: Any) -> None:\n        \"\"\"Initialize the LightningModule with a model and loss function.\n\n        Keyword Args:\n            segmentation_model: Name of the segmentation model type to use\n            encoder_name: Name of the encoder model backbone to use\n            encoder_weights: None or \"imagenet\" to use imagenet pretrained weights in\n                the encoder model\n            in_channels: Number of channels in input image\n            num_classes: Number of semantic classes to predict\n            loss: Name of the loss function\n            ignore_zeros: Whether to ignore the \"0\" class value in the loss and metrics\n\n        Raises:\n            ValueError: if kwargs arguments are invalid\n        \"\"\"\n        super().__init__()\n\n        # Creates `self.hparams` from kwargs\n        self.save_hyperparameters()  # type: ignore[operator]\n        self.hyperparams = cast(Dict[str, Any], self.hparams)\n\n        self.ignore_zeros = None if kwargs[\"ignore_zeros\"] else 0\n\n        self.config_task()\n\n        self.train_metrics = MetricCollection(\n            [\n                Accuracy(\n                    num_classes=self.hyperparams[\"num_classes\"],\n                    ignore_index=self.ignore_zeros,\n                    mdmc_average=\"global\",\n                ),\n                JaccardIndex(\n                    num_classes=self.hyperparams[\"num_classes\"],\n                    ignore_index=self.ignore_zeros,\n                ),\n            ],\n            prefix=\"train_\",\n        )\n        self.val_metrics = self.train_metrics.clone(prefix=\"val_\")\n        self.test_metrics = self.train_metrics.clone(prefix=\"test_\")\n\n    def forward(self, *args: Any, **kwargs: Any) -> Any:\n        \"\"\"Forward pass of the model.\n\n        Args:\n            x: tensor of data to run through the model\n\n        Returns:\n            output from the model\n        \"\"\"\n        return self.model(*args, **kwargs)\n\n    def training_step(self, *args: Any, **kwargs: Any) -> Tensor:\n        \"\"\"Compute and return the training loss.\n\n        Args:\n            batch: the output of your DataLoader\n\n        Returns:\n            training loss\n        \"\"\"\n        batch = args[0]\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self.forward(x)\n        y_hat_hard = y_hat.argmax(dim=1)\n\n        loss = self.loss(y_hat, y)\n\n        # by default, the train step logs every `log_every_n_steps` steps where\n        # `log_every_n_steps` is a parameter to the `Trainer` object\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=False)\n        self.train_metrics(y_hat_hard, y)\n\n        return cast(Tensor, loss)\n\n    def training_epoch_end(self, outputs: Any) -> None:\n        \"\"\"Logs epoch level training metrics.\n\n        Args:\n            outputs: list of items returned by training_step\n        \"\"\"\n        self.log_dict(self.train_metrics.compute())\n        self.train_metrics.reset()\n\n    def validation_step(self, *args: Any, **kwargs: Any) -> None:\n        \"\"\"Compute validation loss and log example predictions.\n\n        Args:\n            batch: the output of your DataLoader\n            batch_idx: the index of this batch\n        \"\"\"\n        batch = args[0]\n        batch_idx = args[1]\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self.forward(x)\n        y_hat_hard = y_hat.argmax(dim=1)\n\n        loss = self.loss(y_hat, y)\n\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True)\n        self.val_metrics(y_hat_hard, y)\n\n        if batch_idx < 10:\n            try:\n                datamodule = self.trainer.datamodule  # type: ignore[union-attr]\n                batch[\"prediction\"] = y_hat_hard\n                for key in [\"image\", \"mask\", \"prediction\"]:\n                    batch[key] = batch[key].cpu()\n                sample = unbind_samples(batch)[0]\n                fig = datamodule.plot(sample)\n                summary_writer = self.logger.experiment  # type: ignore[union-attr]\n                summary_writer.add_figure(\n                    f\"image/{batch_idx}\", fig, global_step=self.global_step\n                )\n            except AttributeError:\n                pass\n\n    def validation_epoch_end(self, outputs: Any) -> None:\n        \"\"\"Logs epoch level validation metrics.\n\n        Args:\n            outputs: list of items returned by validation_step\n        \"\"\"\n        self.log_dict(self.val_metrics.compute())\n        self.val_metrics.reset()\n\n    def test_step(self, *args: Any, **kwargs: Any) -> None:\n        \"\"\"Compute test loss.\n\n        Args:\n            batch: the output of your DataLoader\n        \"\"\"\n        batch = args[0]\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self.forward(x)\n        y_hat_hard = y_hat.argmax(dim=1)\n\n        loss = self.loss(y_hat, y)\n\n        # by default, the test and validation steps only log per *epoch*\n        self.log(\"test_loss\", loss, on_step=False, on_epoch=True)\n        self.test_metrics(y_hat_hard, y)\n\n    def test_epoch_end(self, outputs: Any) -> None:\n        \"\"\"Logs epoch level test metrics.\n\n        Args:\n            outputs: list of items returned by test_step\n        \"\"\"\n        self.log_dict(self.test_metrics.compute())\n        self.test_metrics.reset()\n\n    def configure_optimizers(self) -> Dict[str, Any]:\n        \"\"\"Initialize the optimizer and learning rate scheduler.\n\n        Returns:\n            a \"lr dict\" according to the pytorch lightning documentation --\n            https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers\n        \"\"\"\n        optimizer = torch.optim.Adam(\n            self.model.parameters(), lr=self.hyperparams[\"learning_rate\"]\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": ReduceLROnPlateau(\n                    optimizer,\n                    patience=self.hyperparams[\"learning_rate_schedule_patience\"],\n                ),\n                \"monitor\": \"val_loss\",\n            },\n        }\n", "path": "torchgeo/trainers/segmentation.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"Segmentation tasks.\"\"\"\n\nimport warnings\nfrom typing import Any, Dict, cast\n\nimport segmentation_models_pytorch as smp\nimport torch\nimport torch.nn as nn\nfrom pytorch_lightning.core.lightning import LightningModule\nfrom torch import Tensor\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torch.utils.data import DataLoader\nfrom torchmetrics import Accuracy, JaccardIndex, MetricCollection\n\nfrom ..datasets.utils import unbind_samples\nfrom ..models import FCN\n\n# https://github.com/pytorch/pytorch/issues/60979\n# https://github.com/pytorch/pytorch/pull/61045\nDataLoader.__module__ = \"torch.utils.data\"\n\n\nclass SemanticSegmentationTask(LightningModule):\n    \"\"\"LightningModule for semantic segmentation of images.\"\"\"\n\n    def config_task(self) -> None:\n        \"\"\"Configures the task based on kwargs parameters passed to the constructor.\"\"\"\n        if self.hyperparams[\"segmentation_model\"] == \"unet\":\n            self.model = smp.Unet(\n                encoder_name=self.hyperparams[\"encoder_name\"],\n                encoder_weights=self.hyperparams[\"encoder_weights\"],\n                in_channels=self.hyperparams[\"in_channels\"],\n                classes=self.hyperparams[\"num_classes\"],\n            )\n        elif self.hyperparams[\"segmentation_model\"] == \"deeplabv3+\":\n            self.model = smp.DeepLabV3Plus(\n                encoder_name=self.hyperparams[\"encoder_name\"],\n                encoder_weights=self.hyperparams[\"encoder_weights\"],\n                in_channels=self.hyperparams[\"in_channels\"],\n                classes=self.hyperparams[\"num_classes\"],\n            )\n        elif self.hyperparams[\"segmentation_model\"] == \"fcn\":\n            self.model = FCN(\n                in_channels=self.hyperparams[\"in_channels\"],\n                classes=self.hyperparams[\"num_classes\"],\n                num_filters=self.hyperparams[\"num_filters\"],\n            )\n        else:\n            raise ValueError(\n                f\"Model type '{self.hyperparams['segmentation_model']}' is not valid.\"\n            )\n\n        if self.hyperparams[\"loss\"] == \"ce\":\n            ignore_value = -1000 if self.ignore_index is None else self.ignore_index\n            self.loss = nn.CrossEntropyLoss(ignore_index=ignore_value)\n        elif self.hyperparams[\"loss\"] == \"jaccard\":\n            self.loss = smp.losses.JaccardLoss(\n                mode=\"multiclass\", classes=self.hyperparams[\"num_classes\"]\n            )\n        elif self.hyperparams[\"loss\"] == \"focal\":\n            self.loss = smp.losses.FocalLoss(\n                \"multiclass\", ignore_index=self.ignore_index, normalized=True\n            )\n        else:\n            raise ValueError(f\"Loss type '{self.hyperparams['loss']}' is not valid.\")\n\n    def __init__(self, **kwargs: Any) -> None:\n        \"\"\"Initialize the LightningModule with a model and loss function.\n\n        Keyword Args:\n            segmentation_model: Name of the segmentation model type to use\n            encoder_name: Name of the encoder model backbone to use\n            encoder_weights: None or \"imagenet\" to use imagenet pretrained weights in\n                the encoder model\n            in_channels: Number of channels in input image\n            num_classes: Number of semantic classes to predict\n            loss: Name of the loss function\n            ignore_index: Optional integer class index to ignore in the loss and metrics\n\n        Raises:\n            ValueError: if kwargs arguments are invalid\n        \"\"\"\n        super().__init__()\n\n        # Creates `self.hparams` from kwargs\n        self.save_hyperparameters()  # type: ignore[operator]\n        self.hyperparams = cast(Dict[str, Any], self.hparams)\n\n        if not isinstance(kwargs[\"ignore_index\"], (int, type(None))):\n            raise ValueError(\"ignore_index must be an int or None\")\n        if (kwargs[\"ignore_index\"] is not None) and (kwargs[\"loss\"] == \"jaccard\"):\n            warnings.warn(\n                \"ignore_index has no effect on training when loss='jaccard'\",\n                UserWarning,\n            )\n        self.ignore_index = kwargs[\"ignore_index\"]\n        self.config_task()\n\n        self.train_metrics = MetricCollection(\n            [\n                Accuracy(\n                    num_classes=self.hyperparams[\"num_classes\"],\n                    ignore_index=self.ignore_index,\n                    mdmc_average=\"global\",\n                ),\n                JaccardIndex(\n                    num_classes=self.hyperparams[\"num_classes\"],\n                    ignore_index=self.ignore_index,\n                ),\n            ],\n            prefix=\"train_\",\n        )\n        self.val_metrics = self.train_metrics.clone(prefix=\"val_\")\n        self.test_metrics = self.train_metrics.clone(prefix=\"test_\")\n\n    def forward(self, *args: Any, **kwargs: Any) -> Any:\n        \"\"\"Forward pass of the model.\n\n        Args:\n            x: tensor of data to run through the model\n\n        Returns:\n            output from the model\n        \"\"\"\n        return self.model(*args, **kwargs)\n\n    def training_step(self, *args: Any, **kwargs: Any) -> Tensor:\n        \"\"\"Compute and return the training loss.\n\n        Args:\n            batch: the output of your DataLoader\n\n        Returns:\n            training loss\n        \"\"\"\n        batch = args[0]\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self.forward(x)\n        y_hat_hard = y_hat.argmax(dim=1)\n\n        loss = self.loss(y_hat, y)\n\n        # by default, the train step logs every `log_every_n_steps` steps where\n        # `log_every_n_steps` is a parameter to the `Trainer` object\n        self.log(\"train_loss\", loss, on_step=True, on_epoch=False)\n        self.train_metrics(y_hat_hard, y)\n\n        return cast(Tensor, loss)\n\n    def training_epoch_end(self, outputs: Any) -> None:\n        \"\"\"Logs epoch level training metrics.\n\n        Args:\n            outputs: list of items returned by training_step\n        \"\"\"\n        self.log_dict(self.train_metrics.compute())\n        self.train_metrics.reset()\n\n    def validation_step(self, *args: Any, **kwargs: Any) -> None:\n        \"\"\"Compute validation loss and log example predictions.\n\n        Args:\n            batch: the output of your DataLoader\n            batch_idx: the index of this batch\n        \"\"\"\n        batch = args[0]\n        batch_idx = args[1]\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self.forward(x)\n        y_hat_hard = y_hat.argmax(dim=1)\n\n        loss = self.loss(y_hat, y)\n\n        self.log(\"val_loss\", loss, on_step=False, on_epoch=True)\n        self.val_metrics(y_hat_hard, y)\n\n        if batch_idx < 10:\n            try:\n                datamodule = self.trainer.datamodule  # type: ignore[union-attr]\n                batch[\"prediction\"] = y_hat_hard\n                for key in [\"image\", \"mask\", \"prediction\"]:\n                    batch[key] = batch[key].cpu()\n                sample = unbind_samples(batch)[0]\n                fig = datamodule.plot(sample)\n                summary_writer = self.logger.experiment  # type: ignore[union-attr]\n                summary_writer.add_figure(\n                    f\"image/{batch_idx}\", fig, global_step=self.global_step\n                )\n            except AttributeError:\n                pass\n\n    def validation_epoch_end(self, outputs: Any) -> None:\n        \"\"\"Logs epoch level validation metrics.\n\n        Args:\n            outputs: list of items returned by validation_step\n        \"\"\"\n        self.log_dict(self.val_metrics.compute())\n        self.val_metrics.reset()\n\n    def test_step(self, *args: Any, **kwargs: Any) -> None:\n        \"\"\"Compute test loss.\n\n        Args:\n            batch: the output of your DataLoader\n        \"\"\"\n        batch = args[0]\n        x = batch[\"image\"]\n        y = batch[\"mask\"]\n        y_hat = self.forward(x)\n        y_hat_hard = y_hat.argmax(dim=1)\n\n        loss = self.loss(y_hat, y)\n\n        # by default, the test and validation steps only log per *epoch*\n        self.log(\"test_loss\", loss, on_step=False, on_epoch=True)\n        self.test_metrics(y_hat_hard, y)\n\n    def test_epoch_end(self, outputs: Any) -> None:\n        \"\"\"Logs epoch level test metrics.\n\n        Args:\n            outputs: list of items returned by test_step\n        \"\"\"\n        self.log_dict(self.test_metrics.compute())\n        self.test_metrics.reset()\n\n    def configure_optimizers(self) -> Dict[str, Any]:\n        \"\"\"Initialize the optimizer and learning rate scheduler.\n\n        Returns:\n            a \"lr dict\" according to the pytorch lightning documentation --\n            https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#configure-optimizers\n        \"\"\"\n        optimizer = torch.optim.Adam(\n            self.model.parameters(), lr=self.hyperparams[\"learning_rate\"]\n        )\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\n                \"scheduler\": ReduceLROnPlateau(\n                    optimizer,\n                    patience=self.hyperparams[\"learning_rate_schedule_patience\"],\n                ),\n                \"monitor\": \"val_loss\",\n            },\n        }\n", "path": "torchgeo/trainers/segmentation.py"}]}