Repository: microsoft/torchgeo
Pull request #2219: Trainers: support binary, multiclass, and multilabel tasks
Author: login: adamjstewart
Labels: name: documentation
description: Improvements or additions to documentation
name: datasets
description: Geospatial or benchmark datasets
name: testing
description: Continuous integration testing
name: trainers
description: PyTorch Lightning trainers
name: backwards-incompatible
description: Changes that are not backwards compatible
Created: 2024-08-12T15:08:39Z | Merged: 2025-03-27T18:09:50Z

Instead of having separate trainers for binary, multiclass, and multilabel, let's create a single trainer that can handle all 3.

This applies to both Classification and Semantic Segmentation but not to our other trainers.

### Deprecation Note

`MultiLabelClassificationTask(num_classes=...)` is deprecated, use `ClassificationTask(task='multilabel', num_labels=...)` instead.

### References

These links were really helpful in making sure things are supported correctly:

* https://discuss.pytorch.org/t/single-vs-multi-label-vs-multi-class-training-best-practice/143089/3

Closes #2205 @robmarkcole
Closes #599 @kethan52 @geografif
Closes #513 @kethan52 
Closes #245 @calebrob6

Linked issue #2205: Multiclass Classification: assert num_classes >=2
### Summary

Both segmentation and object detection require that the background be included and there is currently a note on these args: `num_classes: Number of prediction classes (including the background)`. Considering every dataaset must have at least 1 class, the min value of num_classes is 2. I propose adding an assertion, to prevent people (like myself!) from forgetting this and setting `num_classes=1` for datasets with a single class. 

### Rationale

This config error has happened to me several times, and can pass silently

### Implementation

I suppose we add validation to the BaseTask init

### Alternatives

_No response_

### Additional information

_No response_
Linked issue comments: user: login: adamjstewart
created_at: 2024-08-04T11:32:14Z
body: Not to completely derail what should otherwise be a simple fix, but...

This brings up the question of how we want to handle different forms of classification/semantic segmentation:

* Binary
* Multiclass
* Multilabel

Torchmetrics originally had a single class for `Accuracy`. In https://github.com/Lightning-AI/torchmetrics/issues/1001, they proposed and implemented separate classes for each of the 3 above types of classification (`BinaryAccuracy`, etc.). The original plan was to deprecate and remove the old single class, but it seems that plan was aborted at some point.

We should decide whether we want `BinaryClassificationTask`, etc. or whether we want to add a `task='binary'`, etc. parameter to `ClassificationTask`.

We could definitely still add such an assertion for now and change it to `assert num_classes > 1 if task != 'binary'` later if needed.
author_association: MEMBER
user: login: robmarkcole
created_at: 2024-08-04T11:56:54Z
body: As you point out, binary etc are args torchmetrics accepts, so I think it makes sense to have this functionality with the existing task
author_association: COLLABORATOR
user: login: adamjstewart
created_at: 2024-08-04T11:58:43Z
body: Just waiting for clarity on whether torchmetrics is planning on supporting the old metrics forever before deciding, but I was leaning towards that too.
author_association: MEMBER
user: login: adamjstewart
created_at: 2024-08-04T16:29:40Z
body: Looks like I misinterpreted, both are supported.

Is there anything special we need to do in our trainers to support binary and multilabel, or do we literally just need to pass different `task` values to torchmetrics? If the former, we may want to split, but if the latter, I agree we should just keep the current classes and add a `task` parameter.
author_association: MEMBER

Diff:
diff --git a/docs/api/trainers.rst b/docs/api/trainers.rst
index 2cbd6547b94..714b394fb50 100644
--- a/docs/api/trainers.rst
+++ b/docs/api/trainers.rst
@@ -2,3 +2,4 @@ torchgeo.trainers
 =================
 
 .. automodule:: torchgeo.trainers
+   :exclude-members: MultiLabelClassificationTask
diff --git a/tests/conf/bigearthnet_all.yaml b/tests/conf/bigearthnet_all.yaml
index 3841248f2ae..eacec4f74ca 100644
--- a/tests/conf/bigearthnet_all.yaml
+++ b/tests/conf/bigearthnet_all.yaml
@@ -1,10 +1,11 @@
 model:
-  class_path: MultiLabelClassificationTask
+  class_path: ClassificationTask
   init_args:
     loss: 'bce'
     model: 'resnet18'
     in_channels: 14
-    num_classes: 19
+    task: 'multilabel'
+    num_labels: 19
 data:
   class_path: BigEarthNetDataModule
   init_args:
diff --git a/tests/conf/bigearthnet_s1.yaml b/tests/conf/bigearthnet_s1.yaml
index ccbe0a4be10..0762b93e007 100644
--- a/tests/conf/bigearthnet_s1.yaml
+++ b/tests/conf/bigearthnet_s1.yaml
@@ -1,10 +1,11 @@
 model:
-  class_path: MultiLabelClassificationTask
+  class_path: ClassificationTask
   init_args:
     loss: 'bce'
     model: 'resnet18'
     in_channels: 2
-    num_classes: 19
+    task: 'multilabel'
+    num_labels: 19
 data:
   class_path: BigEarthNetDataModule
   init_args:
diff --git a/tests/conf/bigearthnet_s2.yaml b/tests/conf/bigearthnet_s2.yaml
index 030ecb96023..5d22ef2ad22 100644
--- a/tests/conf/bigearthnet_s2.yaml
+++ b/tests/conf/bigearthnet_s2.yaml
@@ -1,10 +1,11 @@
 model:
-  class_path: MultiLabelClassificationTask
+  class_path: ClassificationTask
   init_args:
     loss: 'bce'
     model: 'resnet18'
     in_channels: 12
-    num_classes: 19
+    task: 'multilabel'
+    num_labels: 19
 data:
   class_path: BigEarthNetDataModule
   init_args:
diff --git a/tests/conf/cabuar.yaml b/tests/conf/cabuar.yaml
index 0a9ba151269..1cf1166cd78 100644
--- a/tests/conf/cabuar.yaml
+++ b/tests/conf/cabuar.yaml
@@ -1,13 +1,11 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'unet'
     backbone: 'resnet18'
     in_channels: 24
-    num_classes: 2
-    num_filters: 1
-    ignore_index: null
+    task: 'binary'
+    loss: 'bce'
 data:
   class_path: CaBuArDataModule
   init_args:
diff --git a/tests/conf/chabud.yaml b/tests/conf/chabud.yaml
index 3cff070bf59..03cf1132d85 100644
--- a/tests/conf/chabud.yaml
+++ b/tests/conf/chabud.yaml
@@ -1,13 +1,11 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'unet'
     backbone: 'resnet18'
     in_channels: 24
-    num_classes: 2
-    num_filters: 1
-    ignore_index: null
+    task: 'binary'
+    loss: 'bce'
 data:
   class_path: ChaBuDDataModule
   init_args:
diff --git a/tests/conf/etci2021.yaml b/tests/conf/etci2021.yaml
index 1606a870c52..71351210ef8 100644
--- a/tests/conf/etci2021.yaml
+++ b/tests/conf/etci2021.yaml
@@ -1,12 +1,12 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'unet'
     backbone: 'resnet18'
     in_channels: 6
-    num_classes: 2
+    task: 'binary'
     ignore_index: 0
+    loss: 'bce'
 data:
   class_path: ETCI2021DataModule
   init_args:
diff --git a/tests/conf/ftw.yaml b/tests/conf/ftw.yaml
index 7335d5ab1bf..031d3d36f27 100644
--- a/tests/conf/ftw.yaml
+++ b/tests/conf/ftw.yaml
@@ -1,13 +1,12 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'unet'
     backbone: 'resnet18'
     in_channels: 8
-    num_classes: 2
+    task: 'binary'
     num_filters: 1
-    ignore_index: null
+    loss: 'bce'
 data:
   class_path: FieldsOfTheWorldDataModule
   init_args:
diff --git a/tests/conf/inria.yaml b/tests/conf/inria.yaml
index 01e97d3b503..dd20a2daffc 100644
--- a/tests/conf/inria.yaml
+++ b/tests/conf/inria.yaml
@@ -1,12 +1,11 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'unet'
     backbone: 'resnet18'
     in_channels: 3
-    num_classes: 2
-    ignore_index: null
+    task: 'binary'
+    loss: 'bce'
 data:
   class_path: InriaAerialImageLabelingDataModule
   init_args:
diff --git a/tests/conf/mmflood.yaml b/tests/conf/mmflood.yaml
index 1a6301080a7..de2d3fb7d1c 100644
--- a/tests/conf/mmflood.yaml
+++ b/tests/conf/mmflood.yaml
@@ -1,12 +1,11 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'unet'
     backbone: 'resnet18'
     in_channels: 4
-    num_classes: 2
-    num_filters: 1
+    task: 'binary'
+    loss: 'bce'
     ignore_index: 255
 data:
   class_path: MMFloodDataModule
diff --git a/tests/conf/quakeset.yaml b/tests/conf/quakeset.yaml
index 0f1df7baa4b..e923fd14675 100644
--- a/tests/conf/quakeset.yaml
+++ b/tests/conf/quakeset.yaml
@@ -1,10 +1,10 @@
 model:
   class_path: ClassificationTask
   init_args:
-    loss: 'ce'
+    loss: 'bce'
     model: 'resnet18'
     in_channels: 4
-    num_classes: 2
+    task: 'binary'
 data:
   class_path: QuakeSetDataModule
   init_args:
diff --git a/tests/conf/sentinel2_south_america_soybean.yaml b/tests/conf/sentinel2_south_america_soybean.yaml
index 1c788d1f197..113ec976461 100644
--- a/tests/conf/sentinel2_south_america_soybean.yaml
+++ b/tests/conf/sentinel2_south_america_soybean.yaml
@@ -1,12 +1,11 @@
 model:
   class_path: SemanticSegmentationTask
   init_args:
-    loss: 'ce'
     model: 'deeplabv3+'
     backbone: 'resnet18'
     in_channels: 13
-    num_classes: 2
-    num_filters: 1
+    task: 'binary'
+    loss: 'bce'
 data:
   class_path: Sentinel2SouthAmericaSoybeanDataModule
   init_args:
diff --git a/tests/conf/treesatai.yaml b/tests/conf/treesatai.yaml
index e605b688b82..77cf283ffd4 100644
--- a/tests/conf/treesatai.yaml
+++ b/tests/conf/treesatai.yaml
@@ -1,9 +1,10 @@
 model:
-  class_path: MultiLabelClassificationTask
+  class_path: ClassificationTask
   init_args:
     model: 'resnet18'
     in_channels: 19
-    num_classes: 15
+    task: 'multilabel'
+    num_labels: 15
     loss: 'bce'
 data:
   class_path: TreeSatAIDataModule
diff --git a/tests/trainers/test_classification.py b/tests/trainers/test_classification.py
index 73be3b818e1..7ac788be9f1 100644
--- a/tests/trainers/test_classification.py
+++ b/tests/trainers/test_classification.py
@@ -18,8 +18,9 @@
     BigEarthNetDataModule,
     EuroSATDataModule,
     MisconfigurationException,
+    QuakeSetDataModule,
 )
-from torchgeo.datasets import BigEarthNet, EuroSAT, RGBBandsMissingError
+from torchgeo.datasets import BigEarthNet, EuroSAT, QuakeSet, RGBBandsMissingError
 from torchgeo.main import main
 from torchgeo.models import ResNet18_Weights
 from torchgeo.trainers import ClassificationTask, MultiLabelClassificationTask
@@ -41,12 +42,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
-class PredictClassificationDataModule(EuroSATDataModule):
+class PredictBinaryDataModule(QuakeSetDataModule):
+    def setup(self, stage: str) -> None:
+        self.predict_dataset = QuakeSet(split='test', **self.kwargs)
+
+
+class PredictMulticlassDataModule(EuroSATDataModule):
     def setup(self, stage: str) -> None:
         self.predict_dataset = EuroSAT(split='test', **self.kwargs)
 
 
-class PredictMultiLabelClassificationDataModule(BigEarthNetDataModule):
+class PredictMultilabelDataModule(BigEarthNetDataModule):
     def setup(self, stage: str) -> None:
         self.predict_dataset = BigEarthNet(split='test', **self.kwargs)
 
@@ -67,6 +73,9 @@ class TestClassificationTask:
     @pytest.mark.parametrize(
         'name',
         [
+            'bigearthnet_all',
+            'bigearthnet_s1',
+            'bigearthnet_s2',
             'eurosat',
             'eurosat100',
             'eurosatspatial',
@@ -77,6 +86,7 @@ class TestClassificationTask:
             'so2sat_s1',
             'so2sat_s2',
             'so2sat_rgb',
+            'treesatai',
             'ucmerced',
         ],
     )
@@ -178,11 +188,6 @@ def test_weight_str_download(self, weights: WeightsEnum) -> None:
             num_classes=10,
         )
 
-    def test_invalid_loss(self) -> None:
-        match = "Loss type 'invalid_loss' is not valid."
-        with pytest.raises(ValueError, match=match):
-            ClassificationTask(model='resnet18', loss='invalid_loss')
-
     def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
         monkeypatch.setattr(EuroSATDataModule, 'plot', plot)
         datamodule = EuroSATDataModule(
@@ -211,11 +216,14 @@ def test_no_rgb(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
         )
         trainer.validate(model=model, datamodule=datamodule)
 
-    def test_predict(self, fast_dev_run: bool) -> None:
-        datamodule = PredictClassificationDataModule(
-            root='tests/data/eurosat', batch_size=1, num_workers=0
+    def test_binary_predict(self, fast_dev_run: bool) -> None:
+        pytest.importorskip('h5py', minversion='3.6')
+        datamodule = PredictBinaryDataModule(
+            root='tests/data/quakeset', batch_size=1, num_workers=0
+        )
+        model = ClassificationTask(
+            model='resnet18', in_channels=4, task='binary', loss='bce'
         )
-        model = ClassificationTask(model='resnet18', in_channels=13, num_classes=10)
         trainer = Trainer(
             accelerator='cpu',
             fast_dev_run=fast_dev_run,
@@ -224,65 +232,16 @@ def test_predict(self, fast_dev_run: bool) -> None:
         )
         trainer.predict(model=model, datamodule=datamodule)
 
-    @pytest.mark.parametrize(
-        'model_name', ['resnet18', 'efficientnetv2_s', 'vit_base_patch16_384']
-    )
-    def test_freeze_backbone(self, model_name: str) -> None:
-        model = ClassificationTask(model=model_name, freeze_backbone=True)
-        assert not all([param.requires_grad for param in model.model.parameters()])
-        assert all(
-            [param.requires_grad for param in model.model.get_classifier().parameters()]
-        )
-
-
-class TestMultiLabelClassificationTask:
-    @pytest.mark.parametrize(
-        'name', ['bigearthnet_all', 'bigearthnet_s1', 'bigearthnet_s2', 'treesatai']
-    )
-    def test_trainer(
-        self, monkeypatch: MonkeyPatch, name: str, fast_dev_run: bool
-    ) -> None:
-        if name == 'bigearthnet_s2':
-            pytest.importorskip('pyarrow')
-        config = os.path.join('tests', 'conf', name + '.yaml')
-
-        monkeypatch.setattr(timm, 'create_model', create_model)
-
-        args = [
-            '--config',
-            config,
-            '--trainer.accelerator',
-            'cpu',
-            '--trainer.fast_dev_run',
-            str(fast_dev_run),
-            '--trainer.max_epochs',
-            '1',
-            '--trainer.log_every_n_steps',
-            '1',
-        ]
-
-        main(['fit', *args])
-        try:
-            main(['test', *args])
-        except MisconfigurationException:
-            pass
-        try:
-            main(['predict', *args])
-        except MisconfigurationException:
-            pass
-
-    def test_invalid_loss(self) -> None:
-        match = "Loss type 'invalid_loss' is not valid."
-        with pytest.raises(ValueError, match=match):
-            MultiLabelClassificationTask(model='resnet18', loss='invalid_loss')
-
-    def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
-        monkeypatch.setattr(BigEarthNetDataModule, 'plot', plot)
-        datamodule = BigEarthNetDataModule(
-            root='tests/data/bigearthnet/v1', batch_size=1, num_workers=0
+    def test_multiclass_predict(self, fast_dev_run: bool) -> None:
+        datamodule = PredictMulticlassDataModule(
+            root='tests/data/eurosat', batch_size=1, num_workers=0
         )
-        model = MultiLabelClassificationTask(
-            model='resnet18', in_channels=14, num_classes=19, loss='bce'
+        model = ClassificationTask(
+            model='resnet18',
+            in_channels=13,
+            task='multiclass',
+            num_classes=10,
+            loss='ce',
         )
         trainer = Trainer(
             accelerator='cpu',
@@ -290,15 +249,18 @@ def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> N
             log_every_n_steps=1,
             max_epochs=1,
         )
-        trainer.validate(model=model, datamodule=datamodule)
+        trainer.predict(model=model, datamodule=datamodule)
 
-    def test_no_rgb(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
-        monkeypatch.setattr(BigEarthNetDataModule, 'plot', plot_missing_bands)
-        datamodule = BigEarthNetDataModule(
+    def test_multilabel_predict(self, fast_dev_run: bool) -> None:
+        datamodule = PredictMultilabelDataModule(
             root='tests/data/bigearthnet/v1', batch_size=1, num_workers=0
         )
-        model = MultiLabelClassificationTask(
-            model='resnet18', in_channels=14, num_classes=19, loss='bce'
+        model = ClassificationTask(
+            model='resnet18',
+            in_channels=14,
+            task='multilabel',
+            num_labels=19,
+            loss='bce',
         )
         trainer = Trainer(
             accelerator='cpu',
@@ -306,19 +268,36 @@ def test_no_rgb(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
             log_every_n_steps=1,
             max_epochs=1,
         )
-        trainer.validate(model=model, datamodule=datamodule)
+        trainer.predict(model=model, datamodule=datamodule)
 
-    def test_predict(self, fast_dev_run: bool) -> None:
-        datamodule = PredictMultiLabelClassificationDataModule(
-            root='tests/data/bigearthnet/v1', batch_size=1, num_workers=0
+    @pytest.mark.parametrize(
+        'model_name', ['resnet18', 'efficientnetv2_s', 'vit_base_patch16_384']
+    )
+    def test_freeze_backbone(self, model_name: str) -> None:
+        model = ClassificationTask(
+            model=model_name, num_classes=10, freeze_backbone=True
         )
-        model = MultiLabelClassificationTask(
-            model='resnet18', in_channels=14, num_classes=19, loss='bce'
+        assert not all([param.requires_grad for param in model.model.parameters()])
+        assert all(
+            [param.requires_grad for param in model.model.get_classifier().parameters()]
         )
+
+
+class TestMultiLabelClassificationTask:
+    def test_trainer(self, fast_dev_run: bool) -> None:
+        datamodule = BigEarthNetDataModule(
+            root='tests/data/bigearthnet/v1', batch_size=1, num_workers=0
+        )
+
+        with pytest.deprecated_call():
+            model = MultiLabelClassificationTask(
+                model='resnet18', in_channels=14, num_classes=19, loss='bce'
+            )
         trainer = Trainer(
             accelerator='cpu',
             fast_dev_run=fast_dev_run,
             log_every_n_steps=1,
             max_epochs=1,
         )
-        trainer.predict(model=model, datamodule=datamodule)
+        trainer.fit(model=model, datamodule=datamodule)
+        trainer.test(model=model, datamodule=datamodule)
diff --git a/tests/trainers/test_segmentation.py b/tests/trainers/test_segmentation.py
index 4c56e47652d..e8547c93267 100644
--- a/tests/trainers/test_segmentation.py
+++ b/tests/trainers/test_segmentation.py
@@ -3,7 +3,7 @@
 
 import os
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, Literal, cast
 
 import pytest
 import segmentation_models_pytorch as smp
@@ -141,7 +141,7 @@ def mocked_weights(
     ) -> WeightsEnum:
         path = tmp_path / f'{weights}.pth'
         model = timm.create_model(
-            weights.meta['model'], in_chans=weights.meta['in_chans']
+            weights.meta['model'], in_chans=weights.meta['in_chans'], num_classes=10
         )
         torch.save(model.state_dict(), path)
         try:
@@ -158,6 +158,7 @@ def test_weight_enum(self, mocked_weights: WeightsEnum) -> None:
             backbone=mocked_weights.meta['model'],
             weights=mocked_weights,
             in_channels=mocked_weights.meta['in_chans'],
+            num_classes=10,
         )
 
     def test_weight_str(self, mocked_weights: WeightsEnum) -> None:
@@ -165,6 +166,7 @@ def test_weight_str(self, mocked_weights: WeightsEnum) -> None:
             backbone=mocked_weights.meta['model'],
             weights=str(mocked_weights),
             in_channels=mocked_weights.meta['in_chans'],
+            num_classes=10,
         )
 
     @pytest.mark.slow
@@ -183,16 +185,6 @@ def test_weight_str_download(self, weights: WeightsEnum) -> None:
             in_channels=weights.meta['in_chans'],
         )
 
-    def test_invalid_model(self) -> None:
-        match = "Model type 'invalid_model' is not valid."
-        with pytest.raises(ValueError, match=match):
-            SemanticSegmentationTask(model='invalid_model')
-
-    def test_invalid_loss(self) -> None:
-        match = "Loss type 'invalid_loss' is not valid."
-        with pytest.raises(ValueError, match=match):
-            SemanticSegmentationTask(loss='invalid_loss')
-
     def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
         monkeypatch.setattr(SEN12MSDataModule, 'plot', plot)
         datamodule = SEN12MSDataModule(
@@ -229,9 +221,11 @@ def test_no_rgb(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
     @pytest.mark.parametrize(
         'backbone', ['resnet18', 'mobilenet_v2', 'efficientnet-b0']
     )
-    def test_freeze_backbone(self, model_name: str, backbone: str) -> None:
+    def test_freeze_backbone(
+        self, model_name: Literal['unet', 'deeplabv3+'], backbone: str
+    ) -> None:
         model = SemanticSegmentationTask(
-            model=model_name, backbone=backbone, freeze_backbone=True
+            model=model_name, backbone=backbone, num_classes=10, freeze_backbone=True
         )
         assert all(
             [param.requires_grad is False for param in model.model.encoder.parameters()]
@@ -245,8 +239,10 @@ def test_freeze_backbone(self, model_name: str, backbone: str) -> None:
         )
 
     @pytest.mark.parametrize('model_name', ['unet', 'deeplabv3+'])
-    def test_freeze_decoder(self, model_name: str) -> None:
-        model = SemanticSegmentationTask(model=model_name, freeze_decoder=True)
+    def test_freeze_decoder(self, model_name: Literal['unet', 'deeplabv3+']) -> None:
+        model = SemanticSegmentationTask(
+            model=model_name, num_classes=10, freeze_decoder=True
+        )
         assert all(
             [param.requires_grad is False for param in model.model.decoder.parameters()]
         )
diff --git a/torchgeo/datasets/treesatai.py b/torchgeo/datasets/treesatai.py
index 5f55f158361..260171796a1 100644
--- a/torchgeo/datasets/treesatai.py
+++ b/torchgeo/datasets/treesatai.py
@@ -184,7 +184,7 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
             Data and label at that index.
         """
         file = self.files[index]
-        label = torch.zeros(len(self.classes))
+        label = torch.zeros(len(self.classes), dtype=torch.long)
         for genus, _ in self.labels[file]:
             i = self.classes.index(genus)
             label[i] = 1
diff --git a/torchgeo/trainers/classification.py b/torchgeo/trainers/classification.py
index 28df27d1e96..c25d518680f 100644
--- a/torchgeo/trainers/classification.py
+++ b/torchgeo/trainers/classification.py
@@ -4,25 +4,19 @@
 """Trainers for image classification."""
 
 import os
-from typing import Any
+from typing import Any, Literal
 
 import kornia.augmentation as K
 import matplotlib.pyplot as plt
 import timm
-import torch
 import torch.nn as nn
 from matplotlib.figure import Figure
 from segmentation_models_pytorch.losses import FocalLoss, JaccardLoss
 from torch import Tensor
 from torchmetrics import MetricCollection
-from torchmetrics.classification import (
-    MulticlassAccuracy,
-    MulticlassFBetaScore,
-    MulticlassJaccardIndex,
-    MultilabelAccuracy,
-    MultilabelFBetaScore,
-)
+from torchmetrics.classification import Accuracy, FBetaScore, JaccardIndex
 from torchvision.models._api import WeightsEnum
+from typing_extensions import deprecated
 
 from ..datasets import RGBBandsMissingError, unbind_samples
 from ..models import get_weight
@@ -38,8 +32,10 @@ def __init__(
         model: str = 'resnet50',
         weights: WeightsEnum | str | bool | None = None,
         in_channels: int = 3,
-        num_classes: int = 1000,
-        loss: str = 'ce',
+        task: Literal['binary', 'multiclass', 'multilabel'] = 'multiclass',
+        num_classes: int | None = None,
+        num_labels: int | None = None,
+        loss: Literal['ce', 'bce', 'jaccard', 'focal'] = 'ce',
         class_weights: Tensor | None = None,
         lr: float = 1e-3,
         patience: int = 10,
@@ -54,7 +50,9 @@ def __init__(
                 representation of a weight enum, True for ImageNet weights, False
                 or None for random weights, or the path to a saved model state dict.
             in_channels: Number of input channels to model.
-            num_classes: Number of prediction classes.
+            task: One of 'binary', 'multiclass', or 'multilabel'.
+            num_classes: Number of prediction classes (only for ``task='multiclass'``).
+            num_labels: Number of prediction labels (only for ``task='multilabel'``).
             loss: One of 'ce', 'bce', 'jaccard', or 'focal'.
             class_weights: Optional rescaling weight given to each
                 class and used with 'ce' loss.
@@ -63,8 +61,8 @@ class and used with 'ce' loss.
             freeze_backbone: Freeze the backbone network to linear probe
                 the classifier head.
 
-        .. versionchanged:: 0.4
-           *classification_model* was renamed to *model*.
+        .. versionadded:: 0.7
+           The *task* and *num_labels* parameters.
 
         .. versionadded:: 0.5
            The *class_weights* and *freeze_backbone* parameters.
@@ -72,6 +70,9 @@ class and used with 'ce' loss.
         .. versionchanged:: 0.5
            *learning_rate* and *learning_rate_schedule_patience* were renamed to
            *lr* and *patience*.
+
+        .. versionchanged:: 0.4
+           *classification_model* was renamed to *model*.
         """
         self.weights = weights
         super().__init__()
@@ -83,7 +84,7 @@ def configure_models(self) -> None:
         # Create model
         self.model = timm.create_model(
             self.hparams['model'],
-            num_classes=self.hparams['num_classes'],
+            num_classes=self.hparams['num_classes'] or self.hparams['num_labels'] or 1,
             in_chans=self.hparams['in_channels'],
             pretrained=weights is True,
         )
@@ -106,35 +107,29 @@ def configure_models(self) -> None:
                 param.requires_grad = True
 
     def configure_losses(self) -> None:
-        """Initialize the loss criterion.
-
-        Raises:
-            ValueError: If *loss* is invalid.
-        """
-        loss: str = self.hparams['loss']
-        if loss == 'ce':
-            self.criterion: nn.Module = nn.CrossEntropyLoss(
-                weight=self.hparams['class_weights']
-            )
-        elif loss == 'bce':
-            self.criterion = nn.BCEWithLogitsLoss()
-        elif loss == 'jaccard':
-            self.criterion = JaccardLoss(mode='multiclass')
-        elif loss == 'focal':
-            self.criterion = FocalLoss(mode='multiclass', normalized=True)
-        else:
-            raise ValueError(f"Loss type '{loss}' is not valid.")
+        """Initialize the loss criterion."""
+        match self.hparams['loss']:
+            case 'ce':
+                self.criterion: nn.Module = nn.CrossEntropyLoss(
+                    weight=self.hparams['class_weights']
+                )
+            case 'bce':
+                self.criterion = nn.BCEWithLogitsLoss()
+            case 'jaccard':
+                self.criterion = JaccardLoss(mode=self.hparams['task'])
+            case 'focal':
+                self.criterion = FocalLoss(mode=self.hparams['task'], normalized=True)
 
     def configure_metrics(self) -> None:
         """Initialize the performance metrics.
 
-        * :class:`~torchmetrics.classification.MulticlassAccuracy`: The number of
+        * :class:`~torchmetrics.Accuracy`: The number of
           true positives divided by the dataset size. Both overall accuracy (OA)
           using 'micro' averaging and average accuracy (AA) using 'macro' averaging
           are reported. Higher values are better.
-        * :class:`~torchmetrics.classification.MulticlassJaccardIndex`: Intersection
+        * :class:`~torchmetrics.JaccardIndex`: Intersection
           over union (IoU). Uses 'macro' averaging. Higher valuers are better.
-        * :class:`~torchmetrics.classification.MulticlassFBetaScore`: F1 score.
+        * :class:`~torchmetrics.FBetaScore`: F1 score.
           The harmonic mean of precision and recall. Uses 'micro' averaging.
           Higher values are better.
 
@@ -144,20 +139,17 @@ def configure_metrics(self) -> None:
            * 'Macro' averaging gives equal weight to each class, and is useful for
              balanced performance assessment across imbalanced classes.
         """
+        kwargs = {
+            'task': self.hparams['task'],
+            'num_classes': self.hparams['num_classes'],
+            'num_labels': self.hparams['num_labels'],
+        }
         metrics = MetricCollection(
             {
-                'OverallAccuracy': MulticlassAccuracy(
-                    num_classes=self.hparams['num_classes'], average='micro'
-                ),
-                'AverageAccuracy': MulticlassAccuracy(
-                    num_classes=self.hparams['num_classes'], average='macro'
-                ),
-                'JaccardIndex': MulticlassJaccardIndex(
-                    num_classes=self.hparams['num_classes']
-                ),
-                'F1Score': MulticlassFBetaScore(
-                    num_classes=self.hparams['num_classes'], beta=1.0, average='micro'
-                ),
+                'OverallAccuracy': Accuracy(average='micro', **kwargs),
+                'AverageAccuracy': Accuracy(average='macro', **kwargs),
+                'JaccardIndex': JaccardIndex(**kwargs),
+                'F1Score': FBetaScore(beta=1.0, average='micro', **kwargs),
             }
         )
         self.train_metrics = metrics.clone(prefix='train_')
@@ -180,12 +172,16 @@ def training_step(
         x = batch['image']
         y = batch['label']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss: Tensor = self.criterion(y_hat, y)
-        self.log('train_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.train_metrics(y_hat, y)
         self.log_dict(self.train_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
+
+        loss: Tensor = self.criterion(y_hat, y)
+        self.log('train_loss', loss, batch_size=batch_size)
+
         return loss
 
     def validation_step(
@@ -201,12 +197,16 @@ def validation_step(
         x = batch['image']
         y = batch['label']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss = self.criterion(y_hat, y)
-        self.log('val_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.val_metrics(y_hat, y)
         self.log_dict(self.val_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
+
+        loss = self.criterion(y_hat, y)
+        self.log('val_loss', loss, batch_size=batch_size)
+
         if (
             batch_idx < 10
             and hasattr(self.trainer, 'datamodule')
@@ -222,7 +222,12 @@ def validation_step(
                 keepdim=True,
             )
             batch = aug(batch)
-            batch['prediction'] = y_hat.argmax(dim=-1)
+            match self.hparams['task']:
+                case 'binary' | 'multilabel':
+                    batch['prediction'] = (y_hat >= 0.5).long()
+                case 'multiclass':
+                    batch['prediction'] = y_hat.argmax(dim=1)
+
             for key in ['image', 'label', 'prediction']:
                 batch[key] = batch[key].cpu()
             sample = unbind_samples(batch)[0]
@@ -251,12 +256,16 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         x = batch['image']
         y = batch['label']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss = self.criterion(y_hat, y)
-        self.log('test_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.test_metrics(y_hat, y)
         self.log_dict(self.test_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
+
+        loss = self.criterion(y_hat, y)
+        self.log('test_loss', loss, batch_size=batch_size)
+
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
     ) -> Tensor:
@@ -271,155 +280,23 @@ def predict_step(
             Output predicted probabilities.
         """
         x = batch['image']
-        y_hat: Tensor = self(x).softmax(dim=-1)
+        y_hat: Tensor = self(x)
+
+        match self.hparams['task']:
+            case 'binary' | 'multilabel':
+                y_hat = y_hat.sigmoid()
+            case 'multiclass':
+                y_hat = y_hat.softmax(dim=1)
+
         return y_hat
 
 
+@deprecated('Use torchgeo.trainers.ClassificationTask instead')
 class MultiLabelClassificationTask(ClassificationTask):
     """Multi-label image classification."""
 
-    def configure_metrics(self) -> None:
-        """Initialize the performance metrics.
-
-        * :class:`~torchmetrics.classification.MultilabelAccuracy`: The number of
-          true positives divided by the dataset size. Both overall accuracy (OA)
-          using 'micro' averaging and average accuracy (AA) using 'macro' averaging
-          are reported. Higher values are better.
-        * :class:`~torchmetrics.classification.MultilabelFBetaScore`: F1 score.
-          The harmonic mean of precision and recall. Uses 'micro' averaging.
-          Higher values are better.
-
-        .. note::
-           * 'Micro' averaging suits overall performance evaluation but may not
-             reflect minority class accuracy.
-           * 'Macro' averaging gives equal weight to each class, and is useful for
-             balanced performance assessment across imbalanced classes.
-        """
-        metrics = MetricCollection(
-            {
-                'OverallAccuracy': MultilabelAccuracy(
-                    num_labels=self.hparams['num_classes'], average='micro'
-                ),
-                'AverageAccuracy': MultilabelAccuracy(
-                    num_labels=self.hparams['num_classes'], average='macro'
-                ),
-                'F1Score': MultilabelFBetaScore(
-                    num_labels=self.hparams['num_classes'], beta=1.0, average='micro'
-                ),
-            }
-        )
-        self.train_metrics = metrics.clone(prefix='train_')
-        self.val_metrics = metrics.clone(prefix='val_')
-        self.test_metrics = metrics.clone(prefix='test_')
-
-    def training_step(
-        self, batch: Any, batch_idx: int, dataloader_idx: int = 0
-    ) -> Tensor:
-        """Compute the training loss and additional metrics.
-
-        Args:
-            batch: The output of your DataLoader.
-            batch_idx: Integer displaying index of this batch.
-            dataloader_idx: Index of the current dataloader.
-
-        Returns:
-            The loss tensor.
-        """
-        x = batch['image']
-        y = batch['label']
-        batch_size = x.shape[0]
-        y_hat = self(x)
-        y_hat_hard = torch.sigmoid(y_hat)
-        loss: Tensor = self.criterion(y_hat, y.to(torch.float))
-        self.log('train_loss', loss, batch_size=batch_size)
-        self.train_metrics(y_hat_hard, y)
-        self.log_dict(self.train_metrics)
-
-        return loss
-
-    def validation_step(
-        self, batch: Any, batch_idx: int, dataloader_idx: int = 0
-    ) -> None:
-        """Compute the validation loss and additional metrics.
-
-        Args:
-            batch: The output of your DataLoader.
-            batch_idx: Integer displaying index of this batch.
-            dataloader_idx: Index of the current dataloader.
-        """
-        x = batch['image']
-        y = batch['label']
-        batch_size = x.shape[0]
-        y_hat = self(x)
-        y_hat_hard = torch.sigmoid(y_hat)
-        loss = self.criterion(y_hat, y.to(torch.float))
-        self.log('val_loss', loss, batch_size=batch_size)
-        self.val_metrics(y_hat_hard, y)
-        self.log_dict(self.val_metrics, batch_size=batch_size)
-
-        if (
-            batch_idx < 10
-            and hasattr(self.trainer, 'datamodule')
-            and hasattr(self.trainer.datamodule, 'plot')
-            and self.logger
-            and hasattr(self.logger, 'experiment')
-            and hasattr(self.logger.experiment, 'add_figure')
-        ):
-            datamodule = self.trainer.datamodule
-            aug = K.AugmentationSequential(
-                K.Denormalize(datamodule.mean, datamodule.std),
-                data_keys=None,
-                keepdim=True,
-            )
-            batch = aug(batch)
-            batch['prediction'] = y_hat_hard
-            for key in ['image', 'label', 'prediction']:
-                batch[key] = batch[key].cpu()
-            sample = unbind_samples(batch)[0]
-
-            fig: Figure | None = None
-            try:
-                fig = datamodule.plot(sample)
-            except RGBBandsMissingError:
-                pass
-
-            if fig:
-                summary_writer = self.logger.experiment
-                summary_writer.add_figure(
-                    f'image/{batch_idx}', fig, global_step=self.global_step
-                )
-
-    def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
-        """Compute the test loss and additional metrics.
-
-        Args:
-            batch: The output of your DataLoader.
-            batch_idx: Integer displaying index of this batch.
-            dataloader_idx: Index of the current dataloader.
-        """
-        x = batch['image']
-        y = batch['label']
-        batch_size = x.shape[0]
-        y_hat = self(x)
-        y_hat_hard = torch.sigmoid(y_hat)
-        loss = self.criterion(y_hat, y.to(torch.float))
-        self.log('test_loss', loss, batch_size=batch_size)
-        self.test_metrics(y_hat_hard, y)
-        self.log_dict(self.test_metrics, batch_size=batch_size)
-
-    def predict_step(
-        self, batch: Any, batch_idx: int, dataloader_idx: int = 0
-    ) -> Tensor:
-        """Compute the predicted class probabilities.
-
-        Args:
-            batch: The output of your DataLoader.
-            batch_idx: Integer displaying index of this batch.
-            dataloader_idx: Index of the current dataloader.
-
-        Returns:
-            Output predicted probabilities.
-        """
-        x = batch['image']
-        y_hat = torch.sigmoid(self(x))
-        return y_hat
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Wrapper around torchgeo.trainers.ClassificationTask to massage kwargs."""
+        kwargs['task'] = 'multilabel'
+        kwargs['num_labels'] = kwargs['num_classes']
+        super().__init__(*args, **kwargs)
diff --git a/torchgeo/trainers/segmentation.py b/torchgeo/trainers/segmentation.py
index 954675692e0..e395c217630 100644
--- a/torchgeo/trainers/segmentation.py
+++ b/torchgeo/trainers/segmentation.py
@@ -4,7 +4,7 @@
 """Trainers for semantic segmentation."""
 
 import os
-from typing import Any
+from typing import Any, Literal
 
 import kornia.augmentation as K
 import matplotlib.pyplot as plt
@@ -12,8 +12,7 @@
 import torch.nn as nn
 from matplotlib.figure import Figure
 from torch import Tensor
-from torchmetrics import MetricCollection
-from torchmetrics.classification import MulticlassAccuracy, MulticlassJaccardIndex
+from torchmetrics import Accuracy, JaccardIndex, MetricCollection
 from torchvision.models._api import WeightsEnum
 
 from ..datasets import RGBBandsMissingError, unbind_samples
@@ -27,13 +26,15 @@ class SemanticSegmentationTask(BaseTask):
 
     def __init__(
         self,
-        model: str = 'unet',
+        model: Literal['unet', 'deeplabv3+', 'fcn'] = 'unet',
         backbone: str = 'resnet50',
         weights: WeightsEnum | str | bool | None = None,
         in_channels: int = 3,
-        num_classes: int = 1000,
+        task: Literal['binary', 'multiclass', 'multilabel'] = 'multiclass',
+        num_classes: int | None = None,
+        num_labels: int | None = None,
         num_filters: int = 3,
-        loss: str = 'ce',
+        loss: Literal['ce', 'bce', 'jaccard', 'focal'] = 'ce',
         class_weights: Tensor | None = None,
         ignore_index: int | None = None,
         lr: float = 1e-3,
@@ -55,10 +56,12 @@ def __init__(
                 model does not support pretrained weights. Pretrained ViT weight enums
                 are not supported yet.
             in_channels: Number of input channels to model.
-            num_classes: Number of prediction classes (including the background).
+            task: One of 'binary', 'multiclass', or 'multilabel'.
+            num_classes: Number of prediction classes (only for ``task='multiclass'``).
+            num_labels: Number of prediction labels (only for ``task='multilabel'``).
             num_filters: Number of filters. Only applicable when model='fcn'.
             loss: Name of the loss function, currently supports
-                'ce', 'jaccard' or 'focal' loss.
+                'ce', 'bce', 'jaccard', and 'focal' loss.
             class_weights: Optional rescaling weight given to each
                 class and used with 'ce' loss.
             ignore_index: Optional integer class index to ignore in the loss and
@@ -70,12 +73,11 @@ class and used with 'ce' loss.
             freeze_decoder: Freeze the decoder network to linear probe
                 the segmentation head.
 
-        .. versionchanged:: 0.3
-           *ignore_zeros* was renamed to *ignore_index*.
+        .. versionadded:: 0.7
+           The *task* and *num_labels* parameters.
 
-        .. versionchanged:: 0.4
-           *segmentation_model*, *encoder_name*, and *encoder_weights*
-           were renamed to *model*, *backbone*, and *weights*.
+        .. versionchanged:: 0.6
+           The *ignore_index* parameter now works for jaccard loss.
 
         .. versionadded:: 0.5
            The *class_weights*, *freeze_backbone*, and *freeze_decoder* parameters.
@@ -85,48 +87,48 @@ class and used with 'ce' loss.
            *learning_rate* and *learning_rate_schedule_patience* were renamed to
            *lr* and *patience*.
 
-        .. versionchanged:: 0.6
-           The *ignore_index* parameter now works for jaccard loss.
+        .. versionchanged:: 0.4
+           *segmentation_model*, *encoder_name*, and *encoder_weights*
+           were renamed to *model*, *backbone*, and *weights*.
+
+        .. versionchanged:: 0.3
+           *ignore_zeros* was renamed to *ignore_index*.
         """
         self.weights = weights
         super().__init__()
 
     def configure_models(self) -> None:
-        """Initialize the model.
-
-        Raises:
-            ValueError: If *model* is invalid.
-        """
+        """Initialize the model."""
         model: str = self.hparams['model']
         backbone: str = self.hparams['backbone']
         weights = self.weights
         in_channels: int = self.hparams['in_channels']
-        num_classes: int = self.hparams['num_classes']
+        num_classes: int = (
+            self.hparams['num_classes'] or self.hparams['num_labels'] or 1
+        )
         num_filters: int = self.hparams['num_filters']
 
-        if model == 'unet':
-            self.model = smp.Unet(
-                encoder_name=backbone,
-                encoder_weights='imagenet' if weights is True else None,
-                in_channels=in_channels,
-                classes=num_classes,
-            )
-        elif model == 'deeplabv3+':
-            self.model = smp.DeepLabV3Plus(
-                encoder_name=backbone,
-                encoder_weights='imagenet' if weights is True else None,
-                in_channels=in_channels,
-                classes=num_classes,
-            )
-        elif model == 'fcn':
-            self.model = FCN(
-                in_channels=in_channels, classes=num_classes, num_filters=num_filters
-            )
-        else:
-            raise ValueError(
-                f"Model type '{model}' is not valid. "
-                "Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
-            )
+        match model:
+            case 'unet':
+                self.model = smp.Unet(
+                    encoder_name=backbone,
+                    encoder_weights='imagenet' if weights is True else None,
+                    in_channels=in_channels,
+                    classes=num_classes,
+                )
+            case 'deeplabv3+':
+                self.model = smp.DeepLabV3Plus(
+                    encoder_name=backbone,
+                    encoder_weights='imagenet' if weights is True else None,
+                    in_channels=in_channels,
+                    classes=num_classes,
+                )
+            case 'fcn':
+                self.model = FCN(
+                    in_channels=in_channels,
+                    classes=num_classes,
+                    num_filters=num_filters,
+                )
 
         if model != 'fcn':
             if weights and weights is not True:
@@ -149,43 +151,38 @@ def configure_models(self) -> None:
                 param.requires_grad = False
 
     def configure_losses(self) -> None:
-        """Initialize the loss criterion.
-
-        Raises:
-            ValueError: If *loss* is invalid.
-        """
-        loss: str = self.hparams['loss']
-        ignore_index = self.hparams['ignore_index']
-        if loss == 'ce':
-            ignore_value = -1000 if ignore_index is None else ignore_index
-            self.criterion = nn.CrossEntropyLoss(
-                ignore_index=ignore_value, weight=self.hparams['class_weights']
-            )
-        elif loss == 'jaccard':
-            # 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
-            )
-        else:
-            raise ValueError(
-                f"Loss type '{loss}' is not valid. "
-                "Currently, supports 'ce', 'jaccard' or 'focal' loss."
-            )
+        """Initialize the loss criterion."""
+        ignore_index: int | None = self.hparams['ignore_index']
+        match self.hparams['loss']:
+            case 'ce':
+                ignore_value = -1000 if ignore_index is None else ignore_index
+                self.criterion: nn.Module = nn.CrossEntropyLoss(
+                    ignore_index=ignore_value, weight=self.hparams['class_weights']
+                )
+            case 'bce':
+                self.criterion = nn.BCEWithLogitsLoss()
+            case 'jaccard':
+                # 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
+                )
+            case 'focal':
+                self.criterion = smp.losses.FocalLoss(
+                    'multiclass', ignore_index=ignore_index, normalized=True
+                )
 
     def configure_metrics(self) -> None:
         """Initialize the performance metrics.
 
-        * :class:`~torchmetrics.classification.MulticlassAccuracy`: Overall accuracy
+        * :class:`~torchmetrics.Accuracy`: Overall accuracy
           (OA) using 'micro' averaging. The number of true positives divided by the
           dataset size. Higher values are better.
-        * :class:`~torchmetrics.classification.MulticlassJaccardIndex`: Intersection
+        * :class:`~torchmetrics.JaccardIndex`: Intersection
           over union (IoU). Uses 'micro' averaging. Higher valuers are better.
 
         .. note::
@@ -194,19 +191,16 @@ def configure_metrics(self) -> None:
            * 'Macro' averaging, not used here, gives equal weight to each class, useful
              for balanced performance assessment across imbalanced classes.
         """
-        num_classes: int = self.hparams['num_classes']
-        ignore_index: int | None = self.hparams['ignore_index']
+        kwargs = {
+            'task': self.hparams['task'],
+            'num_classes': self.hparams['num_classes'],
+            'num_labels': self.hparams['num_labels'],
+            'ignore_index': self.hparams['ignore_index'],
+        }
         metrics = MetricCollection(
             [
-                MulticlassAccuracy(
-                    num_classes=num_classes,
-                    ignore_index=ignore_index,
-                    multidim_average='global',
-                    average='micro',
-                ),
-                MulticlassJaccardIndex(
-                    num_classes=num_classes, ignore_index=ignore_index, average='micro'
-                ),
+                Accuracy(multidim_average='global', average='micro', **kwargs),
+                JaccardIndex(average='micro', **kwargs),
             ]
         )
         self.train_metrics = metrics.clone(prefix='train_')
@@ -229,11 +223,16 @@ def training_step(
         x = batch['image']
         y = batch['mask']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss: Tensor = self.criterion(y_hat, y)
-        self.log('train_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.train_metrics(y_hat, y)
         self.log_dict(self.train_metrics, batch_size=batch_size)
+
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
+
+        loss: Tensor = self.criterion(y_hat, y)
+        self.log('train_loss', loss, batch_size=batch_size)
+
         return loss
 
     def validation_step(
@@ -249,12 +248,16 @@ def validation_step(
         x = batch['image']
         y = batch['mask']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss = self.criterion(y_hat, y)
-        self.log('val_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.val_metrics(y_hat, y)
         self.log_dict(self.val_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
+
+        loss = self.criterion(y_hat, y)
+        self.log('val_loss', loss, batch_size=batch_size)
+
         if (
             batch_idx < 10
             and hasattr(self.trainer, 'datamodule')
@@ -270,7 +273,12 @@ def validation_step(
                 keepdim=True,
             )
             batch = aug(batch)
-            batch['prediction'] = y_hat.argmax(dim=1)
+            match self.hparams['task']:
+                case 'binary' | 'multilabel':
+                    batch['prediction'] = (y_hat >= 0.5).long()
+                case 'multiclass':
+                    batch['prediction'] = y_hat.argmax(dim=1)
+
             for key in ['image', 'mask', 'prediction']:
                 batch[key] = batch[key].cpu()
             sample = unbind_samples(batch)[0]
@@ -299,12 +307,16 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         x = batch['image']
         y = batch['mask']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss = self.criterion(y_hat, y)
-        self.log('test_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.test_metrics(y_hat, y)
         self.log_dict(self.test_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
+
+        loss = self.criterion(y_hat, y)
+        self.log('test_loss', loss, batch_size=batch_size)
+
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
     ) -> Tensor:
@@ -319,5 +331,12 @@ def predict_step(
             Output predicted probabilities.
         """
         x = batch['image']
-        y_hat: Tensor = self(x).softmax(dim=1)
+        y_hat: Tensor = self(x)
+
+        match self.hparams['task']:
+            case 'binary' | 'multilabel':
+                y_hat = y_hat.sigmoid()
+            case 'multiclass':
+                y_hat = y_hat.softmax(dim=1)
+
         return y_hat


Reviews:
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-08-12T15:10:20Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2024-08-12T15:11:46Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-11T14:56:41Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-11T14:58:30Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-11T14:59:04Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-11T14:59:27Z
user: login: nilsleh
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-11T15:23:47Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-11T15:26:45Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-16T15:22:48Z
user: login: calebrob6
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T15:13:26Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T16:24:04Z
user: login: isaaccorley
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-27T02:19:22Z
user: login: isaaccorley
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-27T02:33:44Z
user: login: isaaccorley
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-27T02:35:19Z
user: login: isaaccorley
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-27T02:36:02Z
user: login: adamjstewart
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-27T15:12:24Z

Review comments:
diff_hunk: @@ -37,7 +30,9 @@ def __init__(
         model: str = 'resnet50',
         weights: WeightsEnum | str | bool | None = None,
         in_channels: int = 3,
-        num_classes: int = 1000,
+        task: str = 'multiclass',
+        num_classes: int | None = None,
+        num_labels: int | None = None,
path: torchgeo/trainers/classification.py
user: login: adamjstewart
body: The defaults here match torchmetrics. If `task='multiclass'`, only `num_classes` is used. If `task='multilabel'`, only `num_labels` is used. If `task='binary'`, both are ignored. Honestly, we could have a single `num_classes` if we want and simply use it for both.
created_at: 2024-08-12T15:10:19Z
author_association: MEMBER
diff_hunk: @@ -266,147 +262,3 @@ def predict_step(
         x = batch['image']
         y_hat: Tensor = self(x).softmax(dim=-1)
         return y_hat
-
-
-class MultiLabelClassificationTask(ClassificationTask):
path: torchgeo/trainers/classification.py
user: login: adamjstewart
body: The _right_ thing to do would be to deprecate this first and remove it in 0.7.0. Not sure how widely used it is. Deprecation is kind of annoying because you need to change all tests to acknowledge the warning message.
created_at: 2024-08-12T15:11:46Z
author_association: MEMBER
diff_hunk: @@ -224,101 +232,72 @@ def test_predict(self, fast_dev_run: bool) -> None:
         )
         trainer.predict(model=model, datamodule=datamodule)
 
-    @pytest.mark.parametrize(
-        'model_name', ['resnet18', 'efficientnetv2_s', 'vit_base_patch16_384']
-    )
-    def test_freeze_backbone(self, model_name: str) -> None:
-        model = ClassificationTask(model=model_name, freeze_backbone=True)
-        assert not all([param.requires_grad for param in model.model.parameters()])
-        assert all(
-            [param.requires_grad for param in model.model.get_classifier().parameters()]
-        )
-
-
-class TestMultiLabelClassificationTask:
-    @pytest.mark.parametrize(
-        'name', ['bigearthnet_all', 'bigearthnet_s1', 'bigearthnet_s2', 'treesatai']
-    )
-    def test_trainer(
-        self, monkeypatch: MonkeyPatch, name: str, fast_dev_run: bool
-    ) -> None:
-        if name == 'bigearthnet_s2':
-            pytest.importorskip('pyarrow')
path: tests/trainers/test_classification.py
user: login: adamjstewart
body: @nilsleh this was actually a typo, you meant bigearthnet v2 (not s2), but we don't test that here either.
created_at: 2025-03-11T14:56:41Z
author_association: MEMBER
diff_hunk: @@ -179,12 +171,16 @@ def training_step(
         x = batch['image']
         y = batch['label']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss: Tensor = self.criterion(y_hat, y)
-        self.log('train_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
path: torchgeo/trainers/classification.py
user: login: adamjstewart
body: `task='binary'` metrics require us to squash the 1 class dimension
created_at: 2025-03-11T14:58:29Z
author_association: MEMBER
diff_hunk: @@ -179,12 +171,16 @@ def training_step(
         x = batch['image']
         y = batch['label']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss: Tensor = self.criterion(y_hat, y)
-        self.log('train_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.train_metrics(y_hat, y)
         self.log_dict(self.train_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
path: torchgeo/trainers/classification.py
user: login: adamjstewart
body: `BCEWithLogitsLoss` requires float inputs for some reason...
created_at: 2025-03-11T14:59:04Z
author_association: MEMBER
diff_hunk: @@ -214,12 +214,17 @@ def validation_step(
             and hasattr(self.logger, 'experiment')
             and hasattr(self.logger.experiment, 'add_figure')
         ):
-            datamodule = self.trainer.datamodule
-            batch['prediction'] = y_hat.argmax(dim=-1)
+            match self.hparams['task']:
+                case 'binary' | 'multilabel':
+                    batch['prediction'] = (y_hat >= 0.5).long()
path: torchgeo/trainers/classification.py
user: login: adamjstewart
body: 0.5 is also the threshold used by torchmetrics
created_at: 2025-03-11T14:59:27Z
author_association: MEMBER
diff_hunk: @@ -214,12 +214,17 @@ def validation_step(
             and hasattr(self.logger, 'experiment')
             and hasattr(self.logger.experiment, 'add_figure')
         ):
-            datamodule = self.trainer.datamodule
-            batch['prediction'] = y_hat.argmax(dim=-1)
+            match self.hparams['task']:
+                case 'binary' | 'multilabel':
+                    batch['prediction'] = (y_hat >= 0.5).long()
path: torchgeo/trainers/classification.py
user: login: nilsleh
body: https://lightning.ai/docs/torchmetrics/stable/classification/accuracy.html#binaryaccuracy torchmetrics does this automatically, if you pass a float
created_at: 2025-03-11T15:23:47Z
author_association: COLLABORATOR
diff_hunk: @@ -214,12 +214,17 @@ def validation_step(
             and hasattr(self.logger, 'experiment')
             and hasattr(self.logger.experiment, 'add_figure')
         ):
-            datamodule = self.trainer.datamodule
-            batch['prediction'] = y_hat.argmax(dim=-1)
+            match self.hparams['task']:
+                case 'binary' | 'multilabel':
+                    batch['prediction'] = (y_hat >= 0.5).long()
path: torchgeo/trainers/classification.py
user: login: adamjstewart
body: Yep, this part is for plotting, not for metrics.
created_at: 2025-03-11T15:26:45Z
author_association: MEMBER
diff_hunk: @@ -148,43 +150,38 @@ def configure_models(self) -> None:
                 param.requires_grad = False
 
     def configure_losses(self) -> None:
-        """Initialize the loss criterion.
-
-        Raises:
-            ValueError: If *loss* is invalid.
-        """
-        loss: str = self.hparams['loss']
-        ignore_index = self.hparams['ignore_index']
-        if loss == 'ce':
-            ignore_value = -1000 if ignore_index is None else ignore_index
-            self.criterion = nn.CrossEntropyLoss(
-                ignore_index=ignore_value, weight=self.hparams['class_weights']
-            )
-        elif loss == 'jaccard':
-            # 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
-            )
-        else:
-            raise ValueError(
-                f"Loss type '{loss}' is not valid. "
-                "Currently, supports 'ce', 'jaccard' or 'focal' loss."
-            )
+        """Initialize the loss criterion."""
+        ignore_index: int | None = self.hparams['ignore_index']
+        match self.hparams['loss']:
+            case 'ce':
+                ignore_value = -1000 if ignore_index is None else ignore_index
+                self.criterion: nn.Module = nn.CrossEntropyLoss(
+                    ignore_index=ignore_value, weight=self.hparams['class_weights']
+                )
+            case 'bce':
+                self.criterion = nn.BCEWithLogitsLoss()
path: torchgeo/trainers/segmentation.py
user: login: adamjstewart
body: Is there any way we can modify this to support `ignore_index`? Maybe clever use of `pos_weight`? Or will values > 1 be automatically ignored?
created_at: 2025-03-16T15:22:48Z
author_association: MEMBER
diff_hunk: @@ -84,48 +86,48 @@ class and used with 'ce' loss.
            *learning_rate* and *learning_rate_schedule_patience* were renamed to
            *lr* and *patience*.
 
-        .. versionchanged:: 0.6
-           The *ignore_index* parameter now works for jaccard loss.
+        .. versionchanged:: 0.4
+           *segmentation_model*, *encoder_name*, and *encoder_weights*
+           were renamed to *model*, *backbone*, and *weights*.
+
+        .. versionchanged:: 0.3
+           *ignore_zeros* was renamed to *ignore_index*.
         """
         self.weights = weights
         super().__init__()
 
     def configure_models(self) -> None:
-        """Initialize the model.
-
-        Raises:
-            ValueError: If *model* is invalid.
-        """
+        """Initialize the model."""
         model: str = self.hparams['model']
         backbone: str = self.hparams['backbone']
         weights = self.weights
         in_channels: int = self.hparams['in_channels']
-        num_classes: int = self.hparams['num_classes']
+        num_classes: int = (
+            self.hparams['num_classes'] or self.hparams['num_labels'] or 1
+        )
         num_filters: int = self.hparams['num_filters']
 
-        if model == 'unet':
-            self.model = smp.Unet(
-                encoder_name=backbone,
-                encoder_weights='imagenet' if weights is True else None,
-                in_channels=in_channels,
-                classes=num_classes,
-            )
-        elif model == 'deeplabv3+':
-            self.model = smp.DeepLabV3Plus(
-                encoder_name=backbone,
-                encoder_weights='imagenet' if weights is True else None,
-                in_channels=in_channels,
-                classes=num_classes,
-            )
-        elif model == 'fcn':
-            self.model = FCN(
-                in_channels=in_channels, classes=num_classes, num_filters=num_filters
-            )
-        else:
-            raise ValueError(
-                f"Model type '{model}' is not valid. "
-                "Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
-            )
+        match model:
path: torchgeo/trainers/segmentation.py
user: login: calebrob6
body: Do we not throw a ValueError anymore?
created_at: 2025-03-19T15:13:25Z
author_association: MEMBER
diff_hunk: @@ -84,48 +86,48 @@ class and used with 'ce' loss.
            *learning_rate* and *learning_rate_schedule_patience* were renamed to
            *lr* and *patience*.
 
-        .. versionchanged:: 0.6
-           The *ignore_index* parameter now works for jaccard loss.
+        .. versionchanged:: 0.4
+           *segmentation_model*, *encoder_name*, and *encoder_weights*
+           were renamed to *model*, *backbone*, and *weights*.
+
+        .. versionchanged:: 0.3
+           *ignore_zeros* was renamed to *ignore_index*.
         """
         self.weights = weights
         super().__init__()
 
     def configure_models(self) -> None:
-        """Initialize the model.
-
-        Raises:
-            ValueError: If *model* is invalid.
-        """
+        """Initialize the model."""
         model: str = self.hparams['model']
         backbone: str = self.hparams['backbone']
         weights = self.weights
         in_channels: int = self.hparams['in_channels']
-        num_classes: int = self.hparams['num_classes']
+        num_classes: int = (
+            self.hparams['num_classes'] or self.hparams['num_labels'] or 1
+        )
         num_filters: int = self.hparams['num_filters']
 
-        if model == 'unet':
-            self.model = smp.Unet(
-                encoder_name=backbone,
-                encoder_weights='imagenet' if weights is True else None,
-                in_channels=in_channels,
-                classes=num_classes,
-            )
-        elif model == 'deeplabv3+':
-            self.model = smp.DeepLabV3Plus(
-                encoder_name=backbone,
-                encoder_weights='imagenet' if weights is True else None,
-                in_channels=in_channels,
-                classes=num_classes,
-            )
-        elif model == 'fcn':
-            self.model = FCN(
-                in_channels=in_channels, classes=num_classes, num_filters=num_filters
-            )
-        else:
-            raise ValueError(
-                f"Model type '{model}' is not valid. "
-                "Currently, only supports 'unet', 'deeplabv3+' and 'fcn'."
-            )
+        match model:
path: torchgeo/trainers/segmentation.py
user: login: adamjstewart
body: We don't. I want to use `Literal` type hints to more effectively document accepted values, although this makes it a little harder to raise error messages appropriately. We could still test it with invalid inputs and tell mypy to ignore the incorrect inputs if we really want better error messages here. I'm on the fence about this.
created_at: 2025-03-19T16:24:04Z
author_association: MEMBER
diff_hunk: @@ -148,43 +150,38 @@ def configure_models(self) -> None:
                 param.requires_grad = False
 
     def configure_losses(self) -> None:
-        """Initialize the loss criterion.
-
-        Raises:
-            ValueError: If *loss* is invalid.
-        """
-        loss: str = self.hparams['loss']
-        ignore_index = self.hparams['ignore_index']
-        if loss == 'ce':
-            ignore_value = -1000 if ignore_index is None else ignore_index
-            self.criterion = nn.CrossEntropyLoss(
-                ignore_index=ignore_value, weight=self.hparams['class_weights']
-            )
-        elif loss == 'jaccard':
-            # 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
-            )
-        else:
-            raise ValueError(
-                f"Loss type '{loss}' is not valid. "
-                "Currently, supports 'ce', 'jaccard' or 'focal' loss."
-            )
+        """Initialize the loss criterion."""
+        ignore_index: int | None = self.hparams['ignore_index']
+        match self.hparams['loss']:
+            case 'ce':
+                ignore_value = -1000 if ignore_index is None else ignore_index
+                self.criterion: nn.Module = nn.CrossEntropyLoss(
+                    ignore_index=ignore_value, weight=self.hparams['class_weights']
+                )
+            case 'bce':
+                self.criterion = nn.BCEWithLogitsLoss()
path: torchgeo/trainers/segmentation.py
user: login: isaaccorley
body: Probably would need to write a simple wrapper to do the masking of the loss. Something like this should work

```python
class BCE(nn.BCEWithLogitsLoss):
   def __init__(self, ignore_index, *args, **kwargs):
      super().__init__(*args, **kwargs)
      self.ignore_index = ignore_index
   
   def forward(self, input, target):
      ignore_mask = target == ignore_index
      loss = F.binary_cross_entropy_with_logits(
         input,
         target,
         self.weight,
         pos_weight=self.pos_weight,
         reduction='none'
      )
      loss = loss[~ignore_mask]

      if self.reduction = 'mean':
         loss = loss.mean()
      elif self.reduction = 'sum':
         loss = loss.sum()
      else:
         pass

      return loss
```
   
created_at: 2025-03-27T02:19:21Z
author_association: COLLABORATOR
diff_hunk: @@ -37,7 +30,9 @@ def __init__(
         model: str = 'resnet50',
         weights: WeightsEnum | str | bool | None = None,
         in_channels: int = 3,
-        num_classes: int = 1000,
+        task: str = 'multiclass',
+        num_classes: int | None = None,
+        num_labels: int | None = None,
path: torchgeo/trainers/classification.py
user: login: isaaccorley
body: I think it's fine to be explicit
created_at: 2025-03-27T02:33:44Z
author_association: COLLABORATOR
diff_hunk: @@ -179,12 +171,16 @@ def training_step(
         x = batch['image']
         y = batch['label']
         batch_size = x.shape[0]
-        y_hat = self(x)
-        loss: Tensor = self.criterion(y_hat, y)
-        self.log('train_loss', loss, batch_size=batch_size)
+        y_hat = self(x).squeeze(1)
         self.train_metrics(y_hat, y)
         self.log_dict(self.train_metrics, batch_size=batch_size)
 
+        if self.hparams['loss'] == 'bce':
+            y = y.float()
path: torchgeo/trainers/classification.py
user: login: isaaccorley
body: yup it's for the case that you targets are probabilities like [0.2, 0.8, 0.1] etc. and not just ints [0, 1, 0]. This is common for example if you use methods like label smoothing
created_at: 2025-03-27T02:35:19Z
author_association: COLLABORATOR
diff_hunk: @@ -148,43 +150,38 @@ def configure_models(self) -> None:
                 param.requires_grad = False
 
     def configure_losses(self) -> None:
-        """Initialize the loss criterion.
-
-        Raises:
-            ValueError: If *loss* is invalid.
-        """
-        loss: str = self.hparams['loss']
-        ignore_index = self.hparams['ignore_index']
-        if loss == 'ce':
-            ignore_value = -1000 if ignore_index is None else ignore_index
-            self.criterion = nn.CrossEntropyLoss(
-                ignore_index=ignore_value, weight=self.hparams['class_weights']
-            )
-        elif loss == 'jaccard':
-            # 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
-            )
-        else:
-            raise ValueError(
-                f"Loss type '{loss}' is not valid. "
-                "Currently, supports 'ce', 'jaccard' or 'focal' loss."
-            )
+        """Initialize the loss criterion."""
+        ignore_index: int | None = self.hparams['ignore_index']
+        match self.hparams['loss']:
+            case 'ce':
+                ignore_value = -1000 if ignore_index is None else ignore_index
+                self.criterion: nn.Module = nn.CrossEntropyLoss(
+                    ignore_index=ignore_value, weight=self.hparams['class_weights']
+                )
+            case 'bce':
+                self.criterion = nn.BCEWithLogitsLoss()
path: torchgeo/trainers/segmentation.py
user: login: adamjstewart
body: This is uglier than I would like, would rather try to upstream this to PyTorch first.
created_at: 2025-03-27T15:12:23Z
author_association: MEMBER

PR comments:
user: login: adamjstewart
created_at: 2024-09-05T13:41:13Z
body: I think we first need to add a multilabel semantic segmentation dataset to properly test this.
author_association: MEMBER
user: login: adamjstewart
created_at: 2024-09-05T13:42:17Z
body: Alternatively, skip multilabel semantic segmentation and only support multilabel classification.
author_association: MEMBER
user: login: adamjstewart
created_at: 2025-03-13T13:17:07Z
body: Should we add some kind of documentation/warning/error that recommends BCE for binary/multilabel and non-BCE for multiclass? Or warns when trying to combine `num_classes` with `task='multilabel'`, `num_labels` with `task='multiclass'`, or either with `task='binary'`?
author_association: MEMBER
user: login: calebrob6
created_at: 2025-03-19T15:17:27Z
body: I like the change, however it makes me a _little_ nervous to accept this without confirming that this new class can actually train models under the three different modes (some semantic bug that the tests aren't catching). Have you tried this with a multilabel or binary problem?
author_association: MEMBER
user: login: adamjstewart
created_at: 2025-03-19T16:27:55Z
body: Agreed, I was hoping someone could do some quick benchmarking for me. My biggest concern is whether these loss functions and metrics agree on the threshold for binary classification (0/1 with 0.5 threshold vs. -1/1 with 0 threshold).

We have plenty of builtin binary classification/segmentation and a few multilabel classification datasets but no multilabel segmentation datasets (yet). I actually wanted to add one for multicropping (multiple crops grown in the same field during the same year) but there are no existing datasets.

I'll be busy working on Copernicus-Bench for the next few days so don't yet have time to benchmark these myself.
author_association: MEMBER