Add focal loss support to SegmentationHead

rslearn/train/tasks/segmentation.py

@@ -7,7 +7,9 @@
 import numpy.typing as npt
 import torch
 import torchmetrics.classification
+from einops import rearrange
 from torchmetrics import Metric, MetricCollection
+from torchvision.ops import sigmoid_focal_loss
 
 from rslearn.models.component import FeatureMaps, Predictor
 from rslearn.train.metrics import ConfusionMatrixMetric
@@ -310,24 +312,37 @@ class SegmentationHead(Predictor):
 
     def __init__(
         self,
-        weights: list[float] | None = None,
+        class_weights: list[float] | None = None,
         dice_loss: bool = False,
+        focal_loss: bool = False,
+        focal_loss_alpha: float = 0.25,
+        focal_loss_gamma: float = 2.0,
+        loss_weights: tuple[float, float] = (1.0, 1.0),
         temperature: float = 1.0,
     ):
         """Initialize a new SegmentationTask.
 
         Args:
-            weights: weights for cross entropy loss (Tensor of size C)
-            dice_loss: weather to add dice loss to cross entropy
+            class_weights: weights for cross entropy loss (list of length C)
+            dice_loss: whether to add dice loss to cross entropy
+            focal_loss: whether to use focal loss instead of cross entropy
+            focal_loss_alpha: alpha parameter for focal loss (default 0.25)
+            focal_loss_gamma: gamma parameter for focal loss (default 2.0)
+            loss_weights: tuple of (cls_weight, dice_weight) to scale the
+                classification loss (cross-entropy or focal) and dice loss
             temperature: temperature scaling for softmax, does not affect the loss,
                 only the predictor outputs
         """
         super().__init__()
-        if weights is not None:
-            self.register_buffer("weights", torch.Tensor(weights))
+        if class_weights is not None:
+            self.register_buffer("class_weights", torch.Tensor(class_weights))
         else:
-            self.weights = None
+            self.class_weights = None
         self.dice_loss = dice_loss
+        self.focal_loss = focal_loss
+        self.focal_loss_alpha = focal_loss_alpha
+        self.focal_loss_gamma = focal_loss_gamma
+        self.loss_weights = loss_weights
         self.temperature = temperature
 
     def forward(
@@ -367,21 +382,42 @@ def forward(
             mask = torch.stack(
                 [target["valid"].get_hw_tensor() for target in targets], dim=0
             )
-            per_pixel_loss = torch.nn.functional.cross_entropy(
-                logits, labels, weight=self.weights, reduction="none"
-            )
+
+            if self.focal_loss:
+                # Convert labels to one-hot for focal loss
+                num_classes = logits.shape[1]
+                labels_one_hot = rearrange(
+                    torch.nn.functional.one_hot(labels, num_classes).float(),
+                    "b h w c -> b c h w",
+                )
+                per_pixel_loss = sigmoid_focal_loss(
+                    logits,
+                    labels_one_hot,
+                    alpha=self.focal_loss_alpha,
+                    gamma=self.focal_loss_gamma,
+                    reduction="none",
+                )
+                # Sum over classes dimension to get per-pixel loss
+                per_pixel_loss = per_pixel_loss.sum(dim=1)
+            else:
+                per_pixel_loss = torch.nn.functional.cross_entropy(
+                    logits, labels, weight=self.class_weights, reduction="none"
+                )
+
             mask_sum = torch.sum(mask)
             if mask_sum > 0:
                 # Compute average loss over valid pixels.
-                losses["cls"] = torch.sum(per_pixel_loss * mask) / torch.sum(mask)
+                cls_loss = torch.sum(per_pixel_loss * mask) / torch.sum(mask)
             else:
                 # If there are no valid pixels, we avoid dividing by zero and just let
                 # the summed mask loss be zero.
-                losses["cls"] = torch.sum(per_pixel_loss * mask)
+                cls_loss = torch.sum(per_pixel_loss * mask)
+            losses["cls"] = cls_loss * self.loss_weights[0]
+
             if self.dice_loss:
                 softmax_woT = torch.nn.functional.softmax(logits, dim=1)
                 dice_loss = DiceLoss()(softmax_woT, labels, mask)
-                losses["dice"] = dice_loss
+                losses["dice"] = dice_loss * self.loss_weights[1]
 
         return ModelOutput(
             outputs=outputs,
@@ -746,10 +782,9 @@ def forward(
             the mean Dicen Loss across classes
         """
         num_classes = inputs.shape[1]
-        targets_one_hot = (
-            torch.nn.functional.one_hot(targets, num_classes)
-            .permute(0, 3, 1, 2)
-            .float()
+        targets_one_hot = rearrange(
+            torch.nn.functional.one_hot(targets, num_classes).float(),
+            "b h w c -> b c h w",
         )
 
         # Expand mask to [B, C, H, W]

tests/unit/train/tasks/test_segmentation.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import torch
+from torchvision.ops import sigmoid_focal_loss
 
 from rslearn.models.component import FeatureMaps
 from rslearn.train.model_context import ModelContext, RasterImage, SampleMetadata
@@ -381,6 +382,101 @@ def test_segmentation_head_temperature_confidence() -> None:
     assert max_prob_cold > max_prob_hot
 
 
+class TestSegmentationHeadLoss:
+    """Tests for SegmentationHead loss computation (focal loss and loss_weights)."""
+
+    @staticmethod
+    def _make_inputs(
+        num_classes: int = 3,
+        h: int = 4,
+        w: int = 4,
+        mask_value: float = 1.0,
+    ) -> tuple:
+        """Create deterministic inputs for SegmentationHead.
+
+        Returns (FeatureMaps, ModelContext, targets list).
+        """
+        torch.manual_seed(42)
+        logits = torch.randn(1, num_classes, h, w)
+        labels = torch.randint(0, num_classes, (h, w))
+        targets = [
+            {
+                "classes": RasterImage(labels[None, None, :, :].long()),
+                "valid": RasterImage(
+                    torch.full((1, 1, h, w), mask_value, dtype=torch.float32)
+                ),
+            }
+        ]
+        return FeatureMaps([logits]), ModelContext(inputs=[], metadatas=[]), targets
+
+    def test_loss_weights_scale_losses(self) -> None:
+        """Verify loss_weights[0] proportionally scales the classification loss."""
+        fm, ctx, targets = self._make_inputs()
+
+        head_1x = SegmentationHead(loss_weights=(1.0, 1.0))
+        head_3x = SegmentationHead(loss_weights=(3.0, 1.0))
+        head_1x_dice = SegmentationHead(dice_loss=True, loss_weights=(1.0, 1.0))
+        head_3x_dice = SegmentationHead(dice_loss=True, loss_weights=(1.0, 3.0))
+
+        loss_1x = head_1x(fm, ctx, targets).loss_dict["cls"]
+        loss_3x = head_3x(fm, ctx, targets).loss_dict["cls"]
+        loss_1x_dice = head_1x_dice(fm, ctx, targets).loss_dict["dice"]
+        loss_3x_dice = head_3x_dice(fm, ctx, targets).loss_dict["dice"]
+
+        assert loss_3x.item() == pytest.approx(3.0 * loss_1x.item(), rel=1e-5)
+        assert loss_3x_dice.item() == pytest.approx(3.0 * loss_1x_dice.item(), rel=1e-5)
+
+    def test_losses_match_manual_computation(self) -> None:
+        """Verify CE and focal loss match manually computed reference values."""
+        # 2-class, 2x2 image with one masked pixel.
+        logits = torch.tensor(
+            [[[[1.0, -1.0], [0.5, 2.0]], [[-1.0, 1.0], [0.5, -2.0]]]]
+        )  # [1, 2, 2, 2]
+        labels = torch.tensor([[0, 1], [1, 0]])  # [2, 2]
+        mask = torch.tensor([[1.0, 1.0], [1.0, 0.0]])  # pixel (1,1) masked
+
+        targets = [
+            {
+                "classes": RasterImage(labels[None, None, :, :].long()),
+                "valid": RasterImage(mask[None, None, :, :]),
+            }
+        ]
+        fm = FeatureMaps([logits])
+        ctx = ModelContext(inputs=[], metadatas=[])
+
+        # --- Cross-entropy reference ---
+        per_pixel_ce = torch.nn.functional.cross_entropy(
+            logits, labels.unsqueeze(0), reduction="none"
+        )  # [1, 2, 2]
+        expected_ce = (per_pixel_ce * mask.unsqueeze(0)).sum() / mask.sum()
+
+        ce_head = SegmentationHead()
+        actual_ce = ce_head(fm, ctx, targets).loss_dict["cls"]
+        assert actual_ce.item() == pytest.approx(expected_ce.item(), rel=1e-5)
+
+        # --- Focal loss reference ---
+        labels_one_hot = (
+            torch.nn.functional.one_hot(labels, 2).permute(2, 0, 1).float().unsqueeze(0)
+        )  # [1, 2, 2, 2]
+        per_pixel_focal = sigmoid_focal_loss(
+            logits, labels_one_hot, alpha=0.25, gamma=2.0, reduction="none"
+        ).sum(dim=1)  # [1, 2, 2]
+        expected_focal = (per_pixel_focal * mask.unsqueeze(0)).sum() / mask.sum()
+
+        focal_head = SegmentationHead(focal_loss=True)
+        actual_focal = focal_head(fm, ctx, targets).loss_dict["cls"]
+        assert actual_focal.item() == pytest.approx(expected_focal.item(), rel=1e-5)
+
+    def test_focal_loss_all_masked_produces_zero(self) -> None:
+        """Focal loss with all pixels masked should produce zero loss."""
+        fm, ctx, targets = self._make_inputs(mask_value=0.0)
+
+        head = SegmentationHead(focal_loss=True)
+        loss = head(fm, ctx, targets).loss_dict["cls"]
+
+        assert loss.item() == 0.0
+
+
 class TestSegmentationMetrics:
     """Tests for SegmentationTask.get_metrics() and metric computation.