rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/train/prediction_writer.py

@@ -1,40 +1,154 @@
 """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
 
 import numpy as np
+import numpy.typing as npt
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import BasePredictionWriter
 from upath import UPath
 
-from rslearn.config import LayerType, RasterFormatConfig
-from rslearn.dataset import Dataset
+from rslearn.config import (
+    DatasetConfig,
+    LayerConfig,
+    LayerType,
+    StorageConfig,
+)
+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.raster_format import load_raster_format
-from rslearn.utils.vector_format import load_vector_format
+from rslearn.utils.feature import Feature
+from rslearn.utils.geometry import PixelBounds
+from rslearn.utils.raster_format import (
+    RasterFormat,
+    adjust_projection_and_bounds_for_array,
+)
+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__)
+
+
+@dataclass
+class PendingPatchOutput:
+    """A patch output that hasn't been merged yet."""
+
+    bounds: PixelBounds
+    output: Any
 
 
 class PatchPredictionMerger:
     """Base class for merging predictions from multiple patches."""
 
     def merge(
-        self, outputs: Sequence[Any], metadatas: Sequence[Any]
-    ) -> tuple[Sequence[Any], Sequence[Any]]:
-        """Merge the outputs and metadatas.
+        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.
-            metadatas: the metadatas to process.
+            layer_config: the output layer configuration.
 
         Returns:
-            the merged outputs and metadatas.
+            the merged outputs.
         """
         raise NotImplementedError
 
 
+class VectorMerger(PatchPredictionMerger):
+    """Merger for vector data that simply concatenates the features."""
+
+    def merge(
+        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]
+
+
+class RasterMerger(PatchPredictionMerger):
+    """Merger for raster data that copies the rasters to the output."""
+
+    def __init__(self, padding: int | None = None, downsample_factor: int = 1):
+        """Create a new RasterMerger.
+
+        Args:
+            padding: the padding around the individual patch outputs to remove. This is
+                typically used when leveraging overlapping patches. Portions of outputs
+                at the border of the window will still be retained.
+            downsample_factor: the factor by which the rasters output by the task are
+                lower in resolution relative to the window resolution.
+        """
+        self.padding = padding
+        self.downsample_factor = downsample_factor
+
+    def merge(
+        self,
+        window: Window,
+        outputs: Sequence[PendingPatchOutput],
+        layer_config: LayerConfig,
+    ) -> npt.NDArray:
+        """Merge the raster outputs."""
+        num_channels = outputs[0].output.shape[0]
+        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=layer_config.band_sets[0].dtype.get_numpy_dtype(),
+        )
+
+        # Ensure the outputs are sorted by height then width.
+        # This way when we merge we can be sure that outputs that are lower or further
+        # to the right will overwrite earlier outputs.
+        sorted_outputs = sorted(
+            outputs, key=lambda output: (output.bounds[0], output.bounds[1])
+        )
+        for output in sorted_outputs:
+            # So now we just need to compute the src_offset to copy.
+            # If the output is not on the left or top boundary, then we should apply
+            # the padding (if set).
+            src = output.output
+            src_offset = (
+                output.bounds[0] // self.downsample_factor,
+                output.bounds[1] // self.downsample_factor,
+            )
+            if self.padding is not None and output.bounds[0] != window.bounds[0]:
+                src = src[:, :, self.padding :]
+                src_offset = (src_offset[0] + self.padding, src_offset[1])
+            if self.padding is not None and output.bounds[1] != window.bounds[1]:
+                src = src[:, self.padding :, :]
+                src_offset = (src_offset[0], src_offset[1] + self.padding)
+
+            copy_spatial_array(
+                src=src,
+                dst=merged_image,
+                src_offset=src_offset,
+                dst_offset=(
+                    window.bounds[0] // self.downsample_factor,
+                    window.bounds[1] // self.downsample_factor,
+                ),
+            )
+
+        return merged_image
+
+
 class RslearnWriter(BasePredictionWriter):
     """A writer that writes predictions back into the rslearn dataset.
 
@@ -46,9 +160,12 @@ class RslearnWriter(BasePredictionWriter):
         self,
         path: str,
         output_layer: str,
-        path_options: dict[str, Any] = {},
-        selector: list[str] = [],
+        path_options: dict[str, Any] | None = None,
+        selector: list[str] | None = None,
         merger: PatchPredictionMerger | None = None,
+        output_path: str | Path | None = None,
+        layer_config: LayerConfig | None = None,
+        storage_config: StorageConfig | None = None,
     ):
         """Create a new RslearnWriter.
 
@@ -57,42 +174,125 @@ class RslearnWriter(BasePredictionWriter):
             output_layer: which layer to write the outputs under.
             path_options: additional options for path to pass to fsspec
             selector: keys to access the desired output in the output dict if needed.
+                e.g ["key1", "key2"] gets output["key1"]["key2"]
             merger: merger to use to merge outputs from overlapped patches.
+            output_path: optional custom path for writing predictions. If provided,
+                predictions will be written to this path instead of deriving from dataset path.
+            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
-        self.path = UPath(path, **path_options)
-        self.dataset = Dataset(self.path)
-        self.layer_config = self.dataset.layers[self.output_layer]
+        self.selector = selector or []
+        ds_upath = UPath(path, **path_options or {})
+        output_upath = (
+            UPath(output_path, **path_options or {})
+            if output_path is not None
+            else ds_upath
+        )
 
-        if self.layer_config.layer_type == LayerType.RASTER:
-            band_cfg = self.layer_config.band_sets[0]
-            self.format = load_raster_format(
-                RasterFormatConfig(band_cfg.format["name"], band_cfg.format)
+        self.layer_config, self.dataset_storage = (
+            self._get_layer_config_and_dataset_storage(
+                ds_upath, output_upath, layer_config, storage_config
             )
-        elif self.layer_config.layer_type == LayerType.VECTOR:
-            self.format = load_vector_format(self.layer_config.format)
+        )
+
+        self.format: RasterFormat | VectorFormat
+        if self.layer_config.type == LayerType.RASTER:
+            band_cfg = self.layer_config.band_sets[0]
+            self.format = band_cfg.instantiate_raster_format()
+        elif self.layer_config.type == LayerType.VECTOR:
+            self.format = self.layer_config.instantiate_vector_format()
         else:
-            raise ValueError(f"invalid layer type {self.layer_config.layer_type}")
+            raise ValueError(f"invalid layer type {self.layer_config.type}")
 
-        self.merger = merger
+        if merger is not None:
+            self.merger = merger
+        elif self.layer_config.type == LayerType.RASTER:
+            self.merger = RasterMerger()
+        elif self.layer_config.type == LayerType.VECTOR:
+            self.merger = VectorMerger()
 
         # Map from window name to pending data to write.
         # This is used when windows are split up into patches, so the data from all the
         # patches of each window need to be reconstituted.
-        self.pending_outputs = {}
+        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: Sequence[Any],
-        batch_indices: Sequence[Any],
-        batch: Any,
+        prediction: ModelOutput,
+        batch_indices: Sequence[int] | None,
+        batch: tuple[list, list, list],
         batch_idx: int,
         dataloader_idx: int,
-    ):
+    ) -> None:
         """Write a batch of predictions into the rslearn dataset.
 
         Args:
@@ -100,14 +300,38 @@ class RslearnWriter(BasePredictionWriter):
             pl_module: the LightningModule.
             prediction: the prediction to write.
             batch_indices: batch indices.
-            batch: the batch that was input to the model.
+            batch: the batch that was input to the model. It should be a list of
+                (inputs, targets, metadatas).
             batch_idx: the batch index.
             dataloader_idx: the index in the dataloader.
         """
         assert isinstance(pl_module, RslearnLightningModule)
-        metadatas = batch[2]
-        outputs = [
-            pl_module.task.process_output(output, metadata)
+        task = pl_module.task
+        _, _, metadatas = batch
+        self.process_output_batch(task, prediction.outputs, metadatas)
+
+    def process_output_batch(
+        self,
+        task: Task,
+        prediction: Iterable[Any],
+        metadatas: Iterable[SampleMetadata],
+    ) -> None:
+        """Write a prediction batch with simplified API.
+
+        write_on_batch_end wraps this function to work with lightning API, but only a
+        subset of arguments are used.
+
+        Args:
+            task: the Task that we are writing outputs for.
+            prediction: the list of predictions in this batch to write. These outputs
+                will be processed by the task to obtain a vector (list[Feature]) or
+                raster (npt.NDArray) output.
+            metadatas: corresponding list of metadatas from the batch describing the
+                patches that were processed.
+        """
+        # Process the predictions into outputs that can be written.
+        outputs: list = [
+            task.process_output(output, metadata)
             for output, metadata in zip(prediction, metadatas)
         ]
 
@@ -115,64 +339,75 @@ class RslearnWriter(BasePredictionWriter):
             for k in self.selector:
                 output = output[k]
 
-            window_name = metadata["window_name"]
-            cur_bounds = metadata["bounds"]
-            window_bounds = metadata["window_bounds"]
-
-            if self.layer_config.layer_type == LayerType.RASTER:
-                if window_name not in self.pending_outputs:
-                    self.pending_outputs[window_name] = np.zeros(
-                        (
-                            output.shape[0],
-                            window_bounds[3] - window_bounds[1],
-                            window_bounds[2] - window_bounds[0],
-                        ),
-                        dtype=output.dtype,
-                    )
+            window = Window(
+                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_in_window,
+                metadata.patch_bounds,
+                output,
+            )
 
-                # Use copy_spatial_array to handle the copy since, when using patches,
-                # the last column/row of outputs might extend beyond the bounds of the
-                # window.
-                copy_spatial_array(
-                    src=output,
-                    dst=self.pending_outputs[window_name],
-                    src_offset=(cur_bounds[0], cur_bounds[1]),
-                    dst_offset=(window_bounds[0], window_bounds[1]),
-                )
+    def process_output(
+        self,
+        window: Window,
+        patch_idx: int,
+        num_patches: int,
+        cur_bounds: PixelBounds,
+        output: npt.NDArray | list[Feature],
+    ) -> None:
+        """Process one output from the model.
 
-            elif self.layer_config.layer_type == LayerType.VECTOR:
-                if window_name not in self.pending_outputs:
-                    self.pending_outputs[window_name] = []
+        Args:
+            window: the window that the output pertains to.
+            patch_idx: the index of this patch for the window.
+            num_patches: the total number of patches to be processed for the window.
+            cur_bounds: the bounds of the current patch.
+            output: the output data.
+        """
+        # Incorporate the output into our list of pending patch outputs.
+        if window.name not in self.pending_outputs:
+            self.pending_outputs[window.name] = []
+        self.pending_outputs[window.name].append(PendingPatchOutput(cur_bounds, output))
+        logger.debug(
+            f"Stored PendingPatchOutput for patch #{patch_idx}/{num_patches} at window {window.name}"
+        )
 
-                self.pending_outputs[window_name].extend(output)
+        if patch_idx < num_patches - 1:
+            return
 
-            if metadata["patch_idx"] < metadata["num_patches"] - 1:
-                continue
+        # This is the last patch so it's time to write it.
+        # First get the pending output and clear it.
+        pending_output = self.pending_outputs[window.name]
+        del self.pending_outputs[window.name]
 
-            pending_output = self.pending_outputs[window_name]
-            del self.pending_outputs[window_name]
+        # 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, self.layer_config)
 
-            # This is the last patch so it's time to merge outputs from overlapped patches
-            if self.merger is not None:
-                pending_output = self.merger.merge(pending_output)
+        if self.layer_config.type == LayerType.RASTER:
+            raster_dir = window.get_raster_dir(
+                self.output_layer, self.layer_config.band_sets[0].bands
+            )
+            assert isinstance(self.format, RasterFormat)
 
-            # This is the last patch so it's time to write it.
-            layer_dir = (
-                self.dataset.path
-                / "windows"
-                / metadata["group"]
-                / window_name
-                / "layers"
-                / self.output_layer
+            # In case the merged_output is at a different resolution than the window,
+            # get adjusted projection and bounds for writing it.
+            projection, bounds = adjust_projection_and_bounds_for_array(
+                window.projection, window.bounds, merged_output
             )
+            self.format.encode_raster(raster_dir, projection, bounds, merged_output)
 
-            if self.layer_config.layer_type == LayerType.RASTER:
-                band_dir = layer_dir / "_".join(self.layer_config.band_sets[0].bands)
-                self.format.encode_raster(
-                    band_dir, metadata["projection"], window_bounds, pending_output
-                )
+        elif self.layer_config.type == LayerType.VECTOR:
+            layer_dir = window.get_layer_dir(self.output_layer)
+            assert isinstance(self.format, VectorFormat)
+            self.format.encode_vector(layer_dir, merged_output)
 
-            elif self.layer_config.layer_type == LayerType.VECTOR:
-                self.format.encode_vector(
-                    layer_dir, metadata["projection"], pending_output
-                )
+        window.mark_layer_completed(self.output_layer)

rslearn/train/scheduler.py

@@ -0,0 +1,92 @@
+"""Learning rate schedulers for rslearn."""
+
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import (
+    CosineAnnealingLR,
+    CosineAnnealingWarmRestarts,
+    LRScheduler,
+    MultiStepLR,
+    ReduceLROnPlateau,
+)
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class SchedulerFactory(ABC):
+    """A factory class that initializes an LR scheduler given the optimizer."""
+
+    def get_kwargs(self) -> dict:
+        """Get the keyword arguments for the scheduler."""
+        return {k: v for k, v in asdict(self).items() if v is not None}  # type: ignore
+
+    @abstractmethod
+    def build(self, optimizer: Optimizer) -> LRScheduler:
+        """Build the learning rate scheduler configured by this factory class."""
+        logger.info(
+            f"Using scheduler {self.__class__.__name__} with kwargs {self.get_kwargs()}"
+        )
+
+
+@dataclass
+class PlateauScheduler(SchedulerFactory):
+    """Plateau learning rate scheduler."""
+
+    mode: str | None = None
+    factor: float | None = None
+    patience: int | None = None
+    threshold: float | None = None
+    threshold_mode: str | None = None
+    cooldown: int | None = None
+    min_lr: float | None = None
+    eps: float | None = None
+
+    def build(self, optimizer: Optimizer) -> LRScheduler:
+        """Build the ReduceLROnPlateau scheduler."""
+        super().build(optimizer)
+        return ReduceLROnPlateau(optimizer, **self.get_kwargs())
+
+
+@dataclass
+class MultiStepScheduler(SchedulerFactory):
+    """Step learning rate scheduler."""
+
+    milestones: list[int]
+    gamma: float | None = None
+    last_epoch: int | None = None
+
+    def build(self, optimizer: Optimizer) -> LRScheduler:
+        """Build the ReduceLROnPlateau scheduler."""
+        super().build(optimizer)
+        return MultiStepLR(optimizer, **self.get_kwargs())
+
+
+@dataclass
+class CosineAnnealingScheduler(SchedulerFactory):
+    """Cosine annealing learning rate scheduler."""
+
+    T_max: int
+    eta_min: float | None = None
+
+    def build(self, optimizer: Optimizer) -> LRScheduler:
+        """Build the CosineAnnealingLR scheduler."""
+        super().build(optimizer)
+        return CosineAnnealingLR(optimizer, **self.get_kwargs())
+
+
+@dataclass
+class CosineAnnealingWarmRestartsScheduler(SchedulerFactory):
+    """Cosine annealing with warm restarts learning rate scheduler."""
+
+    T_0: int
+    T_mult: int = 1
+    eta_min: float = 0.0
+
+    def build(self, optimizer: Optimizer) -> LRScheduler:
+        """Build the CosineAnnealingWarmRestarts scheduler."""
+        super().build(optimizer)
+        return CosineAnnealingWarmRestarts(optimizer, **self.get_kwargs())