rslearn 0.0.24__py3-none-any.whl → 0.0.26__py3-none-any.whl

rslearn/models/simple_time_series.py

@@ -1,5 +1,6 @@
 """SimpleTimeSeries encoder."""
 
+import warnings
 from typing import Any
 
 import torch
@@ -25,13 +26,14 @@ class SimpleTimeSeries(FeatureExtractor):
     def __init__(
         self,
         encoder: FeatureExtractor,
-        image_channels: int | None = None,
+        num_timesteps_per_forward_pass: int = 1,
         op: str = "max",
         groups: list[list[int]] | None = None,
         num_layers: int | None = None,
         image_key: str = "image",
         backbone_channels: list[tuple[int, int]] | None = None,
-        image_keys: dict[str, int] | None = None,
+        image_keys: list[str] | dict[str, int] | None = None,
+        image_channels: int | None = None,
     ) -> None:
         """Create a new SimpleTimeSeries.
 
@@ -39,9 +41,11 @@ class SimpleTimeSeries(FeatureExtractor):
             encoder: the underlying FeatureExtractor. It must provide get_backbone_channels
                 function that returns the output channels, or backbone_channels must be set.
                 It must output a FeatureMaps.
-            image_channels: the number of channels per image of the time series. The
-                input should have multiple images concatenated on the channel axis, so
-                this parameter is used to distinguish the different images.
+            num_timesteps_per_forward_pass: how many timesteps to pass to the encoder
+                in each forward pass. Defaults to 1 (one timestep per forward pass).
+                Set to a higher value to batch multiple timesteps together, e.g. for
+                pre/post change detection where you want 4 pre and 4 post images
+                processed together.
             op: one of max, mean, convrnn, conv3d, or conv1d
             groups: sets of images for which to combine features. Within each set,
                 features are combined using the specified operation; then, across sets,
@@ -51,28 +55,53 @@ class SimpleTimeSeries(FeatureExtractor):
                 combined before features and the combined after features. groups is a
                 list of sets, and each set is a list of image indices.
             num_layers: the number of layers for convrnn, conv3d, and conv1d ops.
-            image_key: the key to access the images.
+            image_key: the key to access the images (used when image_keys is not set).
             backbone_channels: manually specify the backbone channels. Can be set if
                 the encoder does not provide get_backbone_channels function.
-            image_keys: as an alternative to setting image_channels, map from the key
-                in input dict to the number of channels per timestep for that modality.
-                This way SimpleTimeSeries can be used with multimodal inputs. One of
-                image_channels or image_keys must be specified.
+            image_keys: list of keys in input dict to process as multimodal inputs.
+                All keys use the same num_timesteps_per_forward_pass. If not set,
+                only the single image_key is used. Passing a dict[str, int] is
+                deprecated and will be removed on 2026-04-01.
+            image_channels: Deprecated, use num_timesteps_per_forward_pass instead.
+                Will be removed on 2026-04-01.
         """
-        if (image_channels is None and image_keys is None) or (
-            image_channels is not None and image_keys is not None
-        ):
-            raise ValueError(
-                "exactly one of image_channels and image_keys must be specified"
+        # Handle deprecated image_channels parameter
+        if image_channels is not None:
+            warnings.warn(
+                "image_channels is deprecated and will be removed on 2026-04-01. "
+                "Use num_timesteps_per_forward_pass instead. The new parameter directly "
+                "specifies the number of timesteps per forward pass rather than requiring "
+                "image_channels // actual_channels.",
+                FutureWarning,
+                stacklevel=2,
             )
 
+        # Handle deprecated dict form of image_keys
+        deprecated_image_keys_dict: dict[str, int] | None = None
+        if isinstance(image_keys, dict):
+            warnings.warn(
+                "Passing image_keys as a dict is deprecated and will be removed on "
+                "2026-04-01. Use image_keys as a list[str] and set "
+                "num_timesteps_per_forward_pass instead.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            deprecated_image_keys_dict = image_keys
+            image_keys = None  # Will use deprecated path in forward
+
         super().__init__()
         self.encoder = encoder
-        self.image_channels = image_channels
+        self.num_timesteps_per_forward_pass = num_timesteps_per_forward_pass
+        # Store deprecated parameters for runtime conversion
+        self._deprecated_image_channels = image_channels
+        self._deprecated_image_keys_dict = deprecated_image_keys_dict
         self.op = op
         self.groups = groups
-        self.image_key = image_key
-        self.image_keys = image_keys
+        # Normalize image_key to image_keys list form
+        if image_keys is not None:
+            self.image_keys = image_keys
+        else:
+            self.image_keys = [image_key]
 
         if backbone_channels is not None:
             out_channels = backbone_channels
@@ -163,24 +192,25 @@ class SimpleTimeSeries(FeatureExtractor):
         return out_channels
 
     def _get_batched_images(
-        self, input_dicts: list[dict[str, Any]], image_key: str, image_channels: int
+        self, input_dicts: list[dict[str, Any]], image_key: str, num_timesteps: int
     ) -> list[RasterImage]:
         """Collect and reshape images across input dicts.
 
         The BTCHW image time series are reshaped to (B*T)CHW so they can be passed to
         the forward pass of a per-image (unitemporal) model.
+
+        Args:
+            input_dicts: list of input dictionaries containing RasterImage objects.
+            image_key: the key to access the RasterImage in each input dict.
+            num_timesteps: how many timesteps to batch together per forward pass.
         """
         images = torch.stack(
             [input_dict[image_key].image for input_dict in input_dicts], dim=0
         )  # B, C, T, H, W
         timestamps = [input_dict[image_key].timestamps for input_dict in input_dicts]
-        # if image channels is not equal to the actual number of channels, then
-        # then every N images should be batched together. For example, if the
-        # number of input channels c == 2, and image_channels == 4, then we
-        # want to pass 2 timesteps to the model.
-        # TODO is probably to make this behaviour clearer but lets leave it like
-        # this for now to not break things.
-        num_timesteps = image_channels // images.shape[1]
+        # num_timesteps specifies how many timesteps to batch together per forward pass.
+        # For example, if the input has 8 timesteps and num_timesteps=4, we do 2
+        # forward passes, each with 4 timesteps batched together.
         batched_timesteps = images.shape[2] // num_timesteps
         images = rearrange(
             images,
@@ -222,10 +252,22 @@ class SimpleTimeSeries(FeatureExtractor):
         n_batch = len(context.inputs)
         n_images: int | None = None
 
-        if self.image_keys is not None:
-            for image_key, image_channels in self.image_keys.items():
+        if self._deprecated_image_keys_dict is not None:
+            # Deprecated dict form: each key has its own channels_per_timestep.
+            # The channels_per_timestep could be used to group multiple timesteps,
+            # together, so we need to divide by the actual image channel count to get
+            # the number of timesteps to be grouped.
+            for (
+                image_key,
+                channels_per_timestep,
+            ) in self._deprecated_image_keys_dict.items():
+                # For deprecated image_keys dict, the value is channels per timestep,
+                # so we need to compute num_timesteps from the actual image channels
+                sample_image = context.inputs[0][image_key].image
+                actual_channels = sample_image.shape[0]  # C in CTHW
+                num_timesteps = channels_per_timestep // actual_channels
                 batched_images = self._get_batched_images(
-                    context.inputs, image_key, image_channels
+                    context.inputs, image_key, num_timesteps
                 )
 
                 if batched_inputs is None:
@@ -240,12 +282,32 @@ class SimpleTimeSeries(FeatureExtractor):
                     batched_inputs[i][image_key] = image
 
         else:
-            assert self.image_channels is not None
-            batched_images = self._get_batched_images(
-                context.inputs, self.image_key, self.image_channels
-            )
-            batched_inputs = [{self.image_key: image} for image in batched_images]
-            n_images = len(batched_images) // n_batch
+            # Determine num_timesteps - either from deprecated image_channels or
+            # directly from num_timesteps_per_forward_pass
+            if self._deprecated_image_channels is not None:
+                # Backwards compatibility: compute num_timesteps from image_channels
+                # (which should be a multiple of the actual per-timestep channels).
+                sample_image = context.inputs[0][self.image_keys[0]].image
+                actual_channels = sample_image.shape[0]  # C in CTHW
+                num_timesteps = self._deprecated_image_channels // actual_channels
+            else:
+                num_timesteps = self.num_timesteps_per_forward_pass
+
+            for image_key in self.image_keys:
+                batched_images = self._get_batched_images(
+                    context.inputs, image_key, num_timesteps
+                )
+
+                if batched_inputs is None:
+                    batched_inputs = [{} for _ in batched_images]
+                    n_images = len(batched_images) // n_batch
+                elif n_images != len(batched_images) // n_batch:
+                    raise ValueError(
+                        "expected all modalities to have the same number of timesteps"
+                    )
+
+                for i, image in enumerate(batched_images):
+                    batched_inputs[i][image_key] = image
 
         assert n_images is not None
         # Now we can apply the underlying FeatureExtractor.

rslearn/train/data_module.py

@@ -21,6 +21,7 @@ from .all_patches_dataset import (
 )
 from .dataset import (
     DataInput,
+    IndexMode,
     ModelDataset,
     MultiDataset,
     RetryDataset,
@@ -69,6 +70,7 @@ class RslearnDataModule(L.LightningDataModule):
         name: str | None = None,
         retries: int = 0,
         use_in_memory_all_patches_dataset: bool = False,
+        index_mode: IndexMode = IndexMode.OFF,
     ) -> None:
         """Initialize a new RslearnDataModule.
 
@@ -92,6 +94,7 @@ class RslearnDataModule(L.LightningDataModule):
             retries: number of retries to attempt for getitem calls
             use_in_memory_all_patches_dataset: whether to use InMemoryAllPatchesDataset
                 instead of IterableAllPatchesDataset if load_all_patches is set to true.
+            index_mode: controls dataset index caching behavior (default: IndexMode.OFF)
         """
         super().__init__()
         self.inputs = inputs
@@ -103,6 +106,7 @@ class RslearnDataModule(L.LightningDataModule):
         self.name = name
         self.retries = retries
         self.use_in_memory_all_patches_dataset = use_in_memory_all_patches_dataset
+        self.index_mode = index_mode
         self.split_configs = {
             "train": default_config.update(train_config),
             "val": default_config.update(val_config),
@@ -138,6 +142,7 @@ class RslearnDataModule(L.LightningDataModule):
                 workers=self.init_workers,
                 name=self.name,
                 fix_patch_pick=(split != "train"),
+                index_mode=self.index_mode,
             )
             logger.info(f"got {len(dataset)} examples in split {split}")
             if split_config.get_load_all_patches():

rslearn/train/dataset.py

@@ -9,6 +9,7 @@ import tempfile
 import time
 import uuid
 from datetime import datetime
+from enum import StrEnum
 from typing import Any
 
 import torch
@@ -29,6 +30,7 @@ from rslearn.dataset.window import (
     get_layer_and_group_from_dir_name,
 )
 from rslearn.log_utils import get_logger
+from rslearn.train.dataset_index import DatasetIndex
 from rslearn.train.model_context import RasterImage
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import PixelBounds, ResolutionFactor
@@ -41,6 +43,19 @@ from .transforms import Sequential
 logger = get_logger(__name__)
 
 
+class IndexMode(StrEnum):
+    """Controls dataset index caching behavior."""
+
+    OFF = "off"
+    """No caching - always load windows from dataset."""
+
+    USE = "use"
+    """Use cached index if available, create if not."""
+
+    REFRESH = "refresh"
+    """Ignore existing cache and rebuild."""
+
+
 def get_torch_dtype(dtype: DType) -> torch.dtype:
     """Convert rslearn DType to torch dtype."""
     if dtype == DType.INT32:
@@ -445,6 +460,7 @@ class SplitConfig:
         overlap_ratio: float | None = None,
         load_all_patches: bool | None = None,
         skip_targets: bool | None = None,
+        output_layer_name_skip_inference_if_exists: str | None = None,
     ) -> None:
         """Initialize a new SplitConfig.
 
@@ -467,6 +483,10 @@ class SplitConfig:
                 for each window, read all patches as separate sequential items in the
                 dataset.
             skip_targets: whether to skip targets when loading inputs
+            output_layer_name_skip_inference_if_exists: optional name of the output layer used during prediction.
+                If set, windows that already
+                have this layer completed will be skipped (useful for resuming
+                partial inference runs).
         """
         self.groups = groups
         self.names = names
@@ -477,6 +497,9 @@ class SplitConfig:
         self.sampler = sampler
         self.patch_size = patch_size
         self.skip_targets = skip_targets
+        self.output_layer_name_skip_inference_if_exists = (
+            output_layer_name_skip_inference_if_exists
+        )
 
         # Note that load_all_patches are handled by the RslearnDataModule rather than
         # the ModelDataset.
@@ -504,6 +527,7 @@ class SplitConfig:
             overlap_ratio=self.overlap_ratio,
             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,
         )
         if other.groups:
             result.groups = other.groups
@@ -527,6 +551,10 @@ class SplitConfig:
             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:
+            result.output_layer_name_skip_inference_if_exists = (
+                other.output_layer_name_skip_inference_if_exists
+            )
         return result
 
     def get_patch_size(self) -> tuple[int, int] | None:
@@ -549,16 +577,26 @@ class SplitConfig:
         """Returns whether skip_targets is enabled (default False)."""
         return True if self.skip_targets is True else False
 
+    def get_output_layer_name_skip_inference_if_exists(self) -> str | None:
+        """Returns output layer to use for resume checks (default None)."""
+        return self.output_layer_name_skip_inference_if_exists
+
 
-def check_window(inputs: dict[str, DataInput], window: Window) -> Window | None:
+def check_window(
+    inputs: dict[str, DataInput],
+    window: Window,
+    output_layer_name_skip_inference_if_exists: str | None = None,
+) -> Window | None:
     """Verify that the window has the required layers based on the specified inputs.
 
     Args:
         inputs: the inputs to the dataset.
         window: the window to check.
+        output_layer_name_skip_inference_if_exists: optional name of the output layer to check for existence.
 
     Returns:
-        the window if it has all the required inputs or None otherwise
+        the window if it has all the required inputs and does not need to be skipped
+        due to an existing output layer; or None otherwise
     """
 
     # Make sure window has all the needed layers.
@@ -588,6 +626,16 @@ def check_window(inputs: dict[str, DataInput], window: Window) -> Window | None:
             )
             return None
 
+    # Optionally skip windows that already have the specified output layer completed.
+    if output_layer_name_skip_inference_if_exists is not None:
+        if window.is_layer_completed(output_layer_name_skip_inference_if_exists):
+            logger.debug(
+                "Skipping window %s since output layer '%s' already exists",
+                window.name,
+                output_layer_name_skip_inference_if_exists,
+            )
+            return None
+
     return window
 
 
@@ -603,6 +651,7 @@ class ModelDataset(torch.utils.data.Dataset):
         workers: int,
         name: str | None = None,
         fix_patch_pick: bool = False,
+        index_mode: IndexMode = IndexMode.OFF,
     ) -> None:
         """Instantiate a new ModelDataset.
 
@@ -612,9 +661,10 @@ class ModelDataset(torch.utils.data.Dataset):
             inputs: data to read from the dataset for training
             task: the task to train on
             workers: number of workers to use for initializing the dataset
-            name: name of the dataset (default: None)
+            name: name of the dataset
             fix_patch_pick: if True, fix the patch 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)
         """
         self.dataset = dataset
         self.split_config = split_config
@@ -635,58 +685,14 @@ class ModelDataset(torch.utils.data.Dataset):
         else:
             self.patch_size = split_config.get_patch_size()
 
-        windows = self._get_initial_windows(split_config, workers)
-
         # If targets are not needed, remove them from the inputs.
         if split_config.get_skip_targets():
             for k in list(self.inputs.keys()):
                 if self.inputs[k].is_target:
                     del self.inputs[k]
 
-        # Eliminate windows that are missing either a requisite input layer, or missing
-        # all target layers.
-        new_windows = []
-        if workers == 0:
-            for window in windows:
-                if check_window(self.inputs, window) is None:
-                    continue
-                new_windows.append(window)
-        else:
-            p = multiprocessing.Pool(workers)
-            outputs = star_imap_unordered(
-                p,
-                check_window,
-                [
-                    dict(
-                        inputs=self.inputs,
-                        window=window,
-                    )
-                    for window in windows
-                ],
-            )
-            for window in tqdm.tqdm(
-                outputs, total=len(windows), desc="Checking available layers in windows"
-            ):
-                if window is None:
-                    continue
-                new_windows.append(window)
-            p.close()
-        windows = new_windows
-
-        # Sort the windows to ensure that the dataset is consistent across GPUs.
-        # Inconsistent ordering can lead to a subset of windows being processed during
-        # "model test" / "model predict" when using multiple GPUs.
-        # We use a hash so that functionality like num_samples limit gets a random
-        # subset of windows (with respect to the hash function choice).
-        windows.sort(
-            key=lambda window: hashlib.sha256(window.name.encode()).hexdigest()
-        )
-
-        # Limit windows to num_samples if requested.
-        if split_config.num_samples:
-            # The windows are sorted by hash of window name so this distribution should
-            # be representative of the population.
-            windows = windows[0 : split_config.num_samples]
+        # Load windows (from index if available, otherwise from dataset)
+        windows = self._load_windows(split_config, workers, index_mode)
 
         # Write dataset_examples to a file so that we can load it lazily in the worker
         # processes. Otherwise it takes a long time to transmit it when spawning each
@@ -755,6 +761,137 @@ class ModelDataset(torch.utils.data.Dataset):
 
         return windows
 
+    def _load_windows(
+        self,
+        split_config: SplitConfig,
+        workers: int,
+        index_mode: IndexMode,
+    ) -> list[Window]:
+        """Load windows, using index if available.
+
+        This method handles:
+        1. Loading from index if index_mode is USE and index exists
+        2. Otherwise, loading from dataset, filtering, sorting, limiting
+        3. Saving to index if index_mode is USE or REFRESH
+
+        Args:
+            split_config: the split configuration.
+            workers: number of worker processes.
+            index_mode: controls caching behavior.
+
+        Returns:
+            list of processed windows ready for training.
+        """
+        # Try to load from index
+        index: DatasetIndex | None = None
+
+        if index_mode != IndexMode.OFF:
+            logger.info(f"Checking index for dataset {self.dataset.path}")
+            index = DatasetIndex(
+                storage=self.dataset.storage,
+                dataset_path=self.dataset.path,
+                groups=split_config.groups,
+                names=split_config.names,
+                tags=split_config.tags,
+                num_samples=split_config.num_samples,
+                skip_targets=split_config.get_skip_targets(),
+                inputs=self.inputs,
+            )
+            refresh = index_mode == IndexMode.REFRESH
+            indexed_windows = index.load_windows(refresh)
+
+            if indexed_windows is not None:
+                logger.info(f"Loaded {len(indexed_windows)} windows from index")
+                return indexed_windows
+
+        # No index available, load and process windows from dataset
+        logger.debug("Loading windows from dataset...")
+        windows = self._get_initial_windows(split_config, workers)
+        windows = self._filter_windows_by_layers(windows, workers)
+        windows = self._sort_and_limit_windows(windows, split_config)
+
+        # Save to index if enabled
+        if index is not None:
+            index.save_windows(windows)
+
+        return windows
+
+    def _filter_windows_by_layers(
+        self, windows: list[Window], workers: int
+    ) -> list[Window]:
+        """Filter windows to only include those with required layers.
+
+        Args:
+            windows: list of windows to filter.
+            workers: number of worker processes for parallel filtering.
+
+        Returns:
+            list of windows that have all required input layers.
+        """
+        output_layer_skip = (
+            self.split_config.get_output_layer_name_skip_inference_if_exists()
+        )
+
+        if workers == 0:
+            return [
+                w
+                for w in windows
+                if check_window(
+                    self.inputs,
+                    w,
+                    output_layer_name_skip_inference_if_exists=output_layer_skip,
+                )
+                is not None
+            ]
+
+        p = multiprocessing.Pool(workers)
+        outputs = star_imap_unordered(
+            p,
+            check_window,
+            [
+                dict(
+                    inputs=self.inputs,
+                    window=window,
+                    output_layer_name_skip_inference_if_exists=output_layer_skip,
+                )
+                for window in windows
+            ],
+        )
+        filtered = []
+        for window in tqdm.tqdm(
+            outputs,
+            total=len(windows),
+            desc="Checking available layers in windows",
+        ):
+            if window is not None:
+                filtered.append(window)
+        p.close()
+        return filtered
+
+    def _sort_and_limit_windows(
+        self, windows: list[Window], split_config: SplitConfig
+    ) -> list[Window]:
+        """Sort windows by hash and apply num_samples limit.
+
+        Sorting ensures consistent ordering across GPUs. Using hash gives a
+        pseudo-random but deterministic order for sampling.
+
+        Args:
+            windows: list of windows to sort and limit.
+            split_config: the split configuration with num_samples.
+
+        Returns:
+            sorted and optionally limited list of windows.
+        """
+        windows.sort(
+            key=lambda window: hashlib.sha256(window.name.encode()).hexdigest()
+        )
+
+        if split_config.num_samples:
+            windows = windows[: split_config.num_samples]
+
+        return windows
+
     def _serialize_item(self, example: Window) -> dict[str, Any]:
         return example.get_metadata()