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

rslearn/models/terramind.py

@@ -8,8 +8,11 @@ import torch.nn.functional as F
 from einops import rearrange
 from terratorch.registry import BACKBONE_REGISTRY
 
+from rslearn.train.model_context import ModelContext
 from rslearn.train.transforms.transform import Transform
 
+from .component import FeatureExtractor, FeatureMaps
+
 
 # TerraMind v1 provides two sizes: base and large
 class TerramindSize(str, Enum):
@@ -85,7 +88,7 @@ PRETRAINED_BANDS = {
 }
 
 
-class Terramind(torch.nn.Module):
+class Terramind(FeatureExtractor):
     """Terramind backbones."""
 
     def __init__(
@@ -123,21 +126,25 @@ class Terramind(torch.nn.Module):
         self.modalities = modalities
         self.do_resizing = do_resizing
 
-    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+    def forward(self, context: ModelContext) -> FeatureMaps:
         """Forward pass for the Terramind model.
 
         Args:
-            inputs: input dicts that must include modalities as keys which are defined in the self.modalities list
+            context: the model context. Input dicts must include modalities as keys
+                which are defined in the self.modalities list.
 
         Returns:
-            List[torch.Tensor]: Single-scale feature tensors from the encoder.
+            a FeatureMaps with one feature map from the encoder, at 1/16 of the input
+                resolution.
         """
         model_inputs = {}
         for modality in self.modalities:
             # We assume the all the inputs include the same modalities
-            if modality not in inputs[0]:
+            if modality not in context.inputs[0]:
                 continue
-            cur = torch.stack([inp[modality] for inp in inputs], dim=0)  # (B, C, H, W)
+            cur = torch.stack(
+                [inp[modality] for inp in context.inputs], dim=0
+            )  # (B, C, H, W)
             if self.do_resizing and (
                 cur.shape[2] != IMAGE_SIZE or cur.shape[3] != IMAGE_SIZE
             ):
@@ -159,7 +166,17 @@ class Terramind(torch.nn.Module):
         image_features = self.model(model_inputs)[-1]
         batch_size, num_patches, _ = image_features.shape
         height, width = int(num_patches**0.5), int(num_patches**0.5)
-        return [rearrange(image_features, "b (h w) d -> b d h w", h=height, w=width)]
+        return FeatureMaps(
+            [
+                rearrange(
+                    image_features,
+                    "b (h w) d -> b d h w",
+                    b=batch_size,
+                    h=height,
+                    w=width,
+                )
+            ]
+        )
 
     def get_backbone_channels(self) -> list:
         """Returns the output channels of this model when used as a backbone.

rslearn/models/trunk.py

@@ -7,6 +7,7 @@ import torch
 
 from rslearn.log_utils import get_logger
 from rslearn.models.task_embedding import BaseTaskEmbedding
+from rslearn.train.model_context import ModelOutput
 
 logger = get_logger(__name__)
 
@@ -32,10 +33,11 @@ class DecoderTrunkLayer(torch.nn.Module, ABC):
             dict with key "outputs" (output tensor of shape (batch_size, seq_len, dim))
             and optionally other keys.
         """
+        raise NotImplementedError
 
     @abstractmethod
     def apply_auxiliary_losses(
-        self, trunk_out: dict[str, Any], outs: dict[str, Any]
+        self, trunk_out: dict[str, Any], outs: ModelOutput
     ) -> None:
         """Apply auxiliary losses in-place.
 
@@ -43,6 +45,7 @@ class DecoderTrunkLayer(torch.nn.Module, ABC):
             trunk_out: The output of the trunk.
             outs: The output of the decoders, with key "loss_dict" containing the losses.
         """
+        raise NotImplementedError
 
 
 class DecoderTrunk(torch.nn.Module):
@@ -122,7 +125,7 @@ class DecoderTrunk(torch.nn.Module):
         return out
 
     def apply_auxiliary_losses(
-        self, trunk_out: dict[str, Any], outs: dict[str, Any]
+        self, trunk_out: dict[str, Any], outs: ModelOutput
     ) -> None:
         """Apply auxiliary losses in-place.
 
@@ -130,7 +133,7 @@ class DecoderTrunk(torch.nn.Module):
 
         Args:
             trunk_out: The output of the trunk.
-            outs: The output of the decoders, with key "loss_dict" containing the losses.
+            outs: The output of the decoders.
         """
         for layer in self.layers:
             layer.apply_auxiliary_losses(trunk_out, outs)

rslearn/models/unet.py

@@ -5,8 +5,15 @@ from typing import Any
 import torch
 import torch.nn.functional as F
 
+from rslearn.train.model_context import ModelContext
 
-class UNetDecoder(torch.nn.Module):
+from .component import (
+    FeatureMaps,
+    IntermediateComponent,
+)
+
+
+class UNetDecoder(IntermediateComponent):
     """UNet-style decoder.
 
     It inputs multi-scale features. Starting from last (lowest resolution) feature map,
@@ -143,23 +150,25 @@ class UNetDecoder(torch.nn.Module):
             align_corners=False,
         )
 
-    def forward(
-        self, in_features: list[torch.Tensor], inputs: list[dict[str, Any]]
-    ) -> torch.Tensor:
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
         """Compute output from multi-scale feature map.
 
         Args:
-            in_features: list of feature maps at different resolutions.
-            inputs: original inputs (ignored).
+            intermediates: the output from the previous model component, which must be a FeatureMaps.
+            context: the model context.
 
         Returns:
-            output image
+            output FeatureMaps consisting of one map. The embedding size is equal to the
+                configured out_channels.
         """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to UNetDecoder must be a FeatureMaps")
+
         # Reverse the features since we will pass them in from lowest resolution to highest.
-        in_features = list(reversed(in_features))
+        in_features = list(reversed(intermediates.feature_maps))
         cur_features = self.layers[0](in_features[0])
         for in_feat, layer in zip(in_features[1:], self.layers[1:]):
             cur_features = layer(torch.cat([cur_features, in_feat], dim=1))
         if self.original_size_to_interpolate is not None:
             cur_features = self._resize(cur_features)
-        return cur_features
+        return FeatureMaps([cur_features])

rslearn/models/upsample.py

@@ -1,9 +1,18 @@
 """An upsampling layer."""
 
+from typing import Any
+
 import torch
 
+from rslearn.train.model_context import ModelContext
+
+from .component import (
+    FeatureMaps,
+    IntermediateComponent,
+)
 
-class Upsample(torch.nn.Module):
+
+class Upsample(IntermediateComponent):
     """Upsamples each input feature map by the same factor."""
 
     def __init__(
@@ -20,16 +29,20 @@ class Upsample(torch.nn.Module):
         super().__init__()
         self.layer = torch.nn.Upsample(scale_factor=scale_factor, mode=mode)
 
-    def forward(
-        self, features: list[torch.Tensor], inputs: list[torch.Tensor]
-    ) -> list[torch.Tensor]:
-        """Upsample each feature map.
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Upsample each feature map by scale_factor.
 
         Args:
-            features: list of feature maps at different resolutions.
-            inputs: original inputs (ignored).
+            intermediates: the output from the previous component, which must be a FeatureMaps.
+            context: the model context.
 
         Returns:
-            upsampled feature maps
+            upsampled feature maps.
         """
-        return [self.layer(feat_map) for feat_map in features]
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to Upsample must be a FeatureMaps")
+
+        upsampled_feat_maps = [
+            self.layer(feat_map) for feat_map in intermediates.feature_maps
+        ]
+        return FeatureMaps(upsampled_feat_maps)

rslearn/train/all_patches_dataset.py

@@ -2,6 +2,7 @@
 
 import itertools
 from collections.abc import Iterable, Iterator
+from dataclasses import replace
 from typing import Any
 
 import shapely
@@ -9,6 +10,7 @@ import torch
 
 from rslearn.dataset import Window
 from rslearn.train.dataset import ModelDataset
+from rslearn.train.model_context import SampleMetadata
 from rslearn.utils.geometry import PixelBounds, STGeometry
 
 
@@ -218,7 +220,7 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
 
     def __iter__(
         self,
-    ) -> Iterator[tuple[dict[str, Any], dict[str, Any], dict[str, Any]]]:
+    ) -> Iterator[tuple[dict[str, Any], dict[str, Any], SampleMetadata]]:
         """Iterate over all patches in each element of the underlying ModelDataset."""
         # Iterate over the window IDs until we have returned enough samples.
         window_ids, num_samples_needed = self._get_worker_iteration_data()
@@ -229,7 +231,7 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
                 raw_inputs, passthrough_inputs, metadata = self.dataset.get_raw_inputs(
                     window_id
                 )
-                bounds = metadata["bounds"]
+                bounds = metadata.patch_bounds
 
                 # For simplicity, pad tensors by patch size to ensure that any patch bounds
                 # extending outside the window bounds will not have issues when we slice
@@ -244,7 +246,7 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
                 )
                 for patch_idx, patch_bounds in enumerate(patches):
                     cur_geom = STGeometry(
-                        metadata["projection"], shapely.box(*patch_bounds), None
+                        metadata.projection, shapely.box(*patch_bounds), None
                     )
                     start_offset = (
                         patch_bounds[0] - bounds[0],
@@ -282,10 +284,12 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
                     cur_passthrough_inputs = crop_input_dict(passthrough_inputs)
 
                     # Adjust the metadata as well.
-                    cur_metadata = metadata.copy()
-                    cur_metadata["bounds"] = patch_bounds
-                    cur_metadata["patch_idx"] = patch_idx
-                    cur_metadata["num_patches"] = len(patches)
+                    cur_metadata = replace(
+                        metadata,
+                        patch_bounds=patch_bounds,
+                        patch_idx=patch_idx,
+                        num_patches_in_window=len(patches),
+                    )
 
                     # Now we can compute input and target dicts via the task.
                     input_dict, target_dict = self.dataset.task.process_inputs(
@@ -297,7 +301,6 @@ class IterableAllPatchesDataset(torch.utils.data.IterableDataset):
                     input_dict, target_dict = self.dataset.transforms(
                         input_dict, target_dict
                     )
-                    input_dict["dataset_source"] = self.dataset.name
 
                     if num_samples_returned < num_samples_needed:
                         yield input_dict, target_dict, cur_metadata
@@ -346,7 +349,7 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
         )
         self.windows = self.dataset.get_dataset_examples()
         self.window_cache: dict[
-            int, tuple[dict[str, Any], dict[str, Any], dict[str, Any]]
+            int, tuple[dict[str, Any], dict[str, Any], SampleMetadata]
         ] = {}
 
         # Precompute the batch boundaries for each window
@@ -360,7 +363,7 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
 
     def get_raw_inputs(
         self, index: int
-    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], dict[str, Any], SampleMetadata]:
         """Get the raw inputs for a single patch. Retrieve from cache if possible.
 
         Also crops/pads the tensors by patch size to protect slicing near right/bottom edges.
@@ -410,13 +413,13 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
 
     def __getitem__(
         self, index: int
-    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], dict[str, Any], SampleMetadata]:
         """Return (input_dict, target_dict, metadata) for a single flattened patch."""
         (window_id, patch_bounds, (patch_idx, num_patches)) = self.patches[index]
         raw_inputs, passthrough_inputs, metadata = self.get_raw_inputs(window_id)
-        bounds = metadata["bounds"]
+        bounds = metadata.patch_bounds
 
-        cur_geom = STGeometry(metadata["projection"], shapely.box(*patch_bounds), None)
+        cur_geom = STGeometry(metadata.projection, shapely.box(*patch_bounds), None)
         start_offset = (patch_bounds[0] - bounds[0], patch_bounds[1] - bounds[1])
         end_offset = (patch_bounds[2] - bounds[0], patch_bounds[3] - bounds[1])
 
@@ -428,10 +431,12 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
         )
 
         # Adjust the metadata as well.
-        cur_metadata = metadata.copy()
-        cur_metadata["bounds"] = patch_bounds
-        cur_metadata["patch_idx"] = patch_idx
-        cur_metadata["num_patches"] = num_patches
+        cur_metadata = replace(
+            metadata,
+            patch_bounds=patch_bounds,
+            patch_idx=patch_idx,
+            num_patches_in_window=num_patches,
+        )
 
         # Now we can compute input and target dicts via the task.
         input_dict, target_dict = self.dataset.task.process_inputs(
@@ -441,7 +446,6 @@ class InMemoryAllPatchesDataset(torch.utils.data.Dataset):
         )
         input_dict.update(cur_passthrough_inputs)
         input_dict, target_dict = self.dataset.transforms(input_dict, target_dict)
-        input_dict["dataset_source"] = self.dataset.name
 
         return input_dict, target_dict, cur_metadata
 

rslearn/train/dataset.py

@@ -20,13 +20,15 @@ from rslearn.config import (
     LayerConfig,
 )
 from rslearn.dataset.dataset import Dataset
+from rslearn.dataset.storage.file import FileWindowStorage
 from rslearn.dataset.window import Window, get_layer_and_group_from_dir_name
 from rslearn.log_utils import get_logger
-from rslearn.train.tasks import Task
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.mp import star_imap_unordered
 
+from .model_context import SampleMetadata
+from .tasks import Task
 from .transforms import Sequential
 
 logger = get_logger(__name__)
@@ -575,37 +577,7 @@ class ModelDataset(torch.utils.data.Dataset):
         else:
             self.patch_size = split_config.get_patch_size()
 
-        if split_config.names:
-            windows = self.dataset.load_windows(
-                groups=split_config.groups,
-                names=split_config.names,
-                show_progress=True,
-                workers=workers,
-            )
-        elif split_config.groups:
-            windows = self.dataset.load_windows(
-                groups=split_config.groups, show_progress=True, workers=workers
-            )
-        else:
-            windows = self.dataset.load_windows(show_progress=True, workers=workers)
-
-        if split_config.tags:
-            # Filter the window.options.
-            new_windows = []
-            num_removed: dict[str, int] = {}
-            for window in windows:
-                for k, v in split_config.tags.items():
-                    if k not in window.options or (v and window.options[k] != v):
-                        num_removed[k] = num_removed.get(k, 0) + 1
-                        break
-                else:
-                    new_windows.append(window)
-            logger.info(
-                f"Started with {len(windows)} windows, ended with {len(new_windows)} windows for {self.dataset.path}"
-            )
-            for k, v in num_removed.items():
-                logger.info(f"Removed {v} windows due to tag {k}")
-            windows = new_windows
+        windows = self._get_initial_windows(split_config, workers)
 
         # If targets are not needed, remove them from the inputs.
         if split_config.get_skip_targets():
@@ -615,17 +587,11 @@ class ModelDataset(torch.utils.data.Dataset):
 
         # Eliminate windows that are missing either a requisite input layer, or missing
         # all target layers.
-        # We use only main thread if the index is set, since that can take a long time
-        # to send to the worker threads, it may get serialized for each window.
         new_windows = []
-        if workers == 0 or (len(windows) >= 1 and windows[0].index is not None):
+        if workers == 0:
             for window in windows:
                 if check_window(self.inputs, window) is None:
                     continue
-                # The index may be set, but now that this check is done, from here on
-                # we no longer need it. We set it None so that we don't end up passing
-                # it later to the dataloader workers.
-                window.index = None
                 new_windows.append(window)
         else:
             p = multiprocessing.Pool(workers)
@@ -681,12 +647,62 @@ class ModelDataset(torch.utils.data.Dataset):
         with open(self.dataset_examples_fname, "w") as f:
             json.dump([self._serialize_item(example) for example in windows], f)
 
+    def _get_initial_windows(
+        self, split_config: SplitConfig, workers: int
+    ) -> list[Window]:
+        """Get the initial windows before input layer filtering.
+
+        The windows are filtered based on configured window names, groups, and tags.
+
+        This is a helper for the init function.
+
+        Args:
+            split_config: the split configuration.
+            workers: number of worker processes.
+
+        Returns:
+            list of windows from the dataset after applying the aforementioned filters.
+        """
+        # Load windows from dataset.
+        # If the window storage is FileWindowStorage, we pass the workers/show_progress arguments.
+        kwargs: dict[str, Any] = {}
+        if isinstance(self.dataset.storage, FileWindowStorage):
+            kwargs["workers"] = workers
+            kwargs["show_progress"] = True
+        # We also add the name/group filters to the kwargs.
+        if split_config.names:
+            kwargs["names"] = split_config.names
+        if split_config.groups:
+            kwargs["groups"] = split_config.groups
+
+        windows = self.dataset.load_windows(**kwargs)
+
+        # Filter by tags (if provided) using the window.options.
+        if split_config.tags:
+            new_windows = []
+            num_removed: dict[str, int] = {}
+            for window in windows:
+                for k, v in split_config.tags.items():
+                    if k not in window.options or (v and window.options[k] != v):
+                        num_removed[k] = num_removed.get(k, 0) + 1
+                        break
+                else:
+                    new_windows.append(window)
+            logger.info(
+                f"Started with {len(windows)} windows, ended with {len(new_windows)} windows for {self.dataset.path}"
+            )
+            for k, v in num_removed.items():
+                logger.info(f"Removed {v} windows due to tag {k}")
+            windows = new_windows
+
+        return windows
+
     def _serialize_item(self, example: Window) -> dict[str, Any]:
         return example.get_metadata()
 
     def _deserialize_item(self, d: dict[str, Any]) -> Window:
         return Window.from_metadata(
-            Window.get_window_root(self.dataset.path, d["group"], d["name"]),
+            self.dataset.storage,
             d,
         )
 
@@ -713,7 +729,7 @@ class ModelDataset(torch.utils.data.Dataset):
 
     def get_raw_inputs(
         self, idx: int
-    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], dict[str, Any], SampleMetadata]:
         """Get the raw inputs and base metadata for this example.
 
         This is the raster or vector data before being processed by the Task. So it
@@ -775,21 +791,23 @@ class ModelDataset(torch.utils.data.Dataset):
             if data_input.passthrough:
                 passthrough_inputs[name] = raw_inputs[name]
 
-        metadata = {
-            "group": window.group,
-            "window_name": window.name,
-            "window_bounds": window.bounds,
-            "bounds": bounds,
-            "time_range": window.time_range,
-            "projection": window.projection,
-            "dataset_source": self.name,
-        }
+        metadata = SampleMetadata(
+            window_group=window.group,
+            window_name=window.name,
+            window_bounds=window.bounds,
+            patch_bounds=bounds,
+            patch_idx=0,
+            num_patches_in_window=1,
+            time_range=window.time_range,
+            projection=window.projection,
+            dataset_source=self.name,
+        )
 
         return raw_inputs, passthrough_inputs, metadata
 
     def __getitem__(
         self, idx: int
-    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
+    ) -> tuple[dict[str, Any], dict[str, Any], SampleMetadata]:
         """Read one training example.
 
         Args:
@@ -801,8 +819,6 @@ class ModelDataset(torch.utils.data.Dataset):
         logger.debug("__getitem__ start pid=%d item_idx=%d", os.getpid(), idx)
 
         raw_inputs, passthrough_inputs, metadata = self.get_raw_inputs(idx)
-        metadata["patch_idx"] = 0
-        metadata["num_patches"] = 1
 
         input_dict, target_dict = self.task.process_inputs(
             raw_inputs,
@@ -811,7 +827,6 @@ class ModelDataset(torch.utils.data.Dataset):
         )
         input_dict.update(passthrough_inputs)
         input_dict, target_dict = self.transforms(input_dict, target_dict)
-        input_dict["dataset_source"] = self.name
 
         logger.debug("__getitem__ finish pid=%d item_idx=%d", os.getpid(), idx)