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

rslearn/models/olmoearth_pretrain/model.py

@@ -0,0 +1,421 @@
+"""OlmoEarth model wrapper for fine-tuning in rslearn."""
+
+import json
+import warnings
+from contextlib import nullcontext
+from datetime import datetime
+from typing import Any
+
+import torch
+from einops import rearrange
+from olmoearth_pretrain.config import Config, require_olmo_core
+from olmoearth_pretrain.data.constants import Modality
+from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample, MaskValue
+from olmoearth_pretrain.model_loader import (
+    ModelID,
+    load_model_from_id,
+    load_model_from_path,
+)
+from olmoearth_pretrain.nn.flexihelios import Encoder, TokensAndMasks
+from upath import UPath
+
+from rslearn.log_utils import get_logger
+from rslearn.models.component import FeatureExtractor, FeatureMaps, TokenFeatureMaps
+from rslearn.train.model_context import ModelContext, RasterImage
+
+logger = get_logger(__name__)
+
+MODALITY_NAMES = [
+    "sentinel2_l2a",
+    "sentinel1",
+    "worldcover",
+    "openstreetmap_raster",
+    "landsat",
+]
+
+AUTOCAST_DTYPE_MAP = {
+    "bfloat16": torch.bfloat16,
+    "float16": torch.float16,
+    "float32": torch.float32,
+}
+
+EMBEDDING_SIZES = {
+    ModelID.OLMOEARTH_V1_NANO: 128,
+    ModelID.OLMOEARTH_V1_TINY: 192,
+    ModelID.OLMOEARTH_V1_BASE: 768,
+    ModelID.OLMOEARTH_V1_LARGE: 1024,
+}
+
+
+class OlmoEarth(FeatureExtractor):
+    """A wrapper to support the OlmoEarth model."""
+
+    def __init__(
+        self,
+        patch_size: int,
+        model_id: ModelID | None = None,
+        model_path: str | None = None,
+        checkpoint_path: str | None = None,
+        selector: list[str | int] = ["encoder"],
+        forward_kwargs: dict[str, Any] = {},
+        random_initialization: bool = False,
+        embedding_size: int | None = None,
+        autocast_dtype: str | None = "bfloat16",
+        token_pooling: bool = True,
+        use_legacy_timestamps: bool = True,
+    ):
+        """Create a new OlmoEarth model.
+
+        Args:
+            patch_size: token spatial patch size to use.
+            model_id: the model ID to load. One of model_id or model_path or checkpoint_path must be
+                set.
+            model_path: the path to load the model from. One of model_id or model_path or checkpoint_path must be
+                set. Same structure as the HF-hosted `model_id` models: bundle with a config.json and weights.pth.
+            checkpoint_path: the checkpoint directory to load from, if model_id or model_path is not
+                set. It should contain a distributed checkpoint with a config.json file as well as model_and_optim
+                folder.
+            selector: an optional sequence of attribute names or list indices to select
+                the sub-module that should be applied on the input images. Defaults to
+                ["encoder"] to select only the transformer encoder.
+            forward_kwargs: additional arguments to pass to forward pass besides the
+                 MaskedOlmoEarthSample.
+            random_initialization: whether to skip loading the checkpoint so the
+                weights are randomly initialized. In this case, the checkpoint is only
+                used to define the model architecture.
+            embedding_size: optional embedding size to report via
+                get_backbone_channels (if model_id is not set).
+            autocast_dtype: which dtype to use for autocasting, or set None to disable.
+            token_pooling: whether or not to pool the tokens. If True, the output will be BxCxHxW. If False,
+                there will be an extra dimension, N, (BxCxHxWxN) representing the temporal and channel
+                dimensions.
+            use_legacy_timestamps: In our original implementation of OlmoEarth, we applied timestamps starting
+                from 0 (instead of the actual timestamps of the input). The option to do this is preserved
+                for backwards compatability with finetuned models which were trained against this implementation.
+        """
+        if use_legacy_timestamps:
+            warnings.warn(
+                "For new projects, don't use legacy timesteps.", DeprecationWarning
+            )
+
+        if (
+            sum(
+                [
+                    model_id is not None,
+                    model_path is not None,
+                    checkpoint_path is not None,
+                ]
+            )
+            != 1
+        ):
+            raise ValueError(
+                "exactly one of model_id, model_path, or checkpoint_path must be set"
+            )
+
+        super().__init__()
+        self.patch_size = patch_size
+        self.forward_kwargs = forward_kwargs
+        self.embedding_size = embedding_size
+
+        if autocast_dtype is not None:
+            self.autocast_dtype = AUTOCAST_DTYPE_MAP[autocast_dtype]
+        else:
+            self.autocast_dtype = None
+
+        if model_id is not None:
+            # Load from Hugging Face.
+            model = load_model_from_id(model_id, load_weights=not random_initialization)
+            if self.embedding_size is None and model_id in EMBEDDING_SIZES:
+                self.embedding_size = EMBEDDING_SIZES[model_id]
+
+        elif model_path is not None:
+            # Load from path.
+            model = load_model_from_path(
+                UPath(model_path), load_weights=not random_initialization
+            )
+
+        else:
+            # Load the distributed model checkpoint by path through Olmo Core
+            model = self._load_model_from_checkpoint(
+                UPath(checkpoint_path), random_initialization
+            )
+
+        # Select just the portion of the model that we actually want to use.
+        for part in selector:
+            if isinstance(part, str):
+                model = getattr(model, part)
+            else:
+                model = model[part]
+        self.model = model
+        self.token_pooling = token_pooling
+        self.use_legacy_timestamps = use_legacy_timestamps
+
+    def _load_model_from_checkpoint(
+        self, checkpoint_upath: UPath, random_initialization: bool
+    ) -> torch.nn.Module:
+        """Load the OlmoEarth pre-trained model from a distributed checkpoint folder.
+
+        The folder should contain config.json as well as the model_and_optim folder
+        that contains the distributed checkpoint. This is the format produced by
+        pre-training runs in olmoearth_pretrain.
+        """
+        with (checkpoint_upath / "config.json").open() as f:
+            config_dict = json.load(f)
+            model_config = Config.from_dict(config_dict["model"])
+
+        model = model_config.build()
+
+        # Load the checkpoint (requires olmo_core for distributed checkpoint loading).
+        if not random_initialization:
+            require_olmo_core(
+                "_load_model_from_checkpoint with random_initialization=False"
+            )
+            from olmo_core.distributed.checkpoint import load_model_and_optim_state
+
+            train_module_dir = checkpoint_upath / "model_and_optim"
+            load_model_and_optim_state(str(train_module_dir), model)
+            logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
+
+        return model
+
+    @staticmethod
+    def time_ranges_to_timestamps(
+        time_ranges: list[tuple[datetime, datetime]],
+        max_timestamps: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        """Turn the time ranges stored in a RasterImage to timestamps accepted by OlmoEarth.
+
+        OlmoEarth only uses the month associated with each timestamp, so we take the midpoint
+        the time range. For some inputs (e.g. Sentinel 2) we take an image from a specific
+        time so that start_time == end_time == mid_time.
+        """
+        timestamps = torch.zeros((max_timestamps, 3), dtype=torch.int32, device=device)
+        mid_ranges = [t[0] + ((t[1] - t[0]) / 2) for t in time_ranges]
+        timestamps[: len(time_ranges), 0] = torch.tensor(
+            [d.day for d in mid_ranges], dtype=torch.int32
+        )
+        # months are indexed 0-11
+        timestamps[: len(time_ranges), 1] = torch.tensor(
+            [d.month - 1 for d in mid_ranges], dtype=torch.int32
+        )
+        timestamps[: len(time_ranges), 2] = torch.tensor(
+            [d.year for d in mid_ranges], dtype=torch.int32
+        )
+        return timestamps
+
+    def _prepare_modality_inputs(
+        self, context: ModelContext
+    ) -> tuple[MaskedOlmoEarthSample, list[str], torch.device]:
+        """Prepare modality tensors and masks for the OlmoEarth model.
+
+        Uses a two-pass approach to ensure all modalities have consistent timestep
+        dimensions for position encoding.
+
+        Args:
+            context: the model context with input tensors.
+
+        Returns:
+            tuple of (sample, present_modalities, device)
+        """
+        kwargs = {}
+        present_modalities = []
+        device = None
+
+        # First pass: find global max_timesteps across all modalities and samples
+        # TODO: currently we assume all modalities have the same number of timesteps,
+        # which is not true for all cases, and time series time steps are assumed to
+        # be 1-month apart. It also assumes continuity between available timesteps.
+        # We'll have to fix all that.
+        max_timesteps = 1
+        modality_data = {}
+        # we will just store the longest time range
+        # per instance in the batch. This means it may not be
+        # aligned per modality
+        timestamps_per_instance: list[list[tuple[datetime, datetime]]] = [[]] * len(
+            context.inputs
+        )
+        for modality in MODALITY_NAMES:
+            if modality not in context.inputs[0]:
+                continue
+            present_modalities.append(modality)
+            tensors = []
+            for idx, inp in enumerate(context.inputs):
+                assert isinstance(inp[modality], RasterImage)
+                tensors.append(inp[modality].image)
+                cur_timestamps = inp[modality].timestamps
+                if cur_timestamps is not None and len(cur_timestamps) > len(
+                    timestamps_per_instance[idx]
+                ):
+                    timestamps_per_instance[idx] = cur_timestamps
+            tensors = [inp[modality].image for inp in context.inputs]
+            device = tensors[0].device
+            max_t = max(t.shape[1] for t in tensors)
+            max_timesteps = max(max_timesteps, max_t)
+            modality_data[modality] = (
+                tensors,
+                len(Modality.get(modality).band_sets),
+            )
+
+        # Second pass: pad and process each modality with global max_timesteps
+        for modality in present_modalities:
+            tensors, num_band_sets = modality_data[modality]
+
+            # Pad tensors to target_ch and track original timesteps for masking
+            padded = []
+            original_timesteps = []
+            for t in tensors:
+                orig_t = t.shape[1]
+                original_timesteps.append(orig_t)
+                if orig_t < max_timesteps:
+                    pad = torch.zeros(
+                        t.shape[:1] + (max_timesteps - orig_t,) + t.shape[2:],
+                        dtype=t.dtype,
+                        device=device,
+                    )
+                    t = torch.cat([t, pad], dim=1)
+                padded.append(t)
+
+            cur = torch.stack(padded, dim=0)
+            cur = rearrange(cur, "b c t h w -> b h w t c")
+            kwargs[modality] = cur
+
+            # Create mask: ONLINE_ENCODER for valid, MISSING for padded timesteps
+            b, h, w = cur.shape[0], cur.shape[1], cur.shape[2]
+            mask = torch.full(
+                (b, h, w, max_timesteps, num_band_sets),
+                fill_value=MaskValue.ONLINE_ENCODER.value,
+                dtype=torch.int32,
+                device=device,
+            )
+            for sample_idx, orig_t in enumerate(original_timesteps):
+                if orig_t < max_timesteps:
+                    mask[sample_idx, :, :, orig_t:, :] = MaskValue.MISSING.value
+            kwargs[f"{modality}_mask"] = mask
+
+        if self.use_legacy_timestamps:
+            # Note that only months (0 to 11) are used in OlmoEarth position encoding.
+            timestamps = torch.zeros(
+                (len(context.inputs), max_timesteps, 3),
+                dtype=torch.int32,
+                device=device,
+            )
+            timestamps[:, :, 0] = 1  # day
+            timestamps[:, :, 1] = torch.arange(max_timesteps, device=device)[
+                None, :
+            ]  # month
+            timestamps[:, :, 2] = 2024  # year
+            kwargs["timestamps"] = timestamps
+        else:
+            if max([len(t) for t in timestamps_per_instance]) == 0:
+                # Timestamps is required.
+                raise ValueError("No inputs had timestamps.")
+            # Note that only months (0 to 11) are used in OlmoEarth position encoding.
+            kwargs["timestamps"] = torch.stack(
+                [
+                    self.time_ranges_to_timestamps(time_range, max_timesteps, device)
+                    for time_range in timestamps_per_instance
+                ],
+                dim=0,
+            )
+
+        return MaskedOlmoEarthSample(**kwargs), present_modalities, device
+
+    def forward(self, context: ModelContext) -> FeatureMaps | TokenFeatureMaps:
+        """Compute feature maps from the OlmoEarth backbone.
+
+        Args:
+            context: the model context. Input dicts should include keys corresponding
+                to the modalities that should be passed to the OlmoEarth model.
+
+        Returns:
+            a FeatureMaps consisting of one feature map, at 1/patch_size of the input
+                resolution. Embeddings will be pooled across modalities and timesteps.
+        """
+        sample, present_modalities, device = self._prepare_modality_inputs(context)
+
+        # Decide context based on self.autocast_dtype.
+        if self.autocast_dtype is None:
+            torch_context = nullcontext()
+        else:
+            assert device is not None
+            torch_context = torch.amp.autocast(
+                device_type=device.type, dtype=self.autocast_dtype
+            )
+
+        # Check if we can bypass masks (fast_pass=True)
+        missing_tokens = False
+        for modality in present_modalities:
+            modality_mask = getattr(sample, f"{modality}_mask")
+            if torch.any(modality_mask == MaskValue.MISSING.value):
+                missing_tokens = True
+                break
+
+        with torch_context:
+            # Currently we assume the provided model always returns a TokensAndMasks object.
+            tokens_and_masks: TokensAndMasks
+            if isinstance(self.model, Encoder):
+                # Encoder has a fast_pass argument to indicate mask is not needed.
+                tokens_and_masks = self.model(
+                    sample,
+                    fast_pass=not missing_tokens,
+                    patch_size=self.patch_size,
+                    **self.forward_kwargs,
+                )["tokens_and_masks"]
+            else:
+                # Other models like STEncoder do not have this option supported.
+                tokens_and_masks = self.model(
+                    sample, patch_size=self.patch_size, **self.forward_kwargs
+                )["tokens_and_masks"]
+
+        # Apply temporal/modality pooling so we just have one feature per patch.
+        features = []
+        if self.token_pooling:
+            for modality in present_modalities:
+                modality_features = getattr(tokens_and_masks, modality)  # BHWTSC
+                # If fast_pass is False, we need to mask the missing tokens before pooling.
+                if missing_tokens:
+                    modality_masks = getattr(
+                        tokens_and_masks, f"{modality}_mask"
+                    )  # BHWTS
+                    modality_masks_bool = (
+                        modality_masks != MaskValue.MISSING.value
+                    ).unsqueeze(-1)
+                    count = modality_masks_bool.sum(dim=[3, 4])
+                    # Masked average over band sets and timesteps (BHWTSC -> BHWC).
+                    pooled = (modality_features * modality_masks_bool).sum(
+                        dim=[3, 4]
+                    ) / count.clamp(min=1)
+                else:
+                    # Pool over band sets and timesteps (BHWTSC -> BHWC).
+                    pooled = modality_features.mean(dim=[3, 4])
+                # We want BHWC -> BCHW.
+                pooled = rearrange(pooled, "b h w c -> b c h w")
+                features.append(pooled)
+            # Pool over the modalities, so we get one BCHW feature map.
+            pooled = torch.stack(features, dim=0).mean(dim=0)
+            return FeatureMaps([pooled])
+        else:
+            for modality in present_modalities:
+                modality_features = getattr(tokens_and_masks, modality)
+                # Combine band sets and timesteps into last dim (BHWTSC -> BHWCN).
+                modality_features = rearrange(
+                    modality_features, "b h w t s c -> b c h w (t s)"
+                )
+                features.append(modality_features)
+            pooled = torch.cat(features, dim=-1)
+            return TokenFeatureMaps([pooled])
+
+    def get_backbone_channels(self) -> list:
+        """Returns the output channels of this model when used as a backbone.
+
+        The output channels is a list of (downsample_factor, depth) that corresponds
+        to the feature maps that the backbone returns. For example, an element [2, 32]
+        indicates that the corresponding feature map is 1/2 the input resolution and
+        has 32 channels.
+
+        Returns:
+            the output channels of the backbone as a list of (downsample_factor, depth)
+            tuples.
+        """
+        return [(self.patch_size, self.embedding_size)]

rslearn/models/olmoearth_pretrain/norm.py

@@ -0,0 +1,86 @@
+"""Normalization transforms."""
+
+import json
+from typing import Any
+
+from olmoearth_pretrain.data.normalize import load_computed_config
+
+from rslearn.log_utils import get_logger
+from rslearn.train.transforms.transform import Transform
+
+logger = get_logger(__file__)
+
+
+class OlmoEarthNormalize(Transform):
+    """Normalize using OlmoEarth JSON config.
+
+    For Sentinel-1 data, the values should be converted to decibels before being passed
+    to this transform.
+    """
+
+    def __init__(
+        self,
+        band_names: dict[str, list[str]],
+        std_multiplier: float | None = 2,
+        config_fname: str | None = None,
+    ) -> None:
+        """Initialize a new OlmoEarthNormalize.
+
+        Args:
+            band_names: map from modality name to the list of bands in that modality in
+                the order they are being loaded. Note that this order must match the
+                expected order for the OlmoEarth model.
+            std_multiplier: the std multiplier matching the one used for the model
+                training in OlmoEarth.
+            config_fname: load the normalization configuration from this file, instead
+                of getting it from OlmoEarth.
+        """
+        super().__init__()
+        self.band_names = band_names
+        self.std_multiplier = std_multiplier
+
+        if config_fname is None:
+            self.norm_config = load_computed_config()
+        else:
+            logger.warning(
+                f"Loading normalization config from {config_fname}. This argument is deprecated and will be removed in a future version."
+            )
+            with open(config_fname) as f:
+                self.norm_config = json.load(f)
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Apply normalization over the inputs and targets.
+
+        Args:
+            input_dict: the input
+            target_dict: the target
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        for modality_name, cur_band_names in self.band_names.items():
+            band_norms = self.norm_config[modality_name]
+            image = input_dict[modality_name]
+            # Keep a set of indices to make sure that we normalize all of them.
+            needed_band_indices = set(range(image.image.shape[0]))
+            num_timesteps = image.image.shape[0] // len(cur_band_names)
+
+            for band, norm_dict in band_norms.items():
+                # If multitemporal, normalize each timestep separately.
+                for t in range(num_timesteps):
+                    band_idx = cur_band_names.index(band) + t * len(cur_band_names)
+                    min_val = norm_dict["mean"] - self.std_multiplier * norm_dict["std"]
+                    max_val = norm_dict["mean"] + self.std_multiplier * norm_dict["std"]
+                    image.image[band_idx] = (image.image[band_idx] - min_val) / (
+                        max_val - min_val
+                    )
+                    needed_band_indices.remove(band_idx)
+
+            if len(needed_band_indices) > 0:
+                raise ValueError(
+                    f"for modality {modality_name}, bands {needed_band_indices} were unexpectedly not normalized"
+                )
+
+        return input_dict, target_dict

rslearn/models/panopticon.py

@@ -0,0 +1,170 @@
+"""Wrapper for the Panopticon model."""
+
+import math
+from enum import StrEnum
+from importlib import resources
+
+import torch
+import torch.nn.functional as F
+import yaml
+from einops import rearrange, repeat
+
+from rslearn.log_utils import get_logger
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
+logger = get_logger(__name__)
+
+
+class PanopticonModalities(StrEnum):
+    """Modalities supported by Panopticon.
+
+    These are the keys needed to load the yaml file from panopticon_data/sensors
+    """
+
+    SENTINEL2 = "sentinel2"
+    LANDSAT8 = "landsat8"
+    SENTINEL1 = "sentinel1"
+    # Add more modalities as needed
+
+
+class Panopticon(FeatureExtractor):
+    """Class containing the Panopticon model that can ingest MaskedHeliosSample objects."""
+
+    patch_size: int = 14
+    base_image_size: int = 224
+
+    def __init__(
+        self,
+        band_order: dict[str, list[str]],
+        torchhub_id: str = "panopticon_vitb14",
+    ):
+        """Initialize the Panopticon wrapper.
+
+        Args:
+            band_order: The band order for the panopticon model, must match the specified order in the data config
+            torchhub_id: The torch hub model ID for panopticon
+        """
+        super().__init__()
+        # Load the panopticon model
+        self._load_model(torchhub_id)
+        self.output_dim = self.model.embed_dim
+        self.band_order = band_order
+        self.supported_modalities = list(band_order.keys())
+
+    def _load_model(self, torchhub_id: str) -> None:
+        """Load the panopticon model from torch hub."""
+        import time
+
+        # Hack to get around https://discuss.pytorch.org/t/torch-hub-load-gives-httperror-rate-limit-exceeded/124769
+        torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
+        for attempt in range(2):
+            try:
+                self.model = torch.hub.load(  # nosec B614
+                    "panopticon-FM/panopticon",
+                    torchhub_id,
+                )
+                break
+            except Exception as e:
+                logger.warning(
+                    f"Error loading panopticon model: {e}. Retrying in 5 seconds..."
+                )
+                time.sleep(5)
+        else:
+            raise RuntimeError(
+                f"Failed to load panopticon model {torchhub_id} after retrying."
+            )
+
+    def _process_modality_data(self, data: torch.Tensor) -> torch.Tensor:
+        """Process individual modality data.
+
+        Args:
+            data: Input tensor of shape [B, C, H, W]
+
+        Returns:
+            Processed tensor of shape [B, C, H, W]
+        """
+        original_height = data.shape[2]
+        new_height = self.patch_size if original_height == 1 else self.base_image_size
+
+        data = F.interpolate(
+            data,
+            size=(new_height, new_height),
+            mode="bilinear",
+            align_corners=False,
+        )
+        return data
+
+    def _create_channel_ids(
+        self, modality: str, batch_size: int, device: torch.device
+    ) -> torch.Tensor:
+        """Create channel IDs for the panopticon model."""
+        with resources.open_text(
+            "rslearn.models.panopticon_data.sensors", f"{modality}.yaml"
+        ) as f:
+            sensor_config = yaml.safe_load(f)
+
+        band_order = self.band_order[modality]
+        chn_ids = [
+            sensor_config["bands"][band.upper()]["gaussian"]["mu"]
+            for band in band_order
+        ]
+        chn_ids = torch.tensor(chn_ids, dtype=torch.float32, device=device)
+        chn_ids = repeat(chn_ids, "c -> b c", b=batch_size)
+        return chn_ids
+
+    def prepare_input(
+        self, input_data: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        """Prepare input for the panopticon model from MaskedHeliosSample."""
+        channel_ids_list: list[torch.Tensor] = []
+        processed_data_list: list[torch.Tensor] = []
+        for modality in self.supported_modalities:
+            if modality not in input_data.keys():
+                logger.debug(f"Modality {modality} not found in input data")
+                continue
+            data = input_data[modality]
+            device = data.device
+            processed_data = self._process_modality_data(data)
+            processed_data_list.append(processed_data)
+            batch_size = processed_data.shape[0]
+            chn_ids = self._create_channel_ids(modality, batch_size, device)
+            channel_ids_list.append(chn_ids)
+
+        processed_data = torch.cat(processed_data_list, dim=1)
+        chn_ids = torch.cat(channel_ids_list, dim=1)
+        return {
+            "imgs": processed_data,
+            "chn_ids": chn_ids,
+        }
+
+    def forward(self, context: ModelContext) -> FeatureMaps:
+        """Forward pass through the panopticon model."""
+        batch_inputs = {
+            key: torch.stack(
+                [inp[key].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
+            )
+            for key in context.inputs[0].keys()
+        }
+        panopticon_inputs = self.prepare_input(batch_inputs)
+        output_features = self.model.forward_features(panopticon_inputs)[
+            "x_norm_patchtokens"
+        ]
+
+        num_tokens = output_features.shape[1]
+        height = int(math.sqrt(num_tokens))
+        output_features = rearrange(
+            output_features, "b (h w) d -> b d h w", h=height, w=height
+        )
+        return FeatureMaps([output_features])
+
+    def get_backbone_channels(self) -> list:
+        """Returns the output channels of this model when used as a backbone.
+
+        The output channels is a list of (downsample_factor, depth) that corresponds
+        to the feature maps that the backbone returns. For example, an element [2, 32]
+        indicates that the corresponding feature map is 1/2 the input resolution and
+        has 32 channels.
+        """
+        return [(self.patch_size, self.output_dim)]