rslearn.train.dataset

dataset

Default Dataset for rslearn.

IndexMode

Bases: StrEnum

Controls dataset index caching behavior.

Source code in rslearn/train/dataset.py

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."""

OFF class-attribute instance-attribute

OFF = 'off'

No caching - always load windows from dataset.

USE class-attribute instance-attribute

USE = 'use'

Use cached index if available, create if not.

REFRESH class-attribute instance-attribute

REFRESH = 'refresh'

Ignore existing cache and rebuild.

SamplerFactory

Factory to produce a Sampler.

This enables configuring a sampler without needing to pass the dataset.

Source code in rslearn/train/dataset.py

class SamplerFactory:
    """Factory to produce a Sampler.

    This enables configuring a sampler without needing to pass the dataset.
    """

    def get_sampler(self, dataset: "ModelDataset") -> torch.utils.data.Sampler:
        """Create a sampler for the given dataset.

        Args:
            dataset: the dataset

        Returns:
            a sampler
        """
        raise NotImplementedError

get_sampler

get_sampler(dataset: ModelDataset) -> Sampler

Create a sampler for the given dataset.

Parameters:

Name	Type	Description	Default
dataset	ModelDataset	the dataset	required

Returns:

Type	Description
Sampler	a sampler

Source code in rslearn/train/dataset.py

def get_sampler(self, dataset: "ModelDataset") -> torch.utils.data.Sampler:
    """Create a sampler for the given dataset.

    Args:
        dataset: the dataset

    Returns:
        a sampler
    """
    raise NotImplementedError

RandomSamplerFactory

Bases: SamplerFactory

A sampler factory for RandomSampler.

Source code in rslearn/train/dataset.py

class RandomSamplerFactory(SamplerFactory):
    """A sampler factory for RandomSampler."""

    def __init__(
        self, replacement: bool = False, num_samples: int | None = None
    ) -> None:
        """Initialize a RandomSamplerFactory.

        Args:
            replacement: whether to pick with replacement, default false
            num_samples: optional number of dataset samples to limit iteration to,
                otherwise picks random samples equal to the dataset size
        """
        self.replacement = replacement
        self.num_samples = num_samples

    def get_sampler(self, dataset: "ModelDataset") -> torch.utils.data.Sampler:
        """Create a sampler for the given dataset.

        Args:
            dataset: the dataset

        Returns:
            a RandomSampler
        """
        return torch.utils.data.RandomSampler(
            dataset, replacement=self.replacement, num_samples=self.num_samples
        )

get_sampler

get_sampler(dataset: ModelDataset) -> Sampler

Create a sampler for the given dataset.

Parameters:

Name	Type	Description	Default
dataset	ModelDataset	the dataset	required

Returns:

Type	Description
Sampler	a RandomSampler

Source code in rslearn/train/dataset.py

def get_sampler(self, dataset: "ModelDataset") -> torch.utils.data.Sampler:
    """Create a sampler for the given dataset.

    Args:
        dataset: the dataset

    Returns:
        a RandomSampler
    """
    return torch.utils.data.RandomSampler(
        dataset, replacement=self.replacement, num_samples=self.num_samples
    )

WeightedRandomSamplerFactory

Bases: SamplerFactory

A sampler factory for WeightedRandomSampler.

Source code in rslearn/train/dataset.py

class WeightedRandomSamplerFactory(SamplerFactory):
    """A sampler factory for WeightedRandomSampler."""

    def __init__(
        self, option_key: str, num_samples: int, replacement: bool = True
    ) -> None:
        """Initialize a WeightedRandomSamplerFactory.

        Args:
            option_key: the key in the window option dict containing the weights
            num_samples: number of examples to sample per epoch
            replacement: whether to pick with replacement, default true
        """
        self.option_key = option_key
        self.num_samples = num_samples
        self.replacement = replacement

    def get_sampler(self, dataset: "ModelDataset") -> torch.utils.data.Sampler:
        """Create a sampler for the given dataset.

        Args:
            dataset: the dataset

        Returns:
            a RandomSampler
        """
        weights = []
        for window in dataset.get_dataset_examples():
            weights.append(window.options[self.option_key])
        return torch.utils.data.WeightedRandomSampler(
            weights, self.num_samples, replacement=self.replacement
        )

get_sampler

get_sampler(dataset: ModelDataset) -> Sampler

Create a sampler for the given dataset.

Parameters:

Name	Type	Description	Default
dataset	ModelDataset	the dataset	required

Returns:

Type	Description
Sampler	a RandomSampler

Source code in rslearn/train/dataset.py

def get_sampler(self, dataset: "ModelDataset") -> torch.utils.data.Sampler:
    """Create a sampler for the given dataset.

    Args:
        dataset: the dataset

    Returns:
        a RandomSampler
    """
    weights = []
    for window in dataset.get_dataset_examples():
        weights.append(window.options[self.option_key])
    return torch.utils.data.WeightedRandomSampler(
        weights, self.num_samples, replacement=self.replacement
    )

DataInput

Specification of a piece of data from a window that is needed for training.

The DataInput includes which layer(s) the data can be obtained from for each window.

Note that this class is not a dataclass because jsonargparse does not play well with dataclasses without enabling specialized options which we have not validated will work with the rest of our code.

Source code in rslearn/train/dataset.py

class DataInput:
    """Specification of a piece of data from a window that is needed for training.

    The DataInput includes which layer(s) the data can be obtained from for each window.

    Note that this class is not a dataclass because jsonargparse does not play well
    with dataclasses without enabling specialized options which we have not validated
    will work with the rest of our code.
    """

    def __init__(
        self,
        data_type: str,
        layers: list[str],
        bands: list[str] | None = None,
        use_all_bands_in_order_of_band_set_idx: int | None = None,
        required: bool = True,
        passthrough: bool = False,
        is_target: bool = False,
        dtype: DType = DType.FLOAT32,
        load_all_layers: bool = False,
        load_all_item_groups: bool = False,
        resolution_factor: ResolutionFactor = ResolutionFactor(),
        resampling: Resampling = Resampling.nearest,
    ):
        """Initialize a new DataInput.

        Args:
            data_type: either "raster" or "vector"
            layers: list of layer names or item group specifiers that this input can be
                read from. If load_all_item_groups=False, each entry should be an item
                group specifier (e.g. "sentinel2" for layer_name=sentinel2, group_idx=0
                or "sentinel2.1" for layer_name=sentinel2, group_idx=1). Otherwise,
                each entry should be a layer name. For example, if you have a layer
                "sentinel2" with three item groups: with load_all_item_groups=False,
                set layers=["sentinel2", "sentinel2.1", "sentinel2.2"]; with
                load_all_item_groups=True, set layers=["sentinel2"].
            bands: the bands to read, if this is a raster.
            use_all_bands_in_order_of_band_set_idx: if set, read all bands from the
                specified layer_config band_set index (ordered as listed in that
                band_set). This is useful for large embedding layers where listing all
                band names in model config is cumbersome.
            required: whether examples lacking one of these layers should be skipped
            passthrough: whether to expose this to the model even if it isn't returned
                by any task
            is_target: whether this DataInput represents a target for the task. Targets
                are not read during prediction phase.
            dtype: data type to load the raster as
            load_all_layers: whether to load all of the entries specified in the layers
                list. By default, we randomly pick one entry to read. When reading
                multiple entries, raster images are stacked on the time dimension. This
                option will also cause the dataset to only include windows where all of
                the entries are materialized (by default, only windows with none of the
                entries materialized would be excluded).
            load_all_item_groups: whether to load all item groups in the layer(s) we
                are reading from. By default, we treat layers as a list of item group
                specifiers, and either pick a random item group to read (if
                load_all_layers=False) or stack all of them (if load_all_layers=True). If
                load_all_item_groups=True, we treat layers as a list of layer names,
                and include all item groups within each layer as candidates for
                reading; whether we pick a random item group or stack them is still
                controlled by load_all_layers. Note that, when load_all_layers=True and
                load_all_item_groups=True, we will only exclude windows from training
                that have zero item groups in one of the configured layers; additionally,
                if windows have different numbers of item groups, then we will read
                RasterImages with different numbers of timesteps.
            resolution_factor: controls the resolution at which raster data is loaded for training.
                By default (factor=1), data is loaded at the window resolution.
                E.g. for a 64x64 window at 10 m/pixel with resolution_factor=1/2,
                the resulting tensor is 32x32 (covering the same geographic area at 20 m/pixel).
            resampling: resampling method (default nearest neighbor).
        """
        self.data_type = data_type
        self.layers = layers
        self.required = required
        self.passthrough = passthrough
        self.is_target = is_target
        self.dtype = dtype
        self.load_all_layers = load_all_layers
        self.load_all_item_groups = load_all_item_groups
        self.resolution_factor = resolution_factor
        self.resampling = resampling

        if bands is not None and use_all_bands_in_order_of_band_set_idx is not None:
            raise ValueError(
                "only one of bands and use_all_bands_in_order_of_band_set_idx should be set"
            )
        if (
            self.data_type == "raster"
            and bands is None
            and use_all_bands_in_order_of_band_set_idx is None
        ):
            raise ValueError(
                "for raster DataInputs, one of bands and use_all_bands_in_order_of_band_set_idx must be set"
            )

        self.bands = bands
        self.use_all_bands_in_order_of_band_set_idx = (
            use_all_bands_in_order_of_band_set_idx
        )

SplitConfig

Configuration that can be specified separately for train, val, and test.

Source code in rslearn/train/dataset.py

class SplitConfig:
    """Configuration that can be specified separately for train, val, and test."""

    def __init__(
        self,
        groups: list[str] | None = None,
        names: list[str] | None = None,
        tags: dict[str, Any] | None = None,
        num_samples: int | None = None,
        num_patches: int | None = None,
        transforms: list[torch.nn.Module] | None = None,
        sampler: SamplerFactory | None = None,
        crop_size: int | tuple[int, int] | None = None,
        overlap_pixels: int | None = None,
        load_all_crops: bool | None = None,
        skip_targets: bool | None = None,
        output_layer_name_skip_inference_if_exists: str | None = None,
        # Deprecated parameters (for backwards compatibility)
        patch_size: int | tuple[int, int] | None = None,
        overlap_ratio: float | None = None,
        load_all_patches: bool | None = None,
    ) -> None:
        """Initialize a new SplitConfig.

        Args:
            groups: for this split, only read windows in one of these groups
            names: for this split, read windows with these specific names
            tags: only select windows that have options matching these tags. If key and
                value are set, then window must have an option with the same key and
                value. If value is empty, then only the existince of the key in the
                window options is checked.
            num_samples: limit this split to this many examples
            num_patches: limit this split to this many patches
            transforms: transforms to apply
            sampler: SamplerFactory for this split
            crop_size: an optional square size or (width, height) tuple. If set, read
                crops of this size rather than entire windows.
            overlap_pixels: the number of pixels shared between adjacent crops during
                sliding window inference.
            load_all_crops: with crop_size set, rather than sampling a random crop
                for each window, read all crops 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).
            patch_size: deprecated, use crop_size instead
            overlap_ratio: deprecated, use overlap_pixels instead
            load_all_patches: deprecated, use load_all_crops instead
        """
        self.groups = groups
        self.names = names
        self.tags = tags
        self.num_samples = num_samples
        self.num_patches = num_patches
        self.transforms = transforms
        self.sampler = sampler
        self.skip_targets = skip_targets
        self.output_layer_name_skip_inference_if_exists = (
            output_layer_name_skip_inference_if_exists
        )

        # These have deprecated equivalents -- we store both raw values since we don't
        # have a complete picture until the final merged SplitConfig is computed. We
        # raise deprecation warnings in merge_and_validate and we disambiguate them in
        # get_ functions (so the variables should never be accessed directly).
        self._crop_size = crop_size
        self._patch_size = patch_size
        self._overlap_pixels = overlap_pixels
        self._overlap_ratio = overlap_ratio
        self._load_all_crops = load_all_crops
        self._load_all_patches = load_all_patches

    def _merge(self, other: "SplitConfig") -> "SplitConfig":
        """Merge settings from another SplitConfig into this one.

        Args:
            other: the config to merge in (its non-None values override self's)

        Returns:
            the resulting SplitConfig combining the settings.
        """
        result = SplitConfig(
            groups=self.groups,
            names=self.names,
            tags=self.tags,
            num_samples=self.num_samples,
            num_patches=self.num_patches,
            transforms=self.transforms,
            sampler=self.sampler,
            crop_size=self._crop_size,
            patch_size=self._patch_size,
            overlap_pixels=self._overlap_pixels,
            overlap_ratio=self._overlap_ratio,
            load_all_crops=self._load_all_crops,
            load_all_patches=self._load_all_patches,
            skip_targets=self.skip_targets,
            output_layer_name_skip_inference_if_exists=self.output_layer_name_skip_inference_if_exists,
        )
        if other.groups:
            result.groups = other.groups
        if other.names:
            result.names = other.names
        if other.tags:
            result.tags = other.tags
        if other.num_samples:
            result.num_samples = other.num_samples
        if other.num_patches:
            result.num_patches = other.num_patches
        if other.transforms:
            result.transforms = other.transforms
        if other.sampler:
            result.sampler = other.sampler
        if other._crop_size is not None:
            result._crop_size = other._crop_size
        if other._patch_size is not None:
            result._patch_size = other._patch_size
        if other._overlap_pixels is not None:
            result._overlap_pixels = other._overlap_pixels
        if other._overlap_ratio is not None:
            result._overlap_ratio = other._overlap_ratio
        if other._load_all_crops is not None:
            result._load_all_crops = other._load_all_crops
        if other._load_all_patches is not None:
            result._load_all_patches = other._load_all_patches
        if other.skip_targets is not None:
            result.skip_targets = other.skip_targets
        if other.output_layer_name_skip_inference_if_exists is not None:
            result.output_layer_name_skip_inference_if_exists = (
                other.output_layer_name_skip_inference_if_exists
            )
        return result

    @staticmethod
    def merge_and_validate(configs: list["SplitConfig"]) -> "SplitConfig":
        """Merge a list of SplitConfigs and validate the result.

        Args:
            configs: list of SplitConfig to merge. Later configs override earlier ones.

        Returns:
            the merged and validated SplitConfig.
        """
        if not configs:
            return SplitConfig()

        result = configs[0]
        for config in configs[1:]:
            result = result._merge(config)

        # Emit deprecation warnings
        if result._patch_size is not None:
            warnings.warn(
                "patch_size is deprecated, use crop_size instead",
                FutureWarning,
                stacklevel=2,
            )
        if result._overlap_ratio is not None:
            warnings.warn(
                "overlap_ratio is deprecated, use overlap_pixels instead",
                FutureWarning,
                stacklevel=2,
            )
        if result._load_all_patches is not None:
            warnings.warn(
                "load_all_patches is deprecated, use load_all_crops instead",
                FutureWarning,
                stacklevel=2,
            )

        # Check for conflicting parameters
        if result._crop_size is not None and result._patch_size is not None:
            raise ValueError("Cannot specify both crop_size and patch_size")
        if result._overlap_pixels is not None and result._overlap_ratio is not None:
            raise ValueError("Cannot specify both overlap_pixels and overlap_ratio")
        if result._load_all_crops is not None and result._load_all_patches is not None:
            raise ValueError("Cannot specify both load_all_crops and load_all_patches")

        # Validate overlap_pixels is non-negative
        if result._overlap_pixels is not None and result._overlap_pixels < 0:
            raise ValueError("overlap_pixels must be non-negative")

        # overlap_pixels requires load_all_crops.
        if result.get_overlap_pixels() > 0 and not result.get_load_all_crops():
            raise ValueError(
                "overlap_pixels requires load_all_crops to be True since (overlap is only used during sliding window inference"
            )

        return result

    def get_crop_size(self) -> tuple[int, int] | None:
        """Get crop size as tuple, handling deprecated patch_size."""
        size = self._crop_size if self._crop_size is not None else self._patch_size
        if size is None:
            return None
        if isinstance(size, int):
            return (size, size)
        return size

    def get_overlap_pixels(self) -> int:
        """Get the overlap pixels (default 0), handling deprecated overlap_ratio."""
        if self._overlap_pixels is not None:
            return self._overlap_pixels
        if self._overlap_ratio is not None:
            crop_size = self.get_crop_size()
            if crop_size is None:
                raise ValueError("overlap_ratio requires crop_size to be set")
            return round(crop_size[0] * self._overlap_ratio)
        return 0

    def get_load_all_crops(self) -> bool:
        """Returns whether loading all crops is enabled (default False)."""
        if self._load_all_crops is not None:
            return self._load_all_crops
        if self._load_all_patches is not None:
            return self._load_all_patches
        return False

    def get_skip_targets(self) -> bool:
        """Returns whether skip_targets is enabled (default False)."""
        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

merge_and_validate staticmethod

merge_and_validate(configs: list[SplitConfig]) -> SplitConfig

Merge a list of SplitConfigs and validate the result.

Parameters:

Name	Type	Description	Default
configs	list[SplitConfig]	list of SplitConfig to merge. Later configs override earlier ones.	required

Returns:

Type	Description
SplitConfig	the merged and validated SplitConfig.

Source code in rslearn/train/dataset.py

@staticmethod
def merge_and_validate(configs: list["SplitConfig"]) -> "SplitConfig":
    """Merge a list of SplitConfigs and validate the result.

    Args:
        configs: list of SplitConfig to merge. Later configs override earlier ones.

    Returns:
        the merged and validated SplitConfig.
    """
    if not configs:
        return SplitConfig()

    result = configs[0]
    for config in configs[1:]:
        result = result._merge(config)

    # Emit deprecation warnings
    if result._patch_size is not None:
        warnings.warn(
            "patch_size is deprecated, use crop_size instead",
            FutureWarning,
            stacklevel=2,
        )
    if result._overlap_ratio is not None:
        warnings.warn(
            "overlap_ratio is deprecated, use overlap_pixels instead",
            FutureWarning,
            stacklevel=2,
        )
    if result._load_all_patches is not None:
        warnings.warn(
            "load_all_patches is deprecated, use load_all_crops instead",
            FutureWarning,
            stacklevel=2,
        )

    # Check for conflicting parameters
    if result._crop_size is not None and result._patch_size is not None:
        raise ValueError("Cannot specify both crop_size and patch_size")
    if result._overlap_pixels is not None and result._overlap_ratio is not None:
        raise ValueError("Cannot specify both overlap_pixels and overlap_ratio")
    if result._load_all_crops is not None and result._load_all_patches is not None:
        raise ValueError("Cannot specify both load_all_crops and load_all_patches")

    # Validate overlap_pixels is non-negative
    if result._overlap_pixels is not None and result._overlap_pixels < 0:
        raise ValueError("overlap_pixels must be non-negative")

    # overlap_pixels requires load_all_crops.
    if result.get_overlap_pixels() > 0 and not result.get_load_all_crops():
        raise ValueError(
            "overlap_pixels requires load_all_crops to be True since (overlap is only used during sliding window inference"
        )

    return result

get_crop_size

get_crop_size() -> tuple[int, int] | None

Get crop size as tuple, handling deprecated patch_size.

Source code in rslearn/train/dataset.py

def get_crop_size(self) -> tuple[int, int] | None:
    """Get crop size as tuple, handling deprecated patch_size."""
    size = self._crop_size if self._crop_size is not None else self._patch_size
    if size is None:
        return None
    if isinstance(size, int):
        return (size, size)
    return size

get_overlap_pixels

get_overlap_pixels() -> int

Get the overlap pixels (default 0), handling deprecated overlap_ratio.

Source code in rslearn/train/dataset.py

def get_overlap_pixels(self) -> int:
    """Get the overlap pixels (default 0), handling deprecated overlap_ratio."""
    if self._overlap_pixels is not None:
        return self._overlap_pixels
    if self._overlap_ratio is not None:
        crop_size = self.get_crop_size()
        if crop_size is None:
            raise ValueError("overlap_ratio requires crop_size to be set")
        return round(crop_size[0] * self._overlap_ratio)
    return 0

get_load_all_crops

get_load_all_crops() -> bool

Returns whether loading all crops is enabled (default False).

Source code in rslearn/train/dataset.py

def get_load_all_crops(self) -> bool:
    """Returns whether loading all crops is enabled (default False)."""
    if self._load_all_crops is not None:
        return self._load_all_crops
    if self._load_all_patches is not None:
        return self._load_all_patches
    return False

get_skip_targets

get_skip_targets() -> bool

Returns whether skip_targets is enabled (default False).

Source code in rslearn/train/dataset.py

def get_skip_targets(self) -> bool:
    """Returns whether skip_targets is enabled (default False)."""
    return True if self.skip_targets is True else False

get_output_layer_name_skip_inference_if_exists

get_output_layer_name_skip_inference_if_exists() -> str | None

Returns output layer to use for resume checks (default None).

Source code in rslearn/train/dataset.py

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

CheckWindowResult dataclass

Aggregated counts of why windows were skipped during check_window.

Source code in rslearn/train/dataset.py

@dataclasses.dataclass
class CheckWindowResult:
    """Aggregated counts of why windows were skipped during check_window."""

    missing_data_input_counts: dict[str, int] = dataclasses.field(default_factory=dict)
    has_output_layer_count: int = 0

    def add(self, other: "CheckWindowResult") -> None:
        """Merge another result into this one."""
        for key, count in other.missing_data_input_counts.items():
            self.missing_data_input_counts[key] = (
                self.missing_data_input_counts.get(key, 0) + count
            )
        self.has_output_layer_count += other.has_output_layer_count

add

add(other: CheckWindowResult) -> None

Merge another result into this one.

Source code in rslearn/train/dataset.py

def add(self, other: "CheckWindowResult") -> None:
    """Merge another result into this one."""
    for key, count in other.missing_data_input_counts.items():
        self.missing_data_input_counts[key] = (
            self.missing_data_input_counts.get(key, 0) + count
        )
    self.has_output_layer_count += other.has_output_layer_count

ModelDataset

Bases: Dataset

The default pytorch dataset implementation for rslearn.

Source code in rslearn/train/dataset.py

class ModelDataset(torch.utils.data.Dataset):
    """The default pytorch dataset implementation for rslearn."""

    def __init__(
        self,
        dataset: Dataset,
        split_config: SplitConfig,
        inputs: dict[str, DataInput],
        task: Task,
        workers: int,
        name: str | None = None,
        fix_crop_pick: bool = False,
        index_mode: IndexMode = IndexMode.OFF,
        load_full_windows: bool = False,
    ) -> None:
        """Instantiate a new ModelDataset.

        Args:
            dataset: underlying rslearn dataset to read data from
            split_config: configuration specific to this split
            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
            fix_crop_pick: if True, fix the crop pick to be the same every time
                for a given window. Useful for testing (default: False)
            index_mode: controls dataset index caching behavior (default: IndexMode.OFF)
            load_full_windows: if True, always load entire windows regardless
                of crop_size. Used by in-memory dataset wrappers that handle
                cropping themselves.
        """
        self.dataset = dataset
        self.split_config = split_config
        self.inputs = inputs
        self.task = task
        self.name = name
        self.fix_crop_pick = fix_crop_pick
        if split_config.transforms:
            self.transforms = Sequential(*split_config.transforms)
        else:
            self.transforms = rslearn.train.transforms.transform.Identity()

        # Skip cropping if the caller will handle it (in-memory wrappers,
        # AllCropsDataset, etc.).
        if load_full_windows or split_config.get_load_all_crops():
            self.crop_size = None
        else:
            self.crop_size = split_config.get_crop_size()

        # 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]

        # 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
        # process.
        self.dataset_examples_fname = os.path.join(
            tempfile.gettempdir(),
            "rslearn_dataset_examples",
            f"{os.getpid()}_{uuid.uuid4()}.json",
        )
        self.num_dataset_examples = len(windows)
        self.dataset_examples: list[Window] | None = None
        logger.info(
            f"Writing {len(windows)} dataset examples to {self.dataset_examples_fname}"
        )
        os.makedirs(os.path.dirname(self.dataset_examples_fname), exist_ok=True)
        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 _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,
                data_storage_factory=self.dataset.window_data_storage_factory,
                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()
        )

        result = CheckWindowResult()
        filtered = []

        with make_pool_and_star_imap_unordered(
            workers,
            check_window,
            [
                dict(
                    inputs=self.inputs,
                    window=window,
                    output_layer_name_skip_inference_if_exists=output_layer_skip,
                )
                for window in windows
            ],
        ) as outputs:
            for window, check_result in tqdm.tqdm(
                outputs,
                total=len(windows),
                desc="Checking available layers in windows",
            ):
                result.add(check_result)
                if window is not None:
                    filtered.append(window)

        if result.missing_data_input_counts:
            for key, count in sorted(result.missing_data_input_counts.items()):
                logger.info("Skipped %d windows due to missing input '%s'", count, key)
        if result.has_output_layer_count > 0:
            logger.info(
                "Skipped %d windows due to existing output layer",
                result.has_output_layer_count,
            )

        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()

    def _deserialize_item(self, d: dict[str, Any]) -> Window:
        window = Window.from_metadata(self.dataset.storage, d)
        window._data = self.dataset.window_data_storage_factory.create(window)
        return window

    def get_dataset_examples(self) -> list[Window]:
        """Get a list of examples in the dataset.

        If load_all_crops is False, this is a list of Windows. Otherwise, this is a
        list of (window, crop_bounds, (crop_idx, # crops)) tuples.
        """
        if self.dataset_examples is None:
            logger.debug(
                f"Loading dataset examples from {self.dataset_examples_fname} in process {os.getpid()}"
            )
            with open(self.dataset_examples_fname) as f:
                self.dataset_examples = [
                    self._deserialize_item(d) for d in json.load(f)
                ]
            logger.debug(f"Finished loading dataset examples in process {os.getpid()}")
        return self.dataset_examples

    def __len__(self) -> int:
        """Returns the dataset length."""
        return self.num_dataset_examples

    def get_raw_inputs(
        self, idx: int
    ) -> 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
        should be a RasterImage for raster and list[Feature] for vector.

        Args:
            idx: the index in the dataset.

        Returns:
            a tuple (raw_inputs, passthrough_inputs, metadata).
        """
        dataset_examples = self.get_dataset_examples()
        example = dataset_examples[idx]
        rng = random.Random(idx if self.fix_crop_pick else None)

        # Select bounds to read.
        if self.crop_size:
            window = example

            def get_crop_range(n_crop: int, n_window: int) -> list[int]:
                if n_crop > n_window:
                    # Select arbitrary range containing the entire window.
                    # Basically arbitrarily padding the window to get to crop size.
                    start = rng.randint(n_window - n_crop, 0)
                    return [start, start + n_crop]

                else:
                    # Select arbitrary crop within the window.
                    start = rng.randint(0, n_window - n_crop)
                    return [start, start + n_crop]

            window_size = (
                window.bounds[2] - window.bounds[0],
                window.bounds[3] - window.bounds[1],
            )
            crop_ranges = [
                get_crop_range(self.crop_size[0], window_size[0]),
                get_crop_range(self.crop_size[1], window_size[1]),
            ]
            bounds = (
                window.bounds[0] + crop_ranges[0][0],
                window.bounds[1] + crop_ranges[1][0],
                window.bounds[0] + crop_ranges[0][1],
                window.bounds[1] + crop_ranges[1][1],
            )

        else:
            window = example
            bounds = window.bounds

        assert isinstance(window, Window)

        raw_inputs = {}
        passthrough_inputs = {}
        for name, data_input in self.inputs.items():
            # Skip non-required inputs if their layers are not available
            if not data_input.required and not is_data_input_available(
                data_input, window
            ):
                logger.debug(
                    "Skipping non-required input '%s' for window %s (layers not available)",
                    name,
                    window.name,
                )
                continue

            raw_inputs[name] = read_data_input(
                self.dataset, window, bounds, data_input, rng
            )
            if data_input.passthrough:
                passthrough_inputs[name] = raw_inputs[name]

        metadata = SampleMetadata(
            window_group=window.group,
            window_name=window.name,
            window_bounds=window.bounds,
            crop_bounds=bounds,
            crop_idx=0,
            num_crops_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], SampleMetadata]:
        """Read one training example.

        Args:
            idx: the index in the dataset.

        Returns:
            a tuple (input_dict, target_dict, metadata)
        """
        logger.debug("__getitem__ start pid=%d item_idx=%d", os.getpid(), idx)

        raw_inputs, passthrough_inputs, metadata = self.get_raw_inputs(idx)

        input_dict, target_dict = self.task.process_inputs(
            raw_inputs,
            metadata=metadata,
            load_targets=not self.split_config.get_skip_targets(),
        )
        input_dict.update(passthrough_inputs)
        input_dict, target_dict = self.transforms(input_dict, target_dict)

        logger.debug("__getitem__ finish pid=%d item_idx=%d", os.getpid(), idx)

        return input_dict, target_dict, metadata

    def set_name(self, name: str) -> None:
        """Set the name of the dataset.

        Args:
            name: the name to set.
        """
        self.name = name

get_dataset_examples

get_dataset_examples() -> list[Window]

Get a list of examples in the dataset.

If load_all_crops is False, this is a list of Windows. Otherwise, this is a list of (window, crop_bounds, (crop_idx, # crops)) tuples.

Source code in rslearn/train/dataset.py

def get_dataset_examples(self) -> list[Window]:
    """Get a list of examples in the dataset.

    If load_all_crops is False, this is a list of Windows. Otherwise, this is a
    list of (window, crop_bounds, (crop_idx, # crops)) tuples.
    """
    if self.dataset_examples is None:
        logger.debug(
            f"Loading dataset examples from {self.dataset_examples_fname} in process {os.getpid()}"
        )
        with open(self.dataset_examples_fname) as f:
            self.dataset_examples = [
                self._deserialize_item(d) for d in json.load(f)
            ]
        logger.debug(f"Finished loading dataset examples in process {os.getpid()}")
    return self.dataset_examples

get_raw_inputs

get_raw_inputs(idx: int) -> 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 should be a RasterImage for raster and list[Feature] for vector.

Parameters:

Name	Type	Description	Default
idx	int	the index in the dataset.	required

Returns:

Type	Description
tuple[dict[str, Any], dict[str, Any], SampleMetadata]	a tuple (raw_inputs, passthrough_inputs, metadata).

Source code in rslearn/train/dataset.py

def get_raw_inputs(
    self, idx: int
) -> 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
    should be a RasterImage for raster and list[Feature] for vector.

    Args:
        idx: the index in the dataset.

    Returns:
        a tuple (raw_inputs, passthrough_inputs, metadata).
    """
    dataset_examples = self.get_dataset_examples()
    example = dataset_examples[idx]
    rng = random.Random(idx if self.fix_crop_pick else None)

    # Select bounds to read.
    if self.crop_size:
        window = example

        def get_crop_range(n_crop: int, n_window: int) -> list[int]:
            if n_crop > n_window:
                # Select arbitrary range containing the entire window.
                # Basically arbitrarily padding the window to get to crop size.
                start = rng.randint(n_window - n_crop, 0)
                return [start, start + n_crop]

            else:
                # Select arbitrary crop within the window.
                start = rng.randint(0, n_window - n_crop)
                return [start, start + n_crop]

        window_size = (
            window.bounds[2] - window.bounds[0],
            window.bounds[3] - window.bounds[1],
        )
        crop_ranges = [
            get_crop_range(self.crop_size[0], window_size[0]),
            get_crop_range(self.crop_size[1], window_size[1]),
        ]
        bounds = (
            window.bounds[0] + crop_ranges[0][0],
            window.bounds[1] + crop_ranges[1][0],
            window.bounds[0] + crop_ranges[0][1],
            window.bounds[1] + crop_ranges[1][1],
        )

    else:
        window = example
        bounds = window.bounds

    assert isinstance(window, Window)

    raw_inputs = {}
    passthrough_inputs = {}
    for name, data_input in self.inputs.items():
        # Skip non-required inputs if their layers are not available
        if not data_input.required and not is_data_input_available(
            data_input, window
        ):
            logger.debug(
                "Skipping non-required input '%s' for window %s (layers not available)",
                name,
                window.name,
            )
            continue

        raw_inputs[name] = read_data_input(
            self.dataset, window, bounds, data_input, rng
        )
        if data_input.passthrough:
            passthrough_inputs[name] = raw_inputs[name]

    metadata = SampleMetadata(
        window_group=window.group,
        window_name=window.name,
        window_bounds=window.bounds,
        crop_bounds=bounds,
        crop_idx=0,
        num_crops_in_window=1,
        time_range=window.time_range,
        projection=window.projection,
        dataset_source=self.name,
    )

    return raw_inputs, passthrough_inputs, metadata

set_name

set_name(name: str) -> None

Set the name of the dataset.

Parameters:

Name	Type	Description	Default
name	str	the name to set.	required

Source code in rslearn/train/dataset.py

def set_name(self, name: str) -> None:
    """Set the name of the dataset.

    Args:
        name: the name to set.
    """
    self.name = name

RetryDataset

Bases: Dataset

A dataset wrapper that retries getitem upon encountering error.

Source code in rslearn/train/dataset.py

class RetryDataset(torch.utils.data.Dataset):
    """A dataset wrapper that retries getitem upon encountering error."""

    def __init__(
        self, dataset: ModelDataset, retries: int = 3, delay: float = 5
    ) -> None:
        """Create a new RetryDataset.

        Args:
            dataset: the dataset to wrap.
            retries: the maximum number of tries before raising error.
            delay: how many seconds to sleep before retrying
        """
        self.dataset = dataset
        self.retries = retries
        self.delay = delay

    def set_name(self, name: str) -> None:
        """Set the name of the dataset.

        Args:
            name: the name to set.
        """
        self.dataset.set_name(name)

    def __len__(self) -> int:
        """Return length of the dataset."""
        return len(self.dataset)

    def __getitem__(self, idx: int) -> Any:
        """Get item from the dataset.

        The get operation is performed on the underlying dataset multiple times up to
        the configured maximum number of retries.

        Args:
            idx: the item index.

        Returns:
            the item data.
        """
        for _ in range(self.retries):
            try:
                return self.dataset[idx]
            except Exception as e:
                logger.warning("warning: caught exception loading item %d: %s", idx, e)
            time.sleep(self.delay)

        # One last try -- but don't catch any more errors.
        return self.dataset[idx]

    def get_dataset_examples(self) -> list[Window]:
        """Returns a list of windows in this dataset."""
        return self.dataset.get_dataset_examples()

set_name

set_name(name: str) -> None

Set the name of the dataset.

Parameters:

Name	Type	Description	Default
name	str	the name to set.	required

Source code in rslearn/train/dataset.py

def set_name(self, name: str) -> None:
    """Set the name of the dataset.

    Args:
        name: the name to set.
    """
    self.dataset.set_name(name)

get_dataset_examples

get_dataset_examples() -> list[Window]

Returns a list of windows in this dataset.

Source code in rslearn/train/dataset.py

def get_dataset_examples(self) -> list[Window]:
    """Returns a list of windows in this dataset."""
    return self.dataset.get_dataset_examples()

MultiDataset

Bases: Dataset

A dataset that combines multiple datasets.

Source code in rslearn/train/dataset.py

class MultiDataset(torch.utils.data.Dataset):
    """A dataset that combines multiple datasets."""

    def __init__(self, datasets: dict[str, RetryDataset]) -> None:
        """Create a new MultiDataset.

        Args:
            datasets: map of dataset name to dataset.
        """
        self.datasets = datasets
        self.buckets = {}
        curr_offset = 0
        for name, ds in datasets.items():
            self.buckets[name] = range(curr_offset, curr_offset + len(ds))
            curr_offset += len(ds)

    def __len__(self) -> int:
        """Return length of the dataset."""
        return sum(len(ds) for ds in self.datasets.values())

    def __getitem__(self, idx: int) -> Any:
        """Get item from the dataset.

        Args:
            idx: the item index.

        Returns:
            the item data.
        """
        for name, bucket in self.buckets.items():
            if idx in bucket:
                return self.datasets[name][idx - bucket.start]
        raise IndexError(f"Index {idx} out of range (len={len(self)})")

get_torch_dtype

get_torch_dtype(dtype: DType) -> dtype

Convert rslearn DType to torch dtype.

Source code in rslearn/train/dataset.py

def get_torch_dtype(dtype: DType) -> torch.dtype:
    """Convert rslearn DType to torch dtype."""
    if dtype == DType.INT32:
        return torch.int32
    elif dtype == DType.FLOAT32:
        return torch.float32
    else:
        raise ValueError(f"unable to handle {dtype} as a torch dtype")

resolve_raster_data_input_bands

resolve_raster_data_input_bands(data_input: DataInput, layer_name: str, layer_config: LayerConfig) -> list[str]

Resolve the band list for a raster DataInput.

If data_input.bands is explicitly provided, it is returned as-is. Otherwise, if use_all_bands_in_order_of_band_set_idx is set, all bands are taken from that specific band set in the dataset's LayerConfig.

Source code in rslearn/train/dataset.py

def resolve_raster_data_input_bands(
    data_input: DataInput,
    layer_name: str,
    layer_config: LayerConfig,
) -> list[str]:
    """Resolve the band list for a raster DataInput.

    If data_input.bands is explicitly provided, it is returned as-is. Otherwise, if
    use_all_bands_in_order_of_band_set_idx is set, all bands are taken from that
    specific band set in the dataset's LayerConfig.
    """
    if data_input.bands is not None:
        return data_input.bands

    band_set_index = data_input.use_all_bands_in_order_of_band_set_idx
    if band_set_index is None:
        raise ValueError(
            f"No bands specified for raster input when reading layer '{layer_name}'. "
            "Set `bands: [...]`, or set `use_all_bands_in_order_of_band_set_idx` "
            "to a band set index to use all band names from the dataset layer config."
        )

    if band_set_index < 0 or band_set_index >= len(layer_config.band_sets):
        raise ValueError(
            "Invalid "
            f"use_all_bands_in_order_of_band_set_idx={band_set_index} "
            f"for layer '{layer_name}'. "
            f"Expected a value in [0, {len(layer_config.band_sets) - 1}]."
        )

    return layer_config.band_sets[band_set_index].bands

read_raster_layer_groups_for_data_input

read_raster_layer_groups_for_data_input(window: Window, bounds: PixelBounds, layer_name: str, group_idxs: list[int], layer_config: LayerConfig, data_input: DataInput) -> tuple[list[Tensor], list[list[tuple[datetime, datetime]] | None]]

Read the specified item groups of a raster layer for a DataInput.

Parameters:

Name	Type	Description	Default
window	Window	the window to read from.	required
bounds	PixelBounds	the bounds to read.	required
layer_name	str	the layer name.	required
group_idxs	list[int]	ordered list of item group indices to read.	required
layer_config	LayerConfig	the layer configuration.	required
data_input	DataInput	the DataInput that specifies the bands and dtype.	required

Returns:

Type	Description
list[Tensor]	Tuple of (list of CTHW tensors, list of timestamps), each with length
list[list[tuple[datetime, datetime]] | None]	len(group_idxs).

Source code in rslearn/train/dataset.py

def read_raster_layer_groups_for_data_input(
    window: Window,
    bounds: PixelBounds,
    layer_name: str,
    group_idxs: list[int],
    layer_config: LayerConfig,
    data_input: DataInput,
) -> tuple[
    list[torch.Tensor],
    list[list[tuple[datetime, datetime]] | None],
]:
    """Read the specified item groups of a raster layer for a DataInput.

    Args:
        window: the window to read from.
        bounds: the bounds to read.
        layer_name: the layer name.
        group_idxs: ordered list of item group indices to read.
        layer_config: the layer configuration.
        data_input: the DataInput that specifies the bands and dtype.

    Returns:
        Tuple of (list of CTHW tensors, list of timestamps), each with length
        ``len(group_idxs)``.
    """
    if not group_idxs:
        return [], []

    # Determine which band sets we need to read to cover the bands in the DataInput.
    needed_bands = resolve_raster_data_input_bands(
        data_input=data_input,
        layer_name=layer_name,
        layer_config=layer_config,
    )
    needed_sets_and_indexes = _resolve_needed_band_sets(
        needed_bands=needed_bands,
        layer_config=layer_config,
        layer_name=layer_name,
        window_name=window.name,
    )

    # Get the projection and bounds to read under (multiply window resolution # by
    # the specified resolution factor).
    final_projection = data_input.resolution_factor.multiply_projection(
        window.projection
    )
    final_bounds = data_input.resolution_factor.multiply_bounds(bounds)

    # We don't know T upfront (it depends on the stored rasters) so we allocate the
    # tensors after reading the first band sets.
    images: list[torch.Tensor | None] = [None] * len(group_idxs)
    timestamps_list: list[list[tuple[datetime, datetime]] | None] = [None] * len(
        group_idxs
    )

    for band_set, src_indexes, dst_indexes in needed_sets_and_indexes:
        if band_set.format is None:
            raise ValueError(f"No format specified for {layer_name}")
        raster_format = band_set.instantiate_raster_format()

        # TODO: previously we try to read based on band_set.zoom_offset when possible,
        # and handle zooming in with torch.repeat (if resampling method is nearest
        # neighbor). However, we have not benchmarked whether this actually improves
        # data loading speed, so for simplicity, for now we let rasterio handle the
        # resampling. If it really is much faster to handle it via torch, then it may
        # make sense to bring back that functionality.

        per_group_arrays = window.data.read_rasters(
            layer_name,
            band_set.bands,
            group_idxs,
            raster_format,
            projection=final_projection,
            bounds=final_bounds,
            resampling=Resampling.nearest,
        )

        for slot, raster_array in enumerate(per_group_arrays):
            src = raster_array.array  # (C, T, H, W)
            image = images[slot]
            if image is None:
                t = src.shape[1]
                image = torch.zeros(
                    (
                        len(needed_bands),
                        t,
                        src.shape[2],
                        src.shape[3],
                    ),
                    dtype=get_torch_dtype(data_input.dtype),
                )
                images[slot] = image
                timestamps_list[slot] = raster_array.timestamps

            image[dst_indexes, :, :, :] = torch.as_tensor(
                src[src_indexes, :, :, :].astype(data_input.dtype.get_numpy_dtype())
            )

    final_images: list[torch.Tensor] = []
    for slot, image in enumerate(images):
        if image is None:
            raise RuntimeError(
                f"No band sets were read for layer {layer_name} group "
                f"{group_idxs[slot]} in window {window.name}, but found all "
                f"needed bands"
            )
        final_images.append(image)

    return final_images, timestamps_list

read_layer_time_range

read_layer_time_range(layer_data: WindowLayerData | None, group_idx: int) -> tuple[datetime, datetime] | None

Extract the combined time range from all items in a layer data group.

Returns the min start time and max end time across all items, or None if no items have time ranges.

Raises:

Type	Description
ValueError	If some items have time_range and others don't.

Source code in rslearn/train/dataset.py

def read_layer_time_range(
    layer_data: WindowLayerData | None, group_idx: int
) -> tuple[datetime, datetime] | None:
    """Extract the combined time range from all items in a layer data group.

    Returns the min start time and max end time across all items, or None if
    no items have time ranges.

    Raises:
        ValueError: If some items have time_range and others don't.
    """
    if layer_data is None:
        return None

    serialized_items = layer_data.serialized_item_groups[group_idx]
    if not serialized_items:
        return None

    first_item = Item.deserialize(serialized_items[0])
    if first_item.geometry.time_range is None:
        return None

    # If the first item has a time_range, all items must have one
    time_ranges: list[tuple[datetime, datetime]] = []
    for serialized_item in serialized_items:
        item = Item.deserialize(serialized_item)
        if item.geometry.time_range is None:
            raise ValueError(
                f"Item '{item.name}' has no time_range, but first item does. "
                "All items in a group must consistently have or lack time_range."
            )
        time_ranges.append(item.geometry.time_range)

    return (
        min(tr[0] for tr in time_ranges),
        max(tr[1] for tr in time_ranges),
    )

read_data_input

read_data_input(dataset: Dataset, window: Window, bounds: PixelBounds, data_input: DataInput, rng: Random) -> RasterImage | list[Feature]

Read the data specified by the DataInput from the window.

Parameters:

Name	Type	Description	Default
dataset	Dataset	the dataset, to get layer configs.	required
window	Window	the window to read from.	required
bounds	PixelBounds	the bounds of the patch we are reading.	required
data_input	DataInput	the DataInput that specifies what layers to read.	required
rng	Random	random number generator	required

Returns:

Type	Description
RasterImage | list[Feature]	the raster or vector data.

Source code in rslearn/train/dataset.py

def read_data_input(
    dataset: Dataset,
    window: Window,
    bounds: PixelBounds,
    data_input: DataInput,
    rng: random.Random,
) -> RasterImage | list[Feature]:
    """Read the data specified by the DataInput from the window.

    Args:
        dataset: the dataset, to get layer configs.
        window: the window to read from.
        bounds: the bounds of the patch we are reading.
        data_input: the DataInput that specifies what layers to read.
        rng: random number generator

    Returns:
        the raster or vector data.
    """
    # We first enumerate which item groups are available.
    # If load_all_item_groups is set, we discover all item groups within each layer.
    # Otherwise, we parse each entry in data_input.layers as a (layer_name, group_idx)
    # specifier.
    layer_options: list[tuple[str, int]] = []
    if data_input.load_all_item_groups:
        # We ensure that item groups are ordered within each layer, and across layers
        # we respect the user's given order.
        completed_groups_by_layer: dict[str, list[int]] = {}
        for layer_name, group_idx in window.list_completed_layers():
            if layer_name not in completed_groups_by_layer:
                completed_groups_by_layer[layer_name] = []
            completed_groups_by_layer[layer_name].append(group_idx)

        for layer_name in data_input.layers:
            cur_completed_groups = completed_groups_by_layer[layer_name]
            for group_idx in sorted(cur_completed_groups):
                layer_options.append((layer_name, group_idx))
    else:
        for option in data_input.layers:
            layer_name, group_idx = get_layer_and_group_from_dir_name(option)
            if not window.is_layer_completed(layer_name, group_idx):
                continue
            layer_options.append((layer_name, group_idx))

    # Now determine which item groups we should actually read.
    # We randomly pick one, unless load_all_layers is set, in which case we read all
    # available options.
    layers_to_read: list[tuple[str, int]]
    if data_input.load_all_layers:
        # We assume that the user has ensured the layers are compatible, e.g. raster
        # layers will need to have the same number of bands.
        layers_to_read = layer_options
    else:
        layers_to_read = [rng.choice(layer_options)]

    if not layers_to_read:
        raise ValueError(
            f"No completed layers found for data input with layers={data_input.layers}. "
            f"Available layer options: {layer_options}. "
            f"load_all_layers={data_input.load_all_layers}, "
            f"load_all_item_groups={data_input.load_all_item_groups}. "
            f"Check that the specified layers exist and are completed in the window."
        )

    if data_input.data_type == "raster":
        layer_datas = window.load_layer_datas()

        # Group entries by layer_name for calls to read_raster_layer_groups_for_data_input.
        groups_by_layer: dict[str, set[int]] = {}
        for layer_name, group_idx in layers_to_read:
            groups_by_layer.setdefault(layer_name, set()).add(group_idx)

        per_group_results: dict[
            tuple[str, int],
            tuple[torch.Tensor, list[tuple[datetime, datetime]] | None],
        ] = {}
        for layer_name, group_idxs_set in groups_by_layer.items():
            # Convert the set to a list, while also sorting so that if we are reading
            # all item groups then we go through the fast path.
            sorted_group_idxs = list(sorted(group_idxs_set))
            layer_config = dataset.layers[layer_name]
            layer_images, layer_timestamps = read_raster_layer_groups_for_data_input(
                window,
                bounds,
                layer_name,
                sorted_group_idxs,
                layer_config,
                data_input,
            )
            for group_idx, image, timestamps in zip(
                sorted_group_idxs, layer_images, layer_timestamps, strict=True
            ):
                per_group_results[(layer_name, group_idx)] = (image, timestamps)

        images: list[torch.Tensor] = []  # each is CTHW
        all_timestamps: list[tuple[datetime, datetime]] = []
        has_all_timestamps = True

        for layer_name, group_idx in layers_to_read:
            image, timestamps = per_group_results[(layer_name, group_idx)]
            images.append(image)

            if timestamps is not None:
                all_timestamps.extend(timestamps)
            elif image.shape[1] == 1:
                # For single-timestep RasterImage, fallback to item-level time range
                # for this item group.
                layer_data = layer_datas.get(layer_name)
                time_range = read_layer_time_range(layer_data, group_idx)
                if time_range is not None:
                    all_timestamps.append(time_range)
                    warnings.warn(
                        "Falling back to item-level time range for single-timestep "
                        "RasterImage is deprecated and will be removed after "
                        "2026-05-01. Ensure timestamps are stored with the raster data.",
                        FutureWarning,
                        stacklevel=2,
                    )
                else:
                    # It is okay for single-timestep RasterImage to not have timestamps.
                    has_all_timestamps = False
            else:
                # Multi-timestep RasterImage must have timestamps.
                raise ValueError(
                    f"Expected multi-timestep RasterImage with T={image.shape[1]} to have timestamps"
                )

        stacked = torch.cat(images, dim=1)
        return RasterImage(
            stacked,
            all_timestamps if has_all_timestamps else None,
        )

    elif data_input.data_type == "vector":
        # We don't really support time series for vector data currently, we just
        # concatenate the features together.
        features: list[Feature] = []
        for layer_name, group_idx in layers_to_read:
            layer_config = dataset.layers[layer_name]
            vector_format = layer_config.instantiate_vector_format()
            cur_features = window.data.read_vector(
                layer_name,
                vector_format,
                group_idx=group_idx,
            )
            features.extend(cur_features)

        return features

    else:
        raise ValueError(f"unknown data type {data_input.data_type}")

is_data_input_available

is_data_input_available(data_input: DataInput, window: Window) -> bool

Check if a data input's required item groups are available in a window.

Parameters:

Name	Type	Description	Default
data_input	DataInput	the data input to check.	required
window	Window	the window to check against.	required

Returns:

Type	Description
bool	True if the required item groups are available.

Source code in rslearn/train/dataset.py

def is_data_input_available(data_input: DataInput, window: Window) -> bool:
    """Check if a data input's required item groups are available in a window.

    Args:
        data_input: the data input to check.
        window: the window to check against.

    Returns:
        True if the required item groups are available.
    """
    # If load_all_layers is enabled, we need all entries present. Otherwise, just one.
    is_any_available = False
    are_all_available = True

    for option in data_input.layers:
        if data_input.load_all_item_groups:
            # In this case the option should be a layer name directly.
            # We can check group_idx=0 to verify there is at least one item group
            # present in this layer (since load_all_item_groups=true, the user doesn't
            # care how many item groups are completed).
            layer_name = option
            group_idx = 0
        else:
            # In this case it specifies an item group (like raster_layer.1).
            layer_name, group_idx = get_layer_and_group_from_dir_name(option)

        if window.is_layer_completed(layer_name, group_idx=group_idx):
            is_any_available = True
        else:
            are_all_available = False

    if data_input.load_all_layers:
        return are_all_available
    else:
        return is_any_available

check_window

check_window(inputs: dict[str, DataInput], window: Window, output_layer_name_skip_inference_if_exists: str | None = None) -> tuple[Window | None, CheckWindowResult]

Verify that the window has the required layers based on the specified inputs.

Parameters:

Name	Type	Description	Default
inputs	dict[str, DataInput]	the inputs to the dataset.	required
window	Window	the window to check.	required
output_layer_name_skip_inference_if_exists	str | None	optional name of the output layer to check for existence.	None

Returns:

Type	Description
Window | None	a tuple of (window, result) where window is the window if it passes all
CheckWindowResult	checks or None otherwise, and result records why the window was skipped.

Source code in rslearn/train/dataset.py

def check_window(
    inputs: dict[str, DataInput],
    window: Window,
    output_layer_name_skip_inference_if_exists: str | None = None,
) -> tuple[Window | None, CheckWindowResult]:
    """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:
        a tuple of (window, result) where window is the window if it passes all
        checks or None otherwise, and result records why the window was skipped.
    """
    for key, data_input in inputs.items():
        if not data_input.required:
            continue
        if not is_data_input_available(data_input, window):
            logger.debug(
                "Skipping window %s since check for input '%s' (layers %s) failed",
                window.name,
                key,
                data_input.layers,
            )
            return (None, CheckWindowResult(missing_data_input_counts={key: 1}))

    # 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, CheckWindowResult(has_output_layer_count=1))

    return (window, CheckWindowResult())