rslearn 0.0.17__py3-none-any.whl → 0.0.18__py3-none-any.whl

rslearn/train/lightning_module.py

@@ -12,6 +12,7 @@ from upath import UPath
 
 from rslearn.log_utils import get_logger
 
+from .model_context import ModelContext, ModelOutput
 from .optimizer import AdamW, OptimizerFactory
 from .scheduler import PlateauScheduler, SchedulerFactory
 from .tasks import Task
@@ -231,12 +232,16 @@ class RslearnLightningModule(L.LightningModule):
         Returns:
             The loss tensor.
         """
-        inputs, targets, _ = batch
+        inputs, targets, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
         batch_size = len(inputs)
-        model_outputs = self(inputs, targets)
-        self.on_train_forward(inputs, targets, model_outputs)
+        model_outputs = self(context, targets)
+        self.on_train_forward(context, targets, model_outputs)
 
-        loss_dict = model_outputs["loss_dict"]
+        loss_dict = model_outputs.loss_dict
         train_loss = sum(loss_dict.values())
         self.log_dict(
             {"train_" + k: v for k, v in loss_dict.items()},
@@ -266,13 +271,17 @@ class RslearnLightningModule(L.LightningModule):
             batch_idx: Integer displaying index of this batch.
             dataloader_idx: Index of the current dataloader.
         """
-        inputs, targets, _ = batch
+        inputs, targets, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
         batch_size = len(inputs)
-        model_outputs = self(inputs, targets)
-        self.on_val_forward(inputs, targets, model_outputs)
+        model_outputs = self(context, targets)
+        self.on_val_forward(context, targets, model_outputs)
 
-        loss_dict = model_outputs["loss_dict"]
-        outputs = model_outputs["outputs"]
+        loss_dict = model_outputs.loss_dict
+        outputs = model_outputs.outputs
         val_loss = sum(loss_dict.values())
         self.log_dict(
             {"val_" + k: v for k, v in loss_dict.items()},
@@ -304,12 +313,16 @@ class RslearnLightningModule(L.LightningModule):
             dataloader_idx: Index of the current dataloader.
         """
         inputs, targets, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
         batch_size = len(inputs)
-        model_outputs = self(inputs, targets)
-        self.on_test_forward(inputs, targets, model_outputs)
+        model_outputs = self(context, targets)
+        self.on_test_forward(context, targets, model_outputs)
 
-        loss_dict = model_outputs["loss_dict"]
-        outputs = model_outputs["outputs"]
+        loss_dict = model_outputs.loss_dict
+        outputs = model_outputs.outputs
         test_loss = sum(loss_dict.values())
         self.log_dict(
             {"test_" + k: v for k, v in loss_dict.items()},
@@ -345,7 +358,7 @@ class RslearnLightningModule(L.LightningModule):
 
     def predict_step(
         self, batch: Any, batch_idx: int, dataloader_idx: int = 0
-    ) -> torch.Tensor:
+    ) -> ModelOutput:
         """Compute the predicted class probabilities.
 
         Args:
@@ -356,63 +369,69 @@ class RslearnLightningModule(L.LightningModule):
         Returns:
             Output predicted probabilities.
         """
-        inputs, _, _ = batch
-        model_outputs = self(inputs)
+        inputs, _, metadatas = batch
+        context = ModelContext(
+            inputs=inputs,
+            metadatas=metadatas,
+        )
+        model_outputs = self(context)
         return model_outputs
 
-    def forward(self, *args: Any, **kwargs: Any) -> Any:
+    def forward(
+        self, context: ModelContext, targets: list[dict[str, Any]] | None = None
+    ) -> ModelOutput:
         """Forward pass of the model.
 
         Args:
-            args: Arguments to pass to model.
-            kwargs: Keyword arguments to pass to model.
+            context: the model context.
+            targets: the target dicts.
 
         Returns:
             Output of the model.
         """
-        return self.model(*args, **kwargs)
+        return self.model(context, targets)
 
     def on_train_forward(
         self,
-        inputs: list[dict[str, Any]],
+        context: ModelContext,
         targets: list[dict[str, Any]],
-        model_outputs: dict[str, Any],
+        model_outputs: ModelOutput,
     ) -> None:
         """Hook to run after the forward pass of the model during training.
 
         Args:
-            inputs: The input batch.
+            context: The model context.
             targets: The target batch.
-            model_outputs: The output of the model, with keys "outputs" and "loss_dict", and possibly other keys.
+            model_outputs: The output of the model.
         """
         pass
 
     def on_val_forward(
         self,
-        inputs: list[dict[str, Any]],
+        context: ModelContext,
         targets: list[dict[str, Any]],
-        model_outputs: dict[str, Any],
+        model_outputs: ModelOutput,
     ) -> None:
         """Hook to run after the forward pass of the model during validation.
 
         Args:
-            inputs: The input batch.
+            context: The model context.
             targets: The target batch.
-            model_outputs: The output of the model, with keys "outputs" and "loss_dict", and possibly other keys.
+            model_outputs: The output of the model.
         """
         pass
 
     def on_test_forward(
         self,
-        inputs: list[dict[str, Any]],
+        context: ModelContext,
         targets: list[dict[str, Any]],
-        model_outputs: dict[str, Any],
+        model_outputs: ModelOutput,
     ) -> None:
         """Hook to run after the forward pass of the model during testing.
 
         Args:
-            inputs: The input batch.
+            context: The model context.
             targets: The target batch.
-            model_outputs: The output of the model, with keys "outputs" and "loss_dict", and possibly other keys.
+            model_outputs: The output of the model.
         """
         pass

rslearn/train/model_context.py

@@ -0,0 +1,54 @@
+"""Data classes to provide various context to models."""
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+import torch
+
+from rslearn.utils.geometry import PixelBounds, Projection
+
+
+@dataclass
+class SampleMetadata:
+    """Metadata pertaining to an example."""
+
+    window_group: str
+    window_name: str
+    window_bounds: PixelBounds
+    patch_bounds: PixelBounds
+    patch_idx: int
+    num_patches_in_window: int
+    time_range: tuple[datetime, datetime] | None
+    projection: Projection
+
+    # Task name to differentiate different tasks.
+    dataset_source: str | None
+
+
+@dataclass
+class ModelContext:
+    """Context to pass to all model components."""
+
+    # One input dict per example in the batch.
+    inputs: list[dict[str, torch.Tensor]]
+    # One SampleMetadata per example in the batch.
+    metadatas: list[SampleMetadata]
+    # Arbitrary dict that components can add to.
+    context_dict: dict[str, Any] = field(default_factory=lambda: {})
+
+
+@dataclass
+class ModelOutput:
+    """The output from the Predictor.
+
+    Args:
+        outputs: output compatible with the configured Task.
+        loss_dict: map from loss names to scalar tensors.
+        metadata: arbitrary dict that can be used to store other outputs.
+    """
+
+    outputs: Iterable[Any]
+    loss_dict: dict[str, torch.Tensor]
+    metadata: dict[str, Any] = field(default_factory=lambda: {})

rslearn/train/prediction_writer.py

@@ -1,6 +1,7 @@
 """rslearn PredictionWriter implementation."""
 
-from collections.abc import Sequence
+import json
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
@@ -12,11 +13,15 @@ from lightning.pytorch.callbacks import BasePredictionWriter
 from upath import UPath
 
 from rslearn.config import (
+    DatasetConfig,
     LayerConfig,
     LayerType,
+    StorageConfig,
 )
-from rslearn.dataset import Dataset, Window
+from rslearn.dataset import Window
+from rslearn.dataset.storage.storage import WindowStorage
 from rslearn.log_utils import get_logger
+from rslearn.train.model_context import SampleMetadata
 from rslearn.utils.array import copy_spatial_array
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds
@@ -27,6 +32,7 @@ from rslearn.utils.raster_format import (
 from rslearn.utils.vector_format import VectorFormat
 
 from .lightning_module import RslearnLightningModule
+from .model_context import ModelOutput
 from .tasks.task import Task
 
 logger = get_logger(__name__)
@@ -43,12 +49,18 @@ class PendingPatchOutput:
 class PatchPredictionMerger:
     """Base class for merging predictions from multiple patches."""
 
-    def merge(self, window: Window, outputs: Sequence[PendingPatchOutput]) -> Any:
+    def merge(
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
+    ) -> Any:
         """Merge the outputs.
 
         Args:
             window: the window we are merging the outputs for.
             outputs: the outputs to process.
+            layer_config: the output layer configuration.
 
         Returns:
             the merged outputs.
@@ -60,7 +72,10 @@ class VectorMerger(PatchPredictionMerger):
     """Merger for vector data that simply concatenates the features."""
 
     def merge(
-        self, window: Window, outputs: Sequence[PendingPatchOutput]
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
     ) -> list[Feature]:
         """Concatenate the vector features."""
         return [feat for output in outputs for feat in output.output]
@@ -83,18 +98,20 @@ class RasterMerger(PatchPredictionMerger):
         self.downsample_factor = downsample_factor
 
     def merge(
-        self, window: Window, outputs: Sequence[PendingPatchOutput]
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
     ) -> npt.NDArray:
         """Merge the raster outputs."""
         num_channels = outputs[0].output.shape[0]
-        dtype = outputs[0].output.dtype
         merged_image = np.zeros(
             (
                 num_channels,
                 (window.bounds[3] - window.bounds[1]) // self.downsample_factor,
                 (window.bounds[2] - window.bounds[0]) // self.downsample_factor,
             ),
-            dtype=dtype,
+            dtype=layer_config.band_sets[0].dtype.get_numpy_dtype(),
         )
 
         # Ensure the outputs are sorted by height then width.
@@ -148,6 +165,7 @@ class RslearnWriter(BasePredictionWriter):
         merger: PatchPredictionMerger | None = None,
         output_path: str | Path | None = None,
         layer_config: LayerConfig | None = None,
+        storage_config: StorageConfig | None = None,
     ):
         """Create a new RslearnWriter.
 
@@ -163,28 +181,24 @@ class RslearnWriter(BasePredictionWriter):
             layer_config: optional layer configuration. If provided, this config will be
                 used instead of reading from the dataset config, allowing usage without
                 requiring dataset config at the output path.
+            storage_config: optional storage configuration, needed similar to layer_config
+                if there is no dataset config.
         """
         super().__init__(write_interval="batch")
         self.output_layer = output_layer
         self.selector = selector or []
-        self.path = UPath(path, **path_options or {})
-        self.output_path = (
+        ds_upath = UPath(path, **path_options or {})
+        output_upath = (
             UPath(output_path, **path_options or {})
             if output_path is not None
-            else None
+            else ds_upath
         )
 
-        # Handle dataset and layer config
-        self.layer_config: LayerConfig
-        if layer_config:
-            self.layer_config = layer_config
-        else:
-            dataset = Dataset(self.path)
-            if self.output_layer not in dataset.layers:
-                raise KeyError(
-                    f"Output layer '{self.output_layer}' not found in dataset layers."
-                )
-            self.layer_config = dataset.layers[self.output_layer]
+        self.layer_config, self.dataset_storage = (
+            self._get_layer_config_and_dataset_storage(
+                ds_upath, output_upath, layer_config, storage_config
+            )
+        )
 
         self.format: RasterFormat | VectorFormat
         if self.layer_config.type == LayerType.RASTER:
@@ -207,11 +221,73 @@ class RslearnWriter(BasePredictionWriter):
         # patches of each window need to be reconstituted.
         self.pending_outputs: dict[str, list[PendingPatchOutput]] = {}
 
+    def _get_layer_config_and_dataset_storage(
+        self,
+        ds_upath: UPath,
+        output_upath: UPath,
+        layer_config: LayerConfig | None,
+        storage_config: StorageConfig | None,
+    ) -> tuple[LayerConfig, WindowStorage]:
+        """Get the layer config and dataset storage to use.
+
+        This is a helper function for the init method.
+
+        If layer_config is set, we use that. If storage_config is set, we use it to
+        instantiate a WindowStorage using the output_upath.
+
+        If one of them is not set, we load the config from the ds_upath. Otherwise, we
+        avoid reading the dataset config; this way, RslearnWriter can be used with
+        output directories that do not contain the dataset config, as long as
+        layer_config and storage_config are both provided.
+
+        Args:
+            ds_upath: the dataset path, where a dataset config can be loaded from if
+                layer_config or storage_config is not provided.
+            output_upath: the output directory, which could be different from the
+                dataset path.
+            layer_config: optional LayerConfig to provide.
+            storage_config: optional StorageConfig to provide.
+
+        Returns:
+            a tuple (layer_config, dataset_storage)
+        """
+        dataset_storage: WindowStorage | None = None
+
+        # Instantiate the WindowStorage from the storage_config if provided.
+        if storage_config:
+            dataset_storage = (
+                storage_config.instantiate_window_storage_factory().get_storage(
+                    output_upath
+                )
+            )
+
+        if not layer_config or not dataset_storage:
+            # Need to load dataset config since one of LayerConfig/StorageConfig is missing.
+            # We use DatasetConfig.model_validate instead of initializing the Dataset
+            # because we want to get a WindowStorage that has the dataset path set to
+            # output_upath instead of ds_upath.
+            with (ds_upath / "config.json").open() as f:
+                dataset_config = DatasetConfig.model_validate(json.load(f))
+
+            if not layer_config:
+                if self.output_layer not in dataset_config.layers:
+                    raise KeyError(
+                        f"Output layer '{self.output_layer}' not found in dataset layers."
+                    )
+                layer_config = dataset_config.layers[self.output_layer]
+
+            if not dataset_storage:
+                dataset_storage = dataset_config.storage.instantiate_window_storage_factory().get_storage(
+                    output_upath
+                )
+
+        return (layer_config, dataset_storage)
+
     def write_on_batch_end(
         self,
         trainer: Trainer,
         pl_module: LightningModule,
-        prediction: dict[str, Sequence],
+        prediction: ModelOutput,
         batch_indices: Sequence[int] | None,
         batch: tuple[list, list, list],
         batch_idx: int,
@@ -232,13 +308,13 @@ class RslearnWriter(BasePredictionWriter):
         assert isinstance(pl_module, RslearnLightningModule)
         task = pl_module.task
         _, _, metadatas = batch
-        self.process_output_batch(task, prediction["outputs"], metadatas)
+        self.process_output_batch(task, prediction.outputs, metadatas)
 
     def process_output_batch(
         self,
         task: Task,
-        prediction: Sequence,
-        metadatas: Sequence,
+        prediction: Iterable[Any],
+        metadatas: Iterable[SampleMetadata],
     ) -> None:
         """Write a prediction batch with simplified API.
 
@@ -263,25 +339,19 @@ class RslearnWriter(BasePredictionWriter):
             for k in self.selector:
                 output = output[k]
 
-            # Use custom output_path if provided, otherwise use dataset path
-            window_base_path = (
-                self.output_path if self.output_path is not None else self.path
-            )
             window = Window(
-                path=Window.get_window_root(
-                    window_base_path, metadata["group"], metadata["window_name"]
-                ),
-                group=metadata["group"],
-                name=metadata["window_name"],
-                projection=metadata["projection"],
-                bounds=metadata["window_bounds"],
-                time_range=metadata["time_range"],
+                storage=self.dataset_storage,
+                group=metadata.window_group,
+                name=metadata.window_name,
+                projection=metadata.projection,
+                bounds=metadata.window_bounds,
+                time_range=metadata.time_range,
             )
             self.process_output(
                 window,
-                metadata["patch_idx"],
-                metadata["num_patches"],
-                metadata["bounds"],
+                metadata.patch_idx,
+                metadata.num_patches_in_window,
+                metadata.patch_bounds,
                 output,
             )
 
@@ -320,7 +390,7 @@ class RslearnWriter(BasePredictionWriter):
 
         # Merge outputs from overlapped patches if merger is set.
         logger.debug(f"Merging and writing for window {window.name}")
-        merged_output = self.merger.merge(window, pending_output)
+        merged_output = self.merger.merge(window, pending_output, self.layer_config)
 
         if self.layer_config.type == LayerType.RASTER:
             raster_dir = window.get_raster_dir(

rslearn/train/tasks/classification.py

@@ -15,6 +15,8 @@ from torchmetrics.classification import (
     MulticlassRecall,
 )
 
+from rslearn.models.component import FeatureVector, Predictor
+from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
@@ -98,7 +100,7 @@ class ClassificationTask(BasicTask):
     def process_inputs(
         self,
         raw_inputs: dict[str, torch.Tensor | list[Feature]],
-        metadata: dict[str, Any],
+        metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """Processes the data into targets.
@@ -154,17 +156,25 @@ class ClassificationTask(BasicTask):
         }
 
     def process_output(
-        self, raw_output: Any, metadata: dict[str, Any]
-    ) -> npt.NDArray[Any] | list[Feature]:
+        self, raw_output: Any, metadata: SampleMetadata
+    ) -> list[Feature]:
         """Processes an output into raster or vector data.
 
         Args:
-            raw_output: the output from prediction head.
+            raw_output: the output from prediction head, which must be a tensor
+                containing output probabilities (one dimension).
             metadata: metadata about the patch being read
 
         Returns:
-            either raster or vector data.
+            a list with one Feature corresponding to the input patch extent with a
+                property name containing the predicted class. It will have another
+                property containing the probabilities if prob_property was set.
         """
+        if not isinstance(raw_output, torch.Tensor) or len(raw_output.shape) != 1:
+            raise ValueError(
+                "expected output for ClassificationTask to be a Tensor with one dimension"
+            )
+
         probs = raw_output.cpu().numpy()
         if len(self.classes) == 2 and self.positive_class_threshold != 0.5:
             positive_class_prob = probs[self.positive_class_id]
@@ -184,8 +194,8 @@ class ClassificationTask(BasicTask):
 
         feature = Feature(
             STGeometry(
-                metadata["projection"],
-                shapely.Point(metadata["bounds"][0], metadata["bounds"][1]),
+                metadata.projection,
+                shapely.Point(metadata.patch_bounds[0], metadata.patch_bounds[1]),
                 None,
             ),
             {
@@ -265,25 +275,31 @@ class ClassificationTask(BasicTask):
         return MetricCollection(metrics)
 
 
-class ClassificationHead(torch.nn.Module):
+class ClassificationHead(Predictor):
     """Head for classification task."""
 
     def forward(
         self,
-        logits: torch.Tensor,
-        inputs: list[dict[str, Any]],
+        intermediates: Any,
+        context: ModelContext,
         targets: list[dict[str, Any]] | None = None,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+    ) -> ModelOutput:
         """Compute the classification outputs and loss from logits and targets.
 
         Args:
-            logits: tensor that is (BatchSize, NumClasses) in shape.
-            inputs: original inputs (ignored).
-            targets: should contain class key that stores the class label.
+            intermediates: output from the previous model component, it should be a
+                FeatureVector with a tensor that is (BatchSize, NumClasses) in shape.
+            context: the model context.
+            targets: must contain "class" key that stores the class label, along with
+                "valid" key indicating whether the label is valid for each example.
 
         Returns:
             tuple of outputs and loss dict
         """
+        if not isinstance(intermediates, FeatureVector):
+            raise ValueError("the input to ClassificationHead must be a FeatureVector")
+
+        logits = intermediates.feature_vector
         outputs = torch.nn.functional.softmax(logits, dim=1)
 
         losses = {}
@@ -298,7 +314,10 @@ class ClassificationHead(torch.nn.Module):
             )
             losses["cls"] = torch.mean(loss)
 
-        return outputs, losses
+        return ModelOutput(
+            outputs=outputs,
+            loss_dict=losses,
+        )
 
 
 class ClassificationMetric(Metric):