rslearn 0.0.27__py3-none-any.whl → 0.0.29__py3-none-any.whl

rslearn/train/dataset.py

@@ -496,53 +496,6 @@ class SplitConfig:
             overlap_ratio: deprecated, use overlap_pixels instead
             load_all_patches: deprecated, use load_all_crops instead
         """
-        # Handle deprecated load_all_patches parameter
-        if load_all_patches is not None:
-            warnings.warn(
-                "load_all_patches is deprecated, use load_all_crops instead",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if load_all_crops is not None:
-                raise ValueError(
-                    "Cannot specify both load_all_patches and load_all_crops"
-                )
-            load_all_crops = load_all_patches
-        # Handle deprecated patch_size parameter
-        if patch_size is not None:
-            warnings.warn(
-                "patch_size is deprecated, use crop_size instead",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if crop_size is not None:
-                raise ValueError("Cannot specify both patch_size and crop_size")
-            crop_size = patch_size
-
-        # Normalize crop_size to tuple[int, int] | None
-        self.crop_size: tuple[int, int] | None = None
-        if crop_size is not None:
-            if isinstance(crop_size, int):
-                self.crop_size = (crop_size, crop_size)
-            else:
-                self.crop_size = crop_size
-
-        # Handle deprecated overlap_ratio parameter
-        if overlap_ratio is not None:
-            warnings.warn(
-                "overlap_ratio is deprecated, use overlap_pixels instead",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if overlap_pixels is not None:
-                raise ValueError("Cannot specify both overlap_ratio and overlap_pixels")
-            if self.crop_size is None:
-                raise ValueError("overlap_ratio requires crop_size to be set")
-            overlap_pixels = round(self.crop_size[0] * overlap_ratio)
-
-        if overlap_pixels is not None and overlap_pixels < 0:
-            raise ValueError("overlap_pixels must be non-negative")
-
         self.groups = groups
         self.names = names
         self.tags = tags
@@ -555,13 +508,22 @@ class SplitConfig:
             output_layer_name_skip_inference_if_exists
         )
 
-        # Note that load_all_crops is handled by the RslearnDataModule rather than the
-        # ModelDataset.
-        self.load_all_crops = load_all_crops
-        self.overlap_pixels = overlap_pixels
+        # These have deprecated equivalents -- we store both raw values since we don't
+        # have a complete picture until the final merged SplitConfig is computed. We
+        # raise deprecation warnings in merge_and_validate and we disambiguate them in
+        # get_ functions (so the variables should never be accessed directly).
+        self._crop_size = crop_size
+        self._patch_size = patch_size
+        self._overlap_pixels = overlap_pixels
+        self._overlap_ratio = overlap_ratio
+        self._load_all_crops = load_all_crops
+        self._load_all_patches = load_all_patches
 
-    def update(self, other: "SplitConfig") -> "SplitConfig":
-        """Override settings in this SplitConfig with those in another.
+    def _merge(self, other: "SplitConfig") -> "SplitConfig":
+        """Merge settings from another SplitConfig into this one.
+
+        Args:
+            other: the config to merge in (its non-None values override self's)
 
         Returns:
             the resulting SplitConfig combining the settings.
@@ -574,9 +536,12 @@ class SplitConfig:
             num_patches=self.num_patches,
             transforms=self.transforms,
             sampler=self.sampler,
-            crop_size=self.crop_size,
-            overlap_pixels=self.overlap_pixels,
-            load_all_crops=self.load_all_crops,
+            crop_size=self._crop_size,
+            patch_size=self._patch_size,
+            overlap_pixels=self._overlap_pixels,
+            overlap_ratio=self._overlap_ratio,
+            load_all_crops=self._load_all_crops,
+            load_all_patches=self._load_all_patches,
             skip_targets=self.skip_targets,
             output_layer_name_skip_inference_if_exists=self.output_layer_name_skip_inference_if_exists,
         )
@@ -594,12 +559,18 @@ class SplitConfig:
             result.transforms = other.transforms
         if other.sampler:
             result.sampler = other.sampler
-        if other.crop_size:
-            result.crop_size = other.crop_size
-        if other.overlap_pixels is not None:
-            result.overlap_pixels = other.overlap_pixels
-        if other.load_all_crops is not None:
-            result.load_all_crops = other.load_all_crops
+        if other._crop_size is not None:
+            result._crop_size = other._crop_size
+        if other._patch_size is not None:
+            result._patch_size = other._patch_size
+        if other._overlap_pixels is not None:
+            result._overlap_pixels = other._overlap_pixels
+        if other._overlap_ratio is not None:
+            result._overlap_ratio = other._overlap_ratio
+        if other._load_all_crops is not None:
+            result._load_all_crops = other._load_all_crops
+        if other._load_all_patches is not None:
+            result._load_all_patches = other._load_all_patches
         if other.skip_targets is not None:
             result.skip_targets = other.skip_targets
         if other.output_layer_name_skip_inference_if_exists is not None:
@@ -608,17 +579,90 @@ class SplitConfig:
             )
         return result
 
+    @staticmethod
+    def merge_and_validate(configs: list["SplitConfig"]) -> "SplitConfig":
+        """Merge a list of SplitConfigs and validate the result.
+
+        Args:
+            configs: list of SplitConfig to merge. Later configs override earlier ones.
+
+        Returns:
+            the merged and validated SplitConfig.
+        """
+        if not configs:
+            return SplitConfig()
+
+        result = configs[0]
+        for config in configs[1:]:
+            result = result._merge(config)
+
+        # Emit deprecation warnings
+        if result._patch_size is not None:
+            warnings.warn(
+                "patch_size is deprecated, use crop_size instead",
+                FutureWarning,
+                stacklevel=2,
+            )
+        if result._overlap_ratio is not None:
+            warnings.warn(
+                "overlap_ratio is deprecated, use overlap_pixels instead",
+                FutureWarning,
+                stacklevel=2,
+            )
+        if result._load_all_patches is not None:
+            warnings.warn(
+                "load_all_patches is deprecated, use load_all_crops instead",
+                FutureWarning,
+                stacklevel=2,
+            )
+
+        # Check for conflicting parameters
+        if result._crop_size is not None and result._patch_size is not None:
+            raise ValueError("Cannot specify both crop_size and patch_size")
+        if result._overlap_pixels is not None and result._overlap_ratio is not None:
+            raise ValueError("Cannot specify both overlap_pixels and overlap_ratio")
+        if result._load_all_crops is not None and result._load_all_patches is not None:
+            raise ValueError("Cannot specify both load_all_crops and load_all_patches")
+
+        # Validate overlap_pixels is non-negative
+        if result._overlap_pixels is not None and result._overlap_pixels < 0:
+            raise ValueError("overlap_pixels must be non-negative")
+
+        # overlap_pixels requires load_all_crops.
+        if result.get_overlap_pixels() > 0 and not result.get_load_all_crops():
+            raise ValueError(
+                "overlap_pixels requires load_all_crops to be True since (overlap is only used during sliding window inference"
+            )
+
+        return result
+
     def get_crop_size(self) -> tuple[int, int] | None:
-        """Get crop size as tuple."""
-        return self.crop_size
+        """Get crop size as tuple, handling deprecated patch_size."""
+        size = self._crop_size if self._crop_size is not None else self._patch_size
+        if size is None:
+            return None
+        if isinstance(size, int):
+            return (size, size)
+        return size
 
     def get_overlap_pixels(self) -> int:
-        """Get the overlap pixels (default 0)."""
-        return self.overlap_pixels if self.overlap_pixels is not None else 0
+        """Get the overlap pixels (default 0), handling deprecated overlap_ratio."""
+        if self._overlap_pixels is not None:
+            return self._overlap_pixels
+        if self._overlap_ratio is not None:
+            crop_size = self.get_crop_size()
+            if crop_size is None:
+                raise ValueError("overlap_ratio requires crop_size to be set")
+            return round(crop_size[0] * self._overlap_ratio)
+        return 0
 
     def get_load_all_crops(self) -> bool:
-        """Returns whether loading all patches is enabled (default False)."""
-        return True if self.load_all_crops is True else False
+        """Returns whether loading all crops is enabled (default False)."""
+        if self._load_all_crops is not None:
+            return self._load_all_crops
+        if self._load_all_patches is not None:
+            return self._load_all_patches
+        return False
 
     def get_skip_targets(self) -> bool:
         """Returns whether skip_targets is enabled (default False)."""
@@ -697,7 +741,7 @@ class ModelDataset(torch.utils.data.Dataset):
         task: Task,
         workers: int,
         name: str | None = None,
-        fix_patch_pick: bool = False,
+        fix_crop_pick: bool = False,
         index_mode: IndexMode = IndexMode.OFF,
     ) -> None:
         """Instantiate a new ModelDataset.
@@ -709,7 +753,7 @@ class ModelDataset(torch.utils.data.Dataset):
             task: the task to train on
             workers: number of workers to use for initializing the dataset
             name: name of the dataset
-            fix_patch_pick: if True, fix the patch pick to be the same every time
+            fix_crop_pick: if True, fix the crop pick to be the same every time
                 for a given window. Useful for testing (default: False)
             index_mode: controls dataset index caching behavior (default: IndexMode.OFF)
         """
@@ -718,14 +762,14 @@ class ModelDataset(torch.utils.data.Dataset):
         self.inputs = inputs
         self.task = task
         self.name = name
-        self.fix_patch_pick = fix_patch_pick
+        self.fix_crop_pick = fix_crop_pick
         if split_config.transforms:
             self.transforms = Sequential(*split_config.transforms)
         else:
             self.transforms = rslearn.train.transforms.transform.Identity()
 
         # Get normalized crop size from the SplitConfig.
-        # But if load all patches is enabled, this is handled by AllCropsDataset, so
+        # But if load_all_crops is enabled, this is handled by AllCropsDataset, so
         # here we instead load the entire windows.
         if split_config.get_load_all_crops():
             self.crop_size = None
@@ -952,7 +996,7 @@ class ModelDataset(torch.utils.data.Dataset):
         """Get a list of examples in the dataset.
 
         If load_all_crops is False, this is a list of Windows. Otherwise, this is a
-        list of (window, crop_bounds, (crop_idx, # patches)) tuples.
+        list of (window, crop_bounds, (crop_idx, # crops)) tuples.
         """
         if self.dataset_examples is None:
             logger.debug(
@@ -985,7 +1029,7 @@ class ModelDataset(torch.utils.data.Dataset):
         """
         dataset_examples = self.get_dataset_examples()
         example = dataset_examples[idx]
-        rng = random.Random(idx if self.fix_patch_pick else None)
+        rng = random.Random(idx if self.fix_crop_pick else None)
 
         # Select bounds to read.
         if self.crop_size:

rslearn/train/lightning_module.py

@@ -6,12 +6,14 @@ from typing import Any
 
 import lightning as L
 import torch
+import wandb
 from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
 from PIL import Image
 from upath import UPath
 
 from rslearn.log_utils import get_logger
 
+from .metrics import NonScalarMetricOutput
 from .model_context import ModelContext, ModelOutput
 from .optimizer import AdamW, OptimizerFactory
 from .scheduler import PlateauScheduler, SchedulerFactory
@@ -210,15 +212,53 @@ class RslearnLightningModule(L.LightningModule):
             # Fail silently for single-dataset case, which is okay
             pass
 
+    def _log_non_scalar_metric(self, name: str, value: NonScalarMetricOutput) -> None:
+        """Log a non-scalar metric to wandb.
+
+        Args:
+            name: the metric name (e.g., "val_confusion_matrix")
+            value: the non-scalar metric output
+        """
+        # The non-scalar metrics are logging directly without Lightning
+        # So we need to skip logging during sanity check.
+        if self.trainer.sanity_checking:
+            return
+
+        # Wandb is required for logging non-scalar metrics.
+        if not wandb.run:
+            logger.warning(
+                f"Weights & Biases is not initialized, skipping logging of {name}"
+            )
+            return
+
+        value.log_to_wandb(name)
+
     def on_validation_epoch_end(self) -> None:
         """Compute and log validation metrics at epoch end.
 
         We manually compute and log metrics here (instead of passing the MetricCollection
         to log_dict) because MetricCollection.compute() properly flattens dict-returning
         metrics, while log_dict expects each metric to return a scalar tensor.
+
+        Non-scalar metrics (like confusion matrices) are logged separately using
+        logger-specific APIs.
         """
         metrics = self.val_metrics.compute()
-        self.log_dict(metrics)
+
+        # Separate scalar and non-scalar metrics
+        scalar_metrics = {}
+        for k, v in metrics.items():
+            if isinstance(v, NonScalarMetricOutput):
+                self._log_non_scalar_metric(k, v)
+            elif isinstance(v, torch.Tensor) and v.dim() > 0 and v.numel() > 1:
+                raise ValueError(
+                    f"Metric '{k}' returned a non-scalar tensor with shape {v.shape}. "
+                    "Wrap it in a NonScalarMetricOutput subclass."
+                )
+            else:
+                scalar_metrics[k] = v
+
+        self.log_dict(scalar_metrics)
         self.val_metrics.reset()
 
     def on_test_epoch_end(self) -> None:
@@ -227,14 +267,30 @@ class RslearnLightningModule(L.LightningModule):
         We manually compute and log metrics here (instead of passing the MetricCollection
         to log_dict) because MetricCollection.compute() properly flattens dict-returning
         metrics, while log_dict expects each metric to return a scalar tensor.
+
+        Non-scalar metrics (like confusion matrices) are logged separately.
         """
         metrics = self.test_metrics.compute()
-        self.log_dict(metrics)
+
+        # Separate scalar and non-scalar metrics
+        scalar_metrics = {}
+        for k, v in metrics.items():
+            if isinstance(v, NonScalarMetricOutput):
+                self._log_non_scalar_metric(k, v)
+            elif isinstance(v, torch.Tensor) and v.dim() > 0 and v.numel() > 1:
+                raise ValueError(
+                    f"Metric '{k}' returned a non-scalar tensor with shape {v.shape}. "
+                    "Wrap it in a NonScalarMetricOutput subclass."
+                )
+            else:
+                scalar_metrics[k] = v
+
+        self.log_dict(scalar_metrics)
         self.test_metrics.reset()
 
         if self.metrics_file:
             with open(self.metrics_file, "w") as f:
-                metrics_dict = {k: v.item() for k, v in metrics.items()}
+                metrics_dict = {k: v.item() for k, v in scalar_metrics.items()}
                 json.dump(metrics_dict, f, indent=4)
                 logger.info(f"Saved metrics to {self.metrics_file}")
 

rslearn/train/metrics.py

@@ -0,0 +1,162 @@
+"""Metric output classes for non-scalar metrics."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import torch
+import wandb
+from torchmetrics import Metric
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class NonScalarMetricOutput(ABC):
+    """Base class for non-scalar metric outputs that need special logging.
+
+    Subclasses should implement the log_to_wandb method to define how the metric
+    should be logged (only supports logging to Weights & Biases).
+    """
+
+    @abstractmethod
+    def log_to_wandb(self, name: str) -> None:
+        """Log this metric to wandb.
+
+        Args:
+            name: the metric name
+        """
+        pass
+
+
+@dataclass
+class ConfusionMatrixOutput(NonScalarMetricOutput):
+    """Confusion matrix metric output.
+
+    Args:
+        confusion_matrix: confusion matrix of shape (num_classes, num_classes)
+            where cm[i, j] is the count of samples with true label i and predicted
+            label j.
+        class_names: optional list of class names for axis labels
+    """
+
+    confusion_matrix: torch.Tensor
+    class_names: list[str] | None = None
+
+    def _expand_confusion_matrix(self) -> tuple[list[int], list[int]]:
+        """Expand confusion matrix to (preds, labels) pairs for wandb.
+
+        Returns:
+            Tuple of (preds, labels) as lists of integers.
+        """
+        cm = self.confusion_matrix.detach().cpu()
+
+        # Handle extra dimensions from distributed reduction
+        if cm.dim() > 2:
+            cm = cm.sum(dim=0)
+
+        total = cm.sum().item()
+        if total == 0:
+            return [], []
+
+        preds = []
+        labels = []
+        for true_label in range(cm.shape[0]):
+            for pred_label in range(cm.shape[1]):
+                count = cm[true_label, pred_label].item()
+                if count > 0:
+                    preds.extend([pred_label] * int(count))
+                    labels.extend([true_label] * int(count))
+
+        return preds, labels
+
+    def log_to_wandb(self, name: str) -> None:
+        """Log confusion matrix to wandb.
+
+        Args:
+            name: the metric name (e.g., "val_confusion_matrix")
+        """
+        preds, labels = self._expand_confusion_matrix()
+
+        if len(preds) == 0:
+            logger.warning(f"No samples to log for {name}")
+            return
+
+        num_classes = self.confusion_matrix.shape[0]
+        if self.class_names is None:
+            class_names = [str(i) for i in range(num_classes)]
+        else:
+            class_names = self.class_names
+
+        wandb.log(
+            {
+                name: wandb.plot.confusion_matrix(
+                    preds=preds,
+                    y_true=labels,
+                    class_names=class_names,
+                    title=name,
+                ),
+            },
+        )
+
+
+class ConfusionMatrixMetric(Metric):
+    """Confusion matrix metric that works on flattened inputs.
+
+    Expects preds of shape (N, C) and labels of shape (N,).
+    Should be wrapped by ClassificationMetric or SegmentationMetric
+    which handle the task-specific preprocessing.
+
+    Args:
+        num_classes: number of classes
+        class_names: optional list of class names for labeling
+    """
+
+    def __init__(
+        self,
+        num_classes: int,
+        class_names: list[str] | None = None,
+    ):
+        """Initialize a new ConfusionMatrixMetric.
+
+        Args:
+            num_classes: number of classes
+            class_names: optional list of class names for labeling
+        """
+        super().__init__()
+        self.num_classes = num_classes
+        self.class_names = class_names
+        self.add_state(
+            "confusion_matrix",
+            default=torch.zeros(num_classes, num_classes, dtype=torch.long),
+            dist_reduce_fx="sum",
+        )
+
+    def update(self, preds: torch.Tensor, labels: torch.Tensor) -> None:
+        """Update metric.
+
+        Args:
+            preds: predictions of shape (N, C) - probabilities
+            labels: ground truth of shape (N,) - class indices
+        """
+        if len(preds) == 0:
+            return
+
+        pred_classes = preds.argmax(dim=1)  # (N,)
+
+        for true_label in range(self.num_classes):
+            for pred_label in range(self.num_classes):
+                count = ((labels == true_label) & (pred_classes == pred_label)).sum()
+                self.confusion_matrix[true_label, pred_label] += count
+
+    def compute(self) -> ConfusionMatrixOutput:
+        """Returns the confusion matrix wrapped in ConfusionMatrixOutput."""
+        return ConfusionMatrixOutput(
+            confusion_matrix=self.confusion_matrix,
+            class_names=self.class_names,
+        )
+
+    def reset(self) -> None:
+        """Reset metric."""
+        super().reset()

rslearn/train/tasks/classification.py

@@ -16,6 +16,7 @@ from torchmetrics.classification import (
 )
 
 from rslearn.models.component import FeatureVector, Predictor
+from rslearn.train.metrics import ConfusionMatrixMetric
 from rslearn.train.model_context import (
     ModelContext,
     ModelOutput,
@@ -44,6 +45,7 @@ class ClassificationTask(BasicTask):
         f1_metric_kwargs: dict[str, Any] = {},
         positive_class: str | None = None,
         positive_class_threshold: float = 0.5,
+        enable_confusion_matrix: bool = False,
         **kwargs: Any,
     ):
         """Initialize a new ClassificationTask.
@@ -69,6 +71,8 @@ class ClassificationTask(BasicTask):
             positive_class: positive class name.
             positive_class_threshold: threshold for classifying the positive class in
                 binary classification (default 0.5).
+            enable_confusion_matrix: whether to compute confusion matrix (default false).
+                If true, it requires wandb to be initialized for logging.
             kwargs: other arguments to pass to BasicTask
         """
         super().__init__(**kwargs)
@@ -84,6 +88,7 @@ class ClassificationTask(BasicTask):
         self.f1_metric_kwargs = f1_metric_kwargs
         self.positive_class = positive_class
         self.positive_class_threshold = positive_class_threshold
+        self.enable_confusion_matrix = enable_confusion_matrix
 
         if self.positive_class_threshold != 0.5:
             # Must be binary classification
@@ -278,6 +283,14 @@ class ClassificationTask(BasicTask):
                 )
                 metrics["f1"] = ClassificationMetric(MulticlassF1Score(**kwargs))
 
+        if self.enable_confusion_matrix:
+            metrics["confusion_matrix"] = ClassificationMetric(
+                ConfusionMatrixMetric(
+                    num_classes=len(self.classes),
+                    class_names=self.classes,
+                ),
+            )
+
         return MetricCollection(metrics)
 
 

rslearn/train/tasks/per_pixel_regression.py

@@ -149,22 +149,28 @@ class PerPixelRegressionHead(Predictor):
     """Head for per-pixel regression task."""
 
     def __init__(
-        self, loss_mode: Literal["mse", "l1"] = "mse", use_sigmoid: bool = False
+        self,
+        loss_mode: Literal["mse", "l1", "huber"] = "mse",
+        use_sigmoid: bool = False,
+        huber_delta: float = 1.0,
     ):
-        """Initialize a new RegressionHead.
+        """Initialize a new PerPixelRegressionHead.
 
         Args:
-            loss_mode: the loss function to use, either "mse" (default) or "l1".
+            loss_mode: the loss function to use: "mse" (default), "l1", or "huber".
             use_sigmoid: whether to apply a sigmoid activation on the output. This
                 requires targets to be between 0-1.
+            huber_delta: delta parameter for Huber loss (only used when
+                loss_mode="huber").
         """
         super().__init__()
 
-        if loss_mode not in ["mse", "l1"]:
-            raise ValueError("invalid loss mode")
+        if loss_mode not in ["mse", "l1", "huber"]:
+            raise ValueError(f"invalid loss mode {loss_mode}")
 
         self.loss_mode = loss_mode
         self.use_sigmoid = use_sigmoid
+        self.huber_delta = huber_delta
 
     def forward(
         self,
@@ -217,8 +223,15 @@ class PerPixelRegressionHead(Predictor):
                 scores = torch.square(outputs - labels)
             elif self.loss_mode == "l1":
                 scores = torch.abs(outputs - labels)
+            elif self.loss_mode == "huber":
+                scores = torch.nn.functional.huber_loss(
+                    outputs,
+                    labels,
+                    reduction="none",
+                    delta=self.huber_delta,
+                )
             else:
-                assert False
+                raise ValueError(f"unknown loss mode {self.loss_mode}")
 
             # Compute average but only over valid pixels.
             mask_total = mask.sum()

rslearn/train/tasks/regression.py

@@ -196,18 +196,24 @@ class RegressionHead(Predictor):
     """Head for regression task."""
 
     def __init__(
-        self, loss_mode: Literal["mse", "l1"] = "mse", use_sigmoid: bool = False
+        self,
+        loss_mode: Literal["mse", "l1", "huber"] = "mse",
+        use_sigmoid: bool = False,
+        huber_delta: float = 1.0,
     ):
         """Initialize a new RegressionHead.
 
         Args:
-            loss_mode: the loss function to use, either "mse" (default) or "l1".
+            loss_mode: the loss function to use: "mse" (default), "l1", or "huber".
             use_sigmoid: whether to apply a sigmoid activation on the output. This
                 requires targets to be between 0-1.
+            huber_delta: delta parameter for Huber loss (only used when
+                loss_mode="huber").
         """
         super().__init__()
         self.loss_mode = loss_mode
         self.use_sigmoid = use_sigmoid
+        self.huber_delta = huber_delta
 
     def forward(
         self,
@@ -251,6 +257,16 @@ class RegressionHead(Predictor):
                 losses["regress"] = torch.mean(torch.square(outputs - labels) * mask)
             elif self.loss_mode == "l1":
                 losses["regress"] = torch.mean(torch.abs(outputs - labels) * mask)
+            elif self.loss_mode == "huber":
+                losses["regress"] = torch.mean(
+                    torch.nn.functional.huber_loss(
+                        outputs,
+                        labels,
+                        reduction="none",
+                        delta=self.huber_delta,
+                    )
+                    * mask
+                )
             else:
                 raise ValueError(f"unknown loss mode {self.loss_mode}")