add dice loss, add tests for loss instantiation

Signed-off-by: Carlos Gomes <[email protected]>

Author: CarlosGomes98

src/terratorch/tasks/segmentation_tasks.py

@@ -152,9 +152,11 @@ def configure_losses(self) -> None:
                     f"Jaccard loss does not support ignore_index, but found non-None value of {ignore_index}."
                 )
                 raise RuntimeError(exception_message)
-            self.criterion = smp.losses.JaccardLoss(mode="multiclass", classes=self.hparams["num_classes"])
+            self.criterion = smp.losses.JaccardLoss(mode="multiclass")
         elif loss == "focal":
             self.criterion = smp.losses.FocalLoss("multiclass", ignore_index=ignore_index, normalized=True)
+        elif loss == "dice":
+            self.criterion = smp.losses.DiceLoss("multiclass", ignore_index=ignore_index)
         else:
             exception_message = f"Loss type '{loss}' is not valid. Currently, supports 'ce', 'jaccard' or 'focal' loss."
             raise ValueError(exception_message)

tests/test_prithvi_tasks.py

@@ -22,7 +22,8 @@ def model_input() -> torch.Tensor:
 
 @pytest.mark.parametrize("backbone", ["prithvi_vit_100", "prithvi_vit_300"])
 @pytest.mark.parametrize("decoder", ["FCNDecoder", "UperNetDecoder", "IdentityDecoder"])
-def test_create_segmentation_task(backbone, decoder, model_factory: PrithviModelFactory):
+@pytest.mark.parametrize("loss", ["ce", "jaccard", "focal", "dice"])
+def test_create_segmentation_task(backbone, decoder, loss, model_factory: PrithviModelFactory):
     SemanticSegmentationTask(
         {
             "backbone": backbone,
@@ -33,12 +34,14 @@ def test_create_segmentation_task(backbone, decoder, model_factory: PrithviModel
             "num_classes": NUM_CLASSES,
         },
         model_factory,
+        loss=loss
     )
 
 
 @pytest.mark.parametrize("backbone", ["prithvi_vit_100", "prithvi_vit_300"])
 @pytest.mark.parametrize("decoder", ["FCNDecoder", "UperNetDecoder", "IdentityDecoder"])
-def test_create_regression_task(backbone, decoder, model_factory: PrithviModelFactory):
+@pytest.mark.parametrize("loss", ["mae", "rmse", "huber"])
+def test_create_regression_task(backbone, decoder, loss, model_factory: PrithviModelFactory):
     PixelwiseRegressionTask(
         {
             "backbone": backbone,
@@ -48,12 +51,14 @@ def test_create_regression_task(backbone, decoder, model_factory: PrithviModelFa
             "pretrained": False,
         },
         model_factory,
+        loss=loss
     )
 
 
 @pytest.mark.parametrize("backbone", ["prithvi_vit_100", "prithvi_vit_300"])
 @pytest.mark.parametrize("decoder", ["FCNDecoder", "UperNetDecoder", "IdentityDecoder"])
-def test_create_classification_task(backbone, decoder, model_factory: PrithviModelFactory):
+@pytest.mark.parametrize("loss", ["ce", "bce", "jaccard", "focal"])
+def test_create_classification_task(backbone, decoder, loss, model_factory: PrithviModelFactory):
     ClassificationTask(
         {
             "backbone": backbone,
@@ -64,4 +69,5 @@ def test_create_classification_task(backbone, decoder, model_factory: PrithviMod
             "num_classes": NUM_CLASSES,
         },
         model_factory,
+        loss=loss
     )