`rslearn.models.prithvi`¶

prithvi ¶

Prithvi V2.

This code is adapted from https://github.com/NASA-IMPACT/Prithvi-WxC

The code is released under:

PrithviV2Models ¶

Bases: StrEnum

Names for different Prithvi models on torch hub.

Source code in rslearn/models/prithvi.py

class PrithviV2Models(StrEnum):
    """Names for different Prithvi models on torch hub."""

    VIT_300 = "VIT_300"
    VIT_600 = "VIT_600"

PrithviV2 ¶

Bases: FeatureExtractor

An Rslearn wrapper for Prithvi 2.0.

Source code in rslearn/models/prithvi.py

class PrithviV2(FeatureExtractor):
    """An Rslearn wrapper for Prithvi 2.0."""

    INPUT_KEY = "image"

    def __init__(
        self,
        cache_dir: str | Path | None = None,
        size: PrithviV2Models = PrithviV2Models.VIT_300,
        num_frames: int = 1,
    ):
        """Create a new PrithviV2.

        Args:
            cache_dir: The local folder in which to download the prithvi config and
                weights. If None, it downloads to a temporary folder.
            size: the model size, see class for various models.
            num_frames: The number of input frames (timesteps). The model was trained on 3,
                but if there is just one timestamp examples use 1 (e.g.
                https://github.com/NASA-IMPACT/Prithvi-EO-2.0/blob/main/examples/
                example_landslide4sense.ipynb)

        """
        super().__init__()
        if cache_dir is None:
            cache_dir = DEFAULT_CACHE_DIR
        cache_dir = Path(cache_dir)

        hub_id = MODEL_TO_HF_INFO[size]["hf_hub_id"]
        revision = MODEL_TO_HF_INFO[size]["revision"]
        checkpoint_fname = MODEL_TO_HF_INFO[size]["weights"]

        config = get_config(cache_dir, hub_id, revision)
        config["num_frames"] = num_frames
        self.model = PrithviMAE(**config)

        if not (cache_dir / checkpoint_fname).exists():
            _ = hf_hub_download(
                local_dir=cache_dir,
                repo_id=hub_id,
                filename=checkpoint_fname,
                revision=revision,
            )  # nosec

        state_dict = torch.load(
            cache_dir / checkpoint_fname,
            map_location="cpu",
            weights_only=True,
        )
        # discard fixed pos_embedding weight, following
        # https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M/blob/e4aabdc440c8ee703a749def8af5bf4700dee35b/inference.py#L362
        for k in list(state_dict.keys()):
            if "pos_embed" in k:
                del state_dict[k]
        self.model.load_state_dict(state_dict, strict=False)
        self.image_resolution = config["img_size"]
        self.bands = config["bands"]
        # patch size is a list [t, h, w], where h == w
        self.patch_size = config["patch_size"][-1]

    def _resize_data(self, data: torch.Tensor) -> torch.Tensor:
        """Process individual modality data.

        Args:
            data: Input tensor of shape [B, C, T, H, W]

        Returns:
            list of tensors of shape [B, C, T, H, W]
        """
        # Get original dimensions
        B, C, T, H, W = data.shape
        data = rearrange(data, "b c t h w -> b (c t) h w")
        original_height = H
        new_height = self.patch_size if original_height == 1 else self.image_resolution
        data = F.interpolate(
            data,
            size=(new_height, new_height),
            mode="bilinear",
            align_corners=False,
        )
        data = rearrange(data, "b (c t) h w -> b c t h w", c=C, t=T)
        return data

    def forward(self, context: ModelContext) -> FeatureMaps:
        """Compute feature maps from the Prithvi V2 backbone.

        Args:
            context: the model context. Input dicts must include "image" key containing
                HLS (Harmonized Landsat-Sentinel) data.

        Returns:
            a FeatureMaps with one map of shape [B, H/p_s, W/p_s, 11*1024] that contains stacked
                feature maps across the 11 transformer blocks.
        """
        # x has shape BCTHW
        x = torch.stack([inp[self.INPUT_KEY].image for inp in context.inputs], dim=0)
        x = self._resize_data(x)
        features = self.model.encoder.forward_features(x)
        # prepare_features_for_image_model was slightly modified since we already
        # know the number of timesteps and don't need to recompute it.
        # in addition we average along the time dimension (instead of concatenating)
        # to keep the embeddings reasonably sized.
        result = self.model.encoder.prepare_features_for_image_model(
            [features[-1]], x.shape[2]
        )
        return FeatureMaps(result)

    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 (patch_size, depth) that corresponds
        to the feature maps that the backbone returns.

        Returns:
            the output channels of the backbone as a list of (patch_size, depth) tuples.
        """
        return [(1, 1024)]

forward ¶

forward(context: ModelContext) -> FeatureMaps

Compute feature maps from the Prithvi V2 backbone.

Parameters:

Name	Type	Description	Default
`context`	`ModelContext`	the model context. Input dicts must include "image" key containing HLS (Harmonized Landsat-Sentinel) data.	required

Returns:

Type	Description
`FeatureMaps`	a FeatureMaps with one map of shape [B, H/p_s, W/p_s, 11*1024] that contains stacked feature maps across the 11 transformer blocks.

Source code in rslearn/models/prithvi.py

def forward(self, context: ModelContext) -> FeatureMaps:
    """Compute feature maps from the Prithvi V2 backbone.

    Args:
        context: the model context. Input dicts must include "image" key containing
            HLS (Harmonized Landsat-Sentinel) data.

    Returns:
        a FeatureMaps with one map of shape [B, H/p_s, W/p_s, 11*1024] that contains stacked
            feature maps across the 11 transformer blocks.
    """
    # x has shape BCTHW
    x = torch.stack([inp[self.INPUT_KEY].image for inp in context.inputs], dim=0)
    x = self._resize_data(x)
    features = self.model.encoder.forward_features(x)
    # prepare_features_for_image_model was slightly modified since we already
    # know the number of timesteps and don't need to recompute it.
    # in addition we average along the time dimension (instead of concatenating)
    # to keep the embeddings reasonably sized.
    result = self.model.encoder.prepare_features_for_image_model(
        [features[-1]], x.shape[2]
    )
    return FeatureMaps(result)

get_backbone_channels ¶

get_backbone_channels() -> list

Returns the output channels of this model when used as a backbone.

The output channels is a list of (patch_size, depth) that corresponds to the feature maps that the backbone returns.

Returns:

Type	Description
`list`	the output channels of the backbone as a list of (patch_size, depth) tuples.

Source code in rslearn/models/prithvi.py

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 (patch_size, depth) that corresponds
    to the feature maps that the backbone returns.

    Returns:
        the output channels of the backbone as a list of (patch_size, depth) tuples.
    """
    return [(1, 1024)]

PrithviNormalize ¶

Bases: Transform

Normalize inputs using Prithvi normalization.

Similar to the model, the input should be an image time series under the key "image".

Source code in rslearn/models/prithvi.py

class PrithviNormalize(Transform):
    """Normalize inputs using Prithvi normalization.

    Similar to the model, the input should be an image time series under the key
    "image".
    """

    def __init__(
        self,
        cache_dir: str | Path | None = None,
        size: PrithviV2Models = PrithviV2Models.VIT_300,
    ) -> None:
        """Initialize a new PrithviNormalize.

        Args:
            cache_dir: the local directory to cache the config.json which contains the
                means and standard deviations used in the normalization.
            size: the model size, see class for various models. In this case (and
                for the current hf revision), the config values (mean and std) are the
                same for both the 300M and 600M model, so its safe to not set this.
        """
        super().__init__()
        hub_id = MODEL_TO_HF_INFO[size]["hf_hub_id"]
        revision = MODEL_TO_HF_INFO[size]["revision"]
        if cache_dir is None:
            cache_dir = DEFAULT_CACHE_DIR
        cache_dir = Path(cache_dir)
        config = get_config(cache_dir, hub_id, revision)
        self.normalizer = Normalize(
            mean=config["mean"],
            std=config["std"],
            selectors=[PrithviV2.INPUT_KEY],
        )

    def forward(
        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        """Apply Prithvi normalization on the image.

        Args:
            input_dict: the input, which must contain the "image" key.
            target_dict: the target

        Returns:
            normalized (input_dicts, target_dicts) tuple
        """
        return self.normalizer(input_dict, target_dict)

forward ¶

forward(input_dict: dict[str, Any], target_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]

Apply Prithvi normalization on the image.

Parameters:

Name	Type	Description	Default
`input_dict`	`dict[str, Any]`	the input, which must contain the "image" key.	required
`target_dict`	`dict[str, Any]`	the target	required

Returns:

Type	Description
`tuple[dict[str, Any], dict[str, Any]]`	normalized (input_dicts, target_dicts) tuple

Source code in rslearn/models/prithvi.py

def forward(
    self, input_dict: dict[str, Any], target_dict: dict[str, Any]
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Apply Prithvi normalization on the image.

    Args:
        input_dict: the input, which must contain the "image" key.
        target_dict: the target

    Returns:
        normalized (input_dicts, target_dicts) tuple
    """
    return self.normalizer(input_dict, target_dict)

PatchEmbed ¶

Bases: Module

3D version of timm.models.vision_transformer.PatchEmbed.

Source code in rslearn/models/prithvi.py

class PatchEmbed(nn.Module):
    """3D version of timm.models.vision_transformer.PatchEmbed."""

    def __init__(
        self,
        input_size: tuple[int, int, int] = (1, 224, 224),
        patch_size: tuple[int, int, int] = (1, 16, 16),
        in_chans: int = 3,
        embed_dim: int = 768,
        norm_layer: nn.Module | None = None,
        flatten: bool = True,
        bias: bool = True,
    ) -> None:
        """Init."""
        super().__init__()
        self.input_size = input_size
        self.patch_size = patch_size
        self.grid_size = [s // p for s, p in zip(self.input_size, self.patch_size)]
        assert self.grid_size >= [1, 1, 1], "Patch size is bigger than input size."
        self.num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2]
        self.flatten = flatten

        self.proj = nn.Conv3d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias
        )
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward."""
        B, C, T, H, W = x.shape

        if (
            T / self.patch_size[0] % 1
            or H / self.patch_size[1] % 1
            or W / self.patch_size[2] % 1
        ):
            warnings.warn(
                f"Input {x.shape[-3:]} is not divisible by patch size {self.patch_size}."
                f"The border will be ignored, add backbone_padding for pixel-wise tasks."
            )

        x = self.proj(x)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # B,C,T,H,W -> B,C,L -> B,L,C
        x = self.norm(x)
        return x

forward ¶

forward(x: Tensor) -> Tensor

Forward.

Source code in rslearn/models/prithvi.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward."""
    B, C, T, H, W = x.shape

    if (
        T / self.patch_size[0] % 1
        or H / self.patch_size[1] % 1
        or W / self.patch_size[2] % 1
    ):
        warnings.warn(
            f"Input {x.shape[-3:]} is not divisible by patch size {self.patch_size}."
            f"The border will be ignored, add backbone_padding for pixel-wise tasks."
        )

    x = self.proj(x)
    if self.flatten:
        x = x.flatten(2).transpose(1, 2)  # B,C,T,H,W -> B,C,L -> B,L,C
    x = self.norm(x)
    return x

TemporalEncoder ¶

Bases: Module

TemporalEncoder.

Source code in rslearn/models/prithvi.py

class TemporalEncoder(nn.Module):
    """TemporalEncoder."""

    def __init__(self, embed_dim: int, trainable_scale: bool = False):
        """Init."""
        super().__init__()
        self.embed_dim = embed_dim
        self.year_embed_dim = embed_dim // 2
        self.julian_day_embed_dim = embed_dim - self.year_embed_dim

        # If trainable, initialize scale with small number
        if trainable_scale:
            self.scale = nn.Parameter(torch.full((1,), 0.1))
        else:
            self.register_buffer("scale", torch.ones(1))

    def forward(
        self, temporal_coords: torch.Tensor, tokens_per_frame: int | None = None
    ) -> torch.Tensor:
        """Forward.

        temporal_coords: year and day-of-year info with shape (B, T, 2).
        tokens_per_frame: number of tokens for each frame in the sample. If provided, embeddings will be
            repeated over T dimension, and final shape is (B, T*tokens_per_frame, embed_dim).
        """
        shape = temporal_coords.shape[:2] + (-1,)  # B, T, -1

        year = _get_1d_sincos_embed_from_grid_torch(
            self.year_embed_dim, temporal_coords[:, :, 0].flatten()
        ).reshape(shape)
        julian_day = _get_1d_sincos_embed_from_grid_torch(
            self.julian_day_embed_dim, temporal_coords[:, :, 1].flatten()
        ).reshape(shape)

        embedding = self.scale * torch.cat([year, julian_day], dim=-1)

        if tokens_per_frame is not None:
            embedding = torch.repeat_interleave(embedding, tokens_per_frame, dim=1)

        return embedding  # B, T*tokens_per_frame, embed_dim

forward ¶

forward(temporal_coords: Tensor, tokens_per_frame: int | None = None) -> Tensor

Forward.

temporal_coords: year and day-of-year info with shape (B, T, 2). tokens_per_frame: number of tokens for each frame in the sample. If provided, embeddings will be repeated over T dimension, and final shape is (B, T*tokens_per_frame, embed_dim).

Source code in rslearn/models/prithvi.py

def forward(
    self, temporal_coords: torch.Tensor, tokens_per_frame: int | None = None
) -> torch.Tensor:
    """Forward.

    temporal_coords: year and day-of-year info with shape (B, T, 2).
    tokens_per_frame: number of tokens for each frame in the sample. If provided, embeddings will be
        repeated over T dimension, and final shape is (B, T*tokens_per_frame, embed_dim).
    """
    shape = temporal_coords.shape[:2] + (-1,)  # B, T, -1

    year = _get_1d_sincos_embed_from_grid_torch(
        self.year_embed_dim, temporal_coords[:, :, 0].flatten()
    ).reshape(shape)
    julian_day = _get_1d_sincos_embed_from_grid_torch(
        self.julian_day_embed_dim, temporal_coords[:, :, 1].flatten()
    ).reshape(shape)

    embedding = self.scale * torch.cat([year, julian_day], dim=-1)

    if tokens_per_frame is not None:
        embedding = torch.repeat_interleave(embedding, tokens_per_frame, dim=1)

    return embedding  # B, T*tokens_per_frame, embed_dim

LocationEncoder ¶

Bases: Module

LocationEncoder.

Source code in rslearn/models/prithvi.py

class LocationEncoder(nn.Module):
    """LocationEncoder."""

    def __init__(self, embed_dim: int, trainable_scale: bool = False):
        """Init."""
        super().__init__()
        self.embed_dim = embed_dim
        self.lat_embed_dim = embed_dim // 2
        self.lon_embed_dim = embed_dim - self.lat_embed_dim

        # If trainable, initialize scale with small number
        if trainable_scale:
            self.scale = nn.Parameter(torch.full((1,), 0.1))
        else:
            self.register_buffer("scale", torch.ones(1))

    def forward(self, location_coords: torch.Tensor) -> torch.Tensor:
        """location_coords: lat and lon info with shape (B, 2)."""
        shape = location_coords.shape[:1] + (1, -1)  # B, 1, -1

        lat = _get_1d_sincos_embed_from_grid_torch(
            self.lat_embed_dim, location_coords[:, 0].flatten()
        ).reshape(shape)
        lon = _get_1d_sincos_embed_from_grid_torch(
            self.lon_embed_dim, location_coords[:, 1].flatten()
        ).reshape(shape)

        embedding = self.scale * torch.cat([lat, lon], dim=-1)

        return embedding  # B, 1, embed_dim

forward ¶

forward(location_coords: Tensor) -> Tensor

location_coords: lat and lon info with shape (B, 2).

Source code in rslearn/models/prithvi.py

def forward(self, location_coords: torch.Tensor) -> torch.Tensor:
    """location_coords: lat and lon info with shape (B, 2)."""
    shape = location_coords.shape[:1] + (1, -1)  # B, 1, -1

    lat = _get_1d_sincos_embed_from_grid_torch(
        self.lat_embed_dim, location_coords[:, 0].flatten()
    ).reshape(shape)
    lon = _get_1d_sincos_embed_from_grid_torch(
        self.lon_embed_dim, location_coords[:, 1].flatten()
    ).reshape(shape)

    embedding = self.scale * torch.cat([lat, lon], dim=-1)

    return embedding  # B, 1, embed_dim

PrithviViT ¶

Bases: Module

Prithvi ViT Encoder.

Source code in rslearn/models/prithvi.py

class PrithviViT(nn.Module):
    """Prithvi ViT Encoder."""

    def __init__(
        self,
        img_size: int | tuple[int, int] = 224,
        patch_size: int | tuple[int, int, int] = (1, 16, 16),
        num_frames: int = 1,
        in_chans: int = 3,
        embed_dim: int = 1024,
        depth: int = 24,
        num_heads: int = 16,
        mlp_ratio: float = 4.0,
        norm_layer: nn.Module = nn.LayerNorm,
        coords_encoding: list[str] | None = None,
        coords_scale_learn: bool = False,
        drop_path: float = 0.0,
        **kwargs: Any,
    ) -> None:
        """Init."""
        super().__init__()

        self.in_chans = in_chans
        self.num_frames = num_frames
        self.embed_dim = embed_dim
        self.img_size = to_2tuple(img_size)
        if isinstance(patch_size, int):
            patch_size = (1, patch_size, patch_size)

        # 3D patch embedding
        self.patch_embed = PatchEmbed(
            input_size=(num_frames,) + self.img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
        )
        self.out_channels = [embed_dim * self.patch_embed.grid_size[0]] * depth

        # Optional temporal and location embedding
        coords_encoding = coords_encoding or []
        self.temporal_encoding = "time" in coords_encoding
        self.location_encoding = "location" in coords_encoding
        if self.temporal_encoding:
            assert patch_size[0] == 1, (
                f"With temporal encoding, patch_size[0] must be 1, received {patch_size[0]}"
            )
            self.temporal_embed_enc = TemporalEncoder(embed_dim, coords_scale_learn)
        if self.location_encoding:
            self.location_embed_enc = LocationEncoder(embed_dim, coords_scale_learn)

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.register_buffer(
            "pos_embed", torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim)
        )

        # Transformer layers
        self.blocks = []
        for i in range(depth):
            self.blocks.append(
                Block(
                    embed_dim,
                    num_heads,
                    mlp_ratio,
                    qkv_bias=True,
                    norm_layer=norm_layer,
                    drop_path=drop_path,
                )
            )
        self.blocks = nn.ModuleList(self.blocks)

        self.norm = norm_layer(embed_dim)

        self.initialize_weights()

    def initialize_weights(self) -> None:
        """initialize_weights."""
        # initialize (and freeze) position embeddings by sin-cos embedding
        pos_embed = get_3d_sincos_pos_embed(
            self.pos_embed.shape[-1], self.patch_embed.grid_size, add_cls_token=True
        )
        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

        # initialize patch_embeddings like nn.Linear (instead of nn.Conv2d)
        w = self.patch_embed.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
        torch.nn.init.normal_(self.cls_token, std=0.02)
        self.apply(_init_weights)

    def random_masking(
        self,
        sequence: torch.Tensor,
        mask_ratio: float,
        noise: None | torch.Tensor = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Perform per-sample random masking by per-sample shuffling.

        Per-sample shuffling is done by argsort random
        noise.

        Args:
            sequence: (`torch.FloatTensor` of shape `(batch_size, sequence_length, dim)`)
            mask_ratio: (float): mask ratio to use.
            noise: (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*) which is
                mainly used for testing purposes to control randomness and maintain the reproducibility
        """
        batch_size, seq_length, dim = sequence.shape
        len_keep = int(seq_length * (1 - mask_ratio))

        if noise is None:
            noise = torch.rand(
                batch_size, seq_length, device=sequence.device
            )  # noise in [0, 1]

        # sort noise for each sample
        ids_shuffle = torch.argsort(noise, dim=1).to(
            sequence.device
        )  # ascend: small is keep, large is remove
        ids_restore = torch.argsort(ids_shuffle, dim=1).to(sequence.device)

        # keep the first subset
        ids_keep = ids_shuffle[:, :len_keep]
        sequence_unmasked = torch.gather(
            sequence, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, dim)
        )

        # generate the binary mask: 0 is keep, 1 is remove
        mask = torch.ones([batch_size, seq_length], device=sequence.device)
        mask[:, :len_keep] = 0
        # unshuffle to get the binary mask
        mask = torch.gather(mask, dim=1, index=ids_restore)

        return sequence_unmasked, mask, ids_restore

    def interpolate_pos_encoding(
        self, sample_shape: tuple[int, int, int] | list[int]
    ) -> torch.Tensor:
        """interpolate_pos_encoding."""
        pos_embed = _interpolate_pos_encoding(
            pos_embed=self.pos_embed,
            grid_size=self.patch_embed.grid_size,
            patch_size=self.patch_embed.patch_size,
            shape=sample_shape,
            embed_dim=self.embed_dim,
        )
        return pos_embed

    def forward(
        self,
        x: torch.Tensor,
        temporal_coords: None | torch.Tensor = None,
        location_coords: None | torch.Tensor = None,
        mask_ratio: float = 0.75,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward."""
        if len(x.shape) == 4 and self.patch_embed.input_size[0] == 1:
            # add time dim
            x = x.unsqueeze(2)
        sample_shape = x.shape[-3:]

        # embed patches
        x = self.patch_embed(x)

        pos_embed = self.interpolate_pos_encoding(sample_shape)
        # add pos embed w/o cls token
        x = x + pos_embed[:, 1:, :]

        if self.temporal_encoding and temporal_coords is not None:
            num_tokens_per_frame = x.shape[1] // self.num_frames
            temporal_encoding = self.temporal_embed_enc(
                temporal_coords, num_tokens_per_frame
            )
            x = x + temporal_encoding
        if self.location_encoding and location_coords is not None:
            location_encoding = self.location_embed_enc(location_coords)
            x = x + location_encoding

        # masking: length -> length * mask_ratio
        x, mask, ids_restore = self.random_masking(x, mask_ratio)

        # append cls token
        cls_token = self.cls_token + pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # apply Transformer blocks
        for block in self.blocks:
            x = block(x)
        x = self.norm(x)

        return x, mask, ids_restore

    def forward_features(
        self,
        x: torch.Tensor,
        temporal_coords: None | torch.Tensor = None,
        location_coords: None | torch.Tensor = None,
    ) -> list[torch.Tensor]:
        """forward_features."""
        if len(x.shape) == 4 and self.patch_embed.input_size[0] == 1:
            # add time dim
            x = x.unsqueeze(2)
        sample_shape = x.shape[-3:]

        # embed patches
        x = self.patch_embed(x)

        pos_embed = self.interpolate_pos_encoding(sample_shape)
        # add pos embed w/o cls token
        x = x + pos_embed[:, 1:, :]

        if self.temporal_encoding and temporal_coords is not None:
            num_tokens_per_frame = x.shape[1] // self.num_frames
            temporal_encoding = self.temporal_embed_enc(
                temporal_coords, num_tokens_per_frame
            )
            x = x + temporal_encoding
        if self.location_encoding and location_coords is not None:
            location_encoding = self.location_embed_enc(location_coords)
            x = x + location_encoding

        # append cls token
        cls_token = self.cls_token + pos_embed[:, :1, :]
        cls_tokens = cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls_tokens, x), dim=1)

        # apply Transformer blocks
        out = []
        for block in self.blocks:
            x = block(x)
            out.append(x.clone())

        x = self.norm(x)
        out[-1] = x
        return out

    def prepare_features_for_image_model(
        self, features: list[torch.Tensor], t: int
    ) -> list[torch.Tensor]:
        """prepare_features_for_image_model."""
        out = []
        for x in features:
            x_no_token = x[:, 1:, :]
            number_of_tokens = x_no_token.shape[1]
            tokens_per_timestep = number_of_tokens // t
            h = int(np.sqrt(tokens_per_timestep))
            encoded = rearrange(
                x_no_token,
                "batch (t h w) e -> batch t e h w",
                e=self.embed_dim,
                t=t,
                h=h,
            )
            # mean along the time dimension
            out.append(encoded.mean(dim=1))
        return out

initialize_weights ¶

initialize_weights() -> None

initialize_weights.

Source code in rslearn/models/prithvi.py

def initialize_weights(self) -> None:
    """initialize_weights."""
    # initialize (and freeze) position embeddings by sin-cos embedding
    pos_embed = get_3d_sincos_pos_embed(
        self.pos_embed.shape[-1], self.patch_embed.grid_size, add_cls_token=True
    )
    self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))

    # initialize patch_embeddings like nn.Linear (instead of nn.Conv2d)
    w = self.patch_embed.proj.weight.data
    torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

    # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
    torch.nn.init.normal_(self.cls_token, std=0.02)
    self.apply(_init_weights)

random_masking ¶

random_masking(sequence: Tensor, mask_ratio: float, noise: None | Tensor = None) -> tuple[Tensor, Tensor, Tensor]

Perform per-sample random masking by per-sample shuffling.

Per-sample shuffling is done by argsort random noise.

Parameters:

Name	Type	Description	Default
`sequence`	`Tensor`	(`torch.FloatTensor` of shape `(batch_size, sequence_length, dim)`)	required
`mask_ratio`	`float`	(float): mask ratio to use.	required
`noise`	`None \| Tensor`	(`torch.FloatTensor` of shape `(batch_size, sequence_length)`, optional) which is mainly used for testing purposes to control randomness and maintain the reproducibility	`None`

Source code in rslearn/models/prithvi.py

def random_masking(
    self,
    sequence: torch.Tensor,
    mask_ratio: float,
    noise: None | torch.Tensor = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Perform per-sample random masking by per-sample shuffling.

    Per-sample shuffling is done by argsort random
    noise.

    Args:
        sequence: (`torch.FloatTensor` of shape `(batch_size, sequence_length, dim)`)
        mask_ratio: (float): mask ratio to use.
        noise: (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*) which is
            mainly used for testing purposes to control randomness and maintain the reproducibility
    """
    batch_size, seq_length, dim = sequence.shape
    len_keep = int(seq_length * (1 - mask_ratio))

    if noise is None:
        noise = torch.rand(
            batch_size, seq_length, device=sequence.device
        )  # noise in [0, 1]

    # sort noise for each sample
    ids_shuffle = torch.argsort(noise, dim=1).to(
        sequence.device
    )  # ascend: small is keep, large is remove
    ids_restore = torch.argsort(ids_shuffle, dim=1).to(sequence.device)

    # keep the first subset
    ids_keep = ids_shuffle[:, :len_keep]
    sequence_unmasked = torch.gather(
        sequence, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, dim)
    )

    # generate the binary mask: 0 is keep, 1 is remove
    mask = torch.ones([batch_size, seq_length], device=sequence.device)
    mask[:, :len_keep] = 0
    # unshuffle to get the binary mask
    mask = torch.gather(mask, dim=1, index=ids_restore)

    return sequence_unmasked, mask, ids_restore

interpolate_pos_encoding ¶

interpolate_pos_encoding(sample_shape: tuple[int, int, int] | list[int]) -> Tensor

interpolate_pos_encoding.

Source code in rslearn/models/prithvi.py

def interpolate_pos_encoding(
    self, sample_shape: tuple[int, int, int] | list[int]
) -> torch.Tensor:
    """interpolate_pos_encoding."""
    pos_embed = _interpolate_pos_encoding(
        pos_embed=self.pos_embed,
        grid_size=self.patch_embed.grid_size,
        patch_size=self.patch_embed.patch_size,
        shape=sample_shape,
        embed_dim=self.embed_dim,
    )
    return pos_embed

forward ¶

forward(x: Tensor, temporal_coords: None | Tensor = None, location_coords: None | Tensor = None, mask_ratio: float = 0.75) -> tuple[Tensor, Tensor, Tensor]

Forward.

Source code in rslearn/models/prithvi.py

def forward(
    self,
    x: torch.Tensor,
    temporal_coords: None | torch.Tensor = None,
    location_coords: None | torch.Tensor = None,
    mask_ratio: float = 0.75,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward."""
    if len(x.shape) == 4 and self.patch_embed.input_size[0] == 1:
        # add time dim
        x = x.unsqueeze(2)
    sample_shape = x.shape[-3:]

    # embed patches
    x = self.patch_embed(x)

    pos_embed = self.interpolate_pos_encoding(sample_shape)
    # add pos embed w/o cls token
    x = x + pos_embed[:, 1:, :]

    if self.temporal_encoding and temporal_coords is not None:
        num_tokens_per_frame = x.shape[1] // self.num_frames
        temporal_encoding = self.temporal_embed_enc(
            temporal_coords, num_tokens_per_frame
        )
        x = x + temporal_encoding
    if self.location_encoding and location_coords is not None:
        location_encoding = self.location_embed_enc(location_coords)
        x = x + location_encoding

    # masking: length -> length * mask_ratio
    x, mask, ids_restore = self.random_masking(x, mask_ratio)

    # append cls token
    cls_token = self.cls_token + pos_embed[:, :1, :]
    cls_tokens = cls_token.expand(x.shape[0], -1, -1)
    x = torch.cat((cls_tokens, x), dim=1)

    # apply Transformer blocks
    for block in self.blocks:
        x = block(x)
    x = self.norm(x)

    return x, mask, ids_restore

forward_features ¶

forward_features(x: Tensor, temporal_coords: None | Tensor = None, location_coords: None | Tensor = None) -> list[Tensor]

forward_features.

Source code in rslearn/models/prithvi.py

def forward_features(
    self,
    x: torch.Tensor,
    temporal_coords: None | torch.Tensor = None,
    location_coords: None | torch.Tensor = None,
) -> list[torch.Tensor]:
    """forward_features."""
    if len(x.shape) == 4 and self.patch_embed.input_size[0] == 1:
        # add time dim
        x = x.unsqueeze(2)
    sample_shape = x.shape[-3:]

    # embed patches
    x = self.patch_embed(x)

    pos_embed = self.interpolate_pos_encoding(sample_shape)
    # add pos embed w/o cls token
    x = x + pos_embed[:, 1:, :]

    if self.temporal_encoding and temporal_coords is not None:
        num_tokens_per_frame = x.shape[1] // self.num_frames
        temporal_encoding = self.temporal_embed_enc(
            temporal_coords, num_tokens_per_frame
        )
        x = x + temporal_encoding
    if self.location_encoding and location_coords is not None:
        location_encoding = self.location_embed_enc(location_coords)
        x = x + location_encoding

    # append cls token
    cls_token = self.cls_token + pos_embed[:, :1, :]
    cls_tokens = cls_token.expand(x.shape[0], -1, -1)
    x = torch.cat((cls_tokens, x), dim=1)

    # apply Transformer blocks
    out = []
    for block in self.blocks:
        x = block(x)
        out.append(x.clone())

    x = self.norm(x)
    out[-1] = x
    return out

prepare_features_for_image_model ¶

prepare_features_for_image_model(features: list[Tensor], t: int) -> list[Tensor]

prepare_features_for_image_model.

Source code in rslearn/models/prithvi.py

def prepare_features_for_image_model(
    self, features: list[torch.Tensor], t: int
) -> list[torch.Tensor]:
    """prepare_features_for_image_model."""
    out = []
    for x in features:
        x_no_token = x[:, 1:, :]
        number_of_tokens = x_no_token.shape[1]
        tokens_per_timestep = number_of_tokens // t
        h = int(np.sqrt(tokens_per_timestep))
        encoded = rearrange(
            x_no_token,
            "batch (t h w) e -> batch t e h w",
            e=self.embed_dim,
            t=t,
            h=h,
        )
        # mean along the time dimension
        out.append(encoded.mean(dim=1))
    return out

MAEDecoder ¶

Bases: Module

Transformer Decoder used in the Prithvi MAE.

Source code in rslearn/models/prithvi.py

class MAEDecoder(nn.Module):
    """Transformer Decoder used in the Prithvi MAE."""

    def __init__(
        self,
        patch_size: int | tuple[int, int, int] = (1, 16, 16),
        grid_size: list[int] | tuple[int, int, int] = (3, 14, 14),
        in_chans: int = 3,
        encoder_embed_dim: int = 1024,
        decoder_embed_dim: int = 512,
        depth: int = 8,
        num_heads: int = 16,
        mlp_ratio: float = 4.0,
        norm_layer: nn.Module = nn.LayerNorm,
        coords_encoding: list[str] | None = None,
        coords_scale_learn: bool = False,
    ) -> None:
        """Init."""
        super().__init__()

        self.decoder_embed = nn.Linear(encoder_embed_dim, decoder_embed_dim, bias=True)
        self.decoder_embed_dim = decoder_embed_dim
        self.grid_size = grid_size
        if isinstance(patch_size, int):
            patch_size = (1, patch_size, patch_size)
        self.patch_size = patch_size
        self.num_frames = self.grid_size[0] * patch_size[0]
        num_patches = self.grid_size[0] * self.grid_size[1] * self.grid_size[2]

        # Optional temporal and location embedding
        coords_encoding = coords_encoding or []
        self.temporal_encoding = "time" in coords_encoding
        self.location_encoding = "location" in coords_encoding
        if self.temporal_encoding:
            self.temporal_embed_dec = TemporalEncoder(
                decoder_embed_dim, coords_scale_learn
            )
        if self.location_encoding:
            self.location_embed_dec = LocationEncoder(
                decoder_embed_dim, coords_scale_learn
            )

        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))

        self.register_buffer(
            "decoder_pos_embed", torch.zeros(1, num_patches + 1, decoder_embed_dim)
        )

        self.decoder_blocks = nn.ModuleList(
            [
                Block(
                    decoder_embed_dim,
                    num_heads,
                    mlp_ratio,
                    qkv_bias=True,
                    norm_layer=norm_layer,
                )
                for _ in range(depth)
            ]
        )

        self.decoder_norm = norm_layer(decoder_embed_dim)
        self.decoder_pred = nn.Linear(
            decoder_embed_dim,
            patch_size[0] * patch_size[1] * patch_size[2] * in_chans,
            bias=True,
        )

        self.initialize_weights()

    def initialize_weights(self) -> None:
        """initialize_weights."""
        # initialize (and freeze) position embeddings by sin-cos embedding
        decoder_pos_embed = get_3d_sincos_pos_embed(
            self.decoder_pos_embed.shape[-1], self.grid_size, add_cls_token=True
        )
        self.decoder_pos_embed.data.copy_(
            torch.from_numpy(decoder_pos_embed).float().unsqueeze(0)
        )

        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
        torch.nn.init.normal_(self.mask_token, std=0.02)
        self.apply(_init_weights)

    def interpolate_pos_encoding(
        self, sample_shape: tuple[int, int, int]
    ) -> torch.Tensor:
        """interpolate_pos_encoding."""
        pos_embed = _interpolate_pos_encoding(
            pos_embed=self.decoder_pos_embed,
            grid_size=self.grid_size,
            patch_size=self.patch_size,
            shape=sample_shape,
            embed_dim=self.decoder_embed_dim,
        )

        return pos_embed

    def forward(
        self,
        hidden_states: torch.Tensor,
        ids_restore: torch.Tensor,
        temporal_coords: None | torch.Tensor = None,
        location_coords: None | torch.Tensor = None,
        input_size: list[int] | None = None,
    ) -> torch.Tensor:
        """Forward."""
        # embed tokens
        x = self.decoder_embed(hidden_states)
        cls_token = x[:, :1, :]

        # append mask tokens to sequence
        mask_tokens = self.mask_token.repeat(
            x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1
        )
        x = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
        # unshuffle
        x = torch.gather(
            x,
            dim=1,
            index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]).to(x.device),
        )

        # add pos embed
        decoder_pos_embed = self.interpolate_pos_encoding(input_size[-3:])  # type: ignore
        cls_token = cls_token + decoder_pos_embed[:, :1, :]
        x = x + decoder_pos_embed[:, 1:, :]

        if self.temporal_encoding and temporal_coords is not None:
            num_tokens_per_frame = x.shape[1] // self.num_frames
            temporal_encoding = self.temporal_embed_dec(
                temporal_coords, num_tokens_per_frame
            )
            # Add temporal encoding w/o cls token
            x = x + temporal_encoding
        if self.location_encoding and location_coords is not None:
            location_encoding = self.location_embed_dec(location_coords)
            # Add location encoding w/o cls token
            x = x + location_encoding

        # append cls token
        x = torch.cat([cls_token, x], dim=1)

        # apply Transformer layers (blocks)
        for block in self.decoder_blocks:
            x = block(x)
        x = self.decoder_norm(x)

        # predictor projection
        pred = self.decoder_pred(x)

        # remove cls token
        pred = pred[:, 1:, :]

        return pred

initialize_weights ¶

initialize_weights() -> None

initialize_weights.

Source code in rslearn/models/prithvi.py

def initialize_weights(self) -> None:
    """initialize_weights."""
    # initialize (and freeze) position embeddings by sin-cos embedding
    decoder_pos_embed = get_3d_sincos_pos_embed(
        self.decoder_pos_embed.shape[-1], self.grid_size, add_cls_token=True
    )
    self.decoder_pos_embed.data.copy_(
        torch.from_numpy(decoder_pos_embed).float().unsqueeze(0)
    )

    # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
    torch.nn.init.normal_(self.mask_token, std=0.02)
    self.apply(_init_weights)

interpolate_pos_encoding ¶

interpolate_pos_encoding(sample_shape: tuple[int, int, int]) -> Tensor

interpolate_pos_encoding.

Source code in rslearn/models/prithvi.py

def interpolate_pos_encoding(
    self, sample_shape: tuple[int, int, int]
) -> torch.Tensor:
    """interpolate_pos_encoding."""
    pos_embed = _interpolate_pos_encoding(
        pos_embed=self.decoder_pos_embed,
        grid_size=self.grid_size,
        patch_size=self.patch_size,
        shape=sample_shape,
        embed_dim=self.decoder_embed_dim,
    )

    return pos_embed

forward ¶

forward(hidden_states: Tensor, ids_restore: Tensor, temporal_coords: None | Tensor = None, location_coords: None | Tensor = None, input_size: list[int] | None = None) -> Tensor

Forward.

Source code in rslearn/models/prithvi.py

def forward(
    self,
    hidden_states: torch.Tensor,
    ids_restore: torch.Tensor,
    temporal_coords: None | torch.Tensor = None,
    location_coords: None | torch.Tensor = None,
    input_size: list[int] | None = None,
) -> torch.Tensor:
    """Forward."""
    # embed tokens
    x = self.decoder_embed(hidden_states)
    cls_token = x[:, :1, :]

    # append mask tokens to sequence
    mask_tokens = self.mask_token.repeat(
        x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1
    )
    x = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
    # unshuffle
    x = torch.gather(
        x,
        dim=1,
        index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]).to(x.device),
    )

    # add pos embed
    decoder_pos_embed = self.interpolate_pos_encoding(input_size[-3:])  # type: ignore
    cls_token = cls_token + decoder_pos_embed[:, :1, :]
    x = x + decoder_pos_embed[:, 1:, :]

    if self.temporal_encoding and temporal_coords is not None:
        num_tokens_per_frame = x.shape[1] // self.num_frames
        temporal_encoding = self.temporal_embed_dec(
            temporal_coords, num_tokens_per_frame
        )
        # Add temporal encoding w/o cls token
        x = x + temporal_encoding
    if self.location_encoding and location_coords is not None:
        location_encoding = self.location_embed_dec(location_coords)
        # Add location encoding w/o cls token
        x = x + location_encoding

    # append cls token
    x = torch.cat([cls_token, x], dim=1)

    # apply Transformer layers (blocks)
    for block in self.decoder_blocks:
        x = block(x)
    x = self.decoder_norm(x)

    # predictor projection
    pred = self.decoder_pred(x)

    # remove cls token
    pred = pred[:, 1:, :]

    return pred

PrithviMAE ¶

Bases: Module

Prithvi Masked Autoencoder.

Source code in rslearn/models/prithvi.py

class PrithviMAE(nn.Module):
    """Prithvi Masked Autoencoder."""

    def __init__(
        self,
        img_size: int | tuple[int, int] = 224,
        patch_size: int | tuple[int, int, int] = (1, 16, 16),
        num_frames: int = 4,
        in_chans: int = 6,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        decoder_embed_dim: int = 512,
        decoder_depth: int = 8,
        decoder_num_heads: int = 16,
        mlp_ratio: float = 4.0,
        norm_layer: nn.Module = nn.LayerNorm,
        norm_pix_loss: bool = False,
        coords_encoding: list[str] | None = None,
        coords_scale_learn: bool = False,
        drop_path: float = 0.0,
        mask_ratio: float = 0.75,
        **kwargs: Any,
    ):
        """Init."""
        super().__init__()

        self.encoder = PrithviViT(
            img_size=img_size,
            num_frames=num_frames,
            patch_size=patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
            depth=depth,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            norm_layer=norm_layer,
            coords_encoding=coords_encoding,
            coords_scale_learn=coords_scale_learn,
            drop_path=drop_path,
        )

        self.decoder = MAEDecoder(
            patch_size=patch_size,
            grid_size=self.encoder.patch_embed.grid_size,
            in_chans=in_chans,
            encoder_embed_dim=embed_dim,
            decoder_embed_dim=decoder_embed_dim,
            depth=decoder_depth,
            num_heads=decoder_num_heads,
            mlp_ratio=mlp_ratio,
            norm_layer=norm_layer,
            coords_encoding=coords_encoding,
            coords_scale_learn=coords_scale_learn,
        )

        self.mask_ratio = mask_ratio
        self.norm_pix_loss = norm_pix_loss
        self.out_channels = self.encoder.out_channels

    def patchify(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """Patchify.

        Args:
            pixel_values: (torch.FloatTensor of shape `(batch_size, num_channels, time, height, width)`):
                Pixel values.

        Returns:
            torch.FloatTensor of shape
                `(batch_size, num_patches, patch_size[0]*patch_size[1]*patch_size[2] * num_channels)`:
                Patchified pixel values.
        """
        patch_size_t, patch_size_h, patch_size_w = self.encoder.patch_embed.patch_size
        num_channels = self.encoder.in_chans

        # patchify
        patchified_pixel_values = rearrange(
            pixel_values,
            "b c (t s) (h p) (w q) -> b (t h w) (s p q c)",
            c=num_channels,
            s=patch_size_t,
            p=patch_size_h,
            q=patch_size_w,
        )

        return patchified_pixel_values

    def unpatchify(
        self,
        patchified_pixel_values: torch.Tensor,
        image_size: tuple[int, int] | None = None,
    ) -> torch.Tensor:
        """Unpatchify.

        Args:
            patchified_pixel_values: (`torch.FloatTensor` of shape
                `(batch_size, num_patches, patch_size[0]*patch_size[1]*patch_size[2] * num_channels))`:
                Patchified pixel values.
            image_size: (`tuple[int, int]`, *optional*):
                Original image size.

        Returns:
            `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`:
                Pixel values.
        """
        patch_size_t, patch_size_h, patch_size_w = self.encoder.patch_embed.patch_size
        image_size = (
            to_2tuple(image_size) if image_size is not None else self.encoder.img_size
        )
        original_height, original_width = image_size
        num_patches_h = original_height // patch_size_h
        num_patches_w = original_width // patch_size_w
        num_channels = self.encoder.in_chans

        pixel_values = rearrange(
            patchified_pixel_values,
            "b (t h w) (s p q c) -> b c (t s) (h p) (w q)",
            c=num_channels,
            h=num_patches_h,
            w=num_patches_w,
            s=patch_size_t,
            p=patch_size_h,
            q=patch_size_w,
        )
        return pixel_values

    def forward_loss(
        self, pixel_values: torch.Tensor, pred: torch.Tensor, mask: torch.Tensor
    ) -> torch.Tensor:
        """forward_loss.

        Args:
            pixel_values: (`torch.FloatTensor` of shape `(batch_size, num_channels, time, height, width)`):
                Pixel values.
            pred: (`torch.FloatTensor` of shape
                `(batch_size, num_patches, patch_size[0]*patch_size[1]*patch_size[2] * num_channels)`:
                Predicted pixel values.
            mask: (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
                Tensor indicating which patches are masked (1) and which are not (0).

        Returns:
            `torch.FloatTensor`: Pixel reconstruction loss.
        """
        target = self.patchify(pixel_values)
        if self.norm_pix_loss:
            mean = target.mean(dim=-1, keepdim=True)
            var = target.var(dim=-1, keepdim=True)
            target = (target - mean) / (var + 1.0e-6) ** 0.5

        loss = (pred - target) ** 2
        loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
        loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
        return loss

    def forward(
        self,
        pixel_values: torch.Tensor,
        temporal_coords: None | torch.Tensor = None,
        location_coords: None | torch.Tensor = None,
        mask_ratio: float | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward."""
        if len(pixel_values.shape) == 4 and self.encoder.patch_embed.input_size[0] == 1:
            # add time dim
            pixel_values = pixel_values.unsqueeze(2)

        mask_ratio = mask_ratio or self.mask_ratio
        latent, mask, ids_restore = self.encoder(
            pixel_values, temporal_coords, location_coords, mask_ratio
        )
        pred = self.decoder(
            latent,
            ids_restore,
            temporal_coords,
            location_coords,
            input_size=pixel_values.shape,
        )
        loss = self.forward_loss(pixel_values, pred, mask)
        return loss, pred, mask

    def forward_features(
        self,
        x: torch.Tensor,
        temporal_coords: None | torch.Tensor = None,
        location_coords: None | torch.Tensor = None,
    ) -> list[torch.Tensor]:
        """forward_features."""
        return self.encoder.forward_features(x, temporal_coords, location_coords)

patchify ¶

patchify(pixel_values: Tensor) -> Tensor

Patchify.

Parameters:

Name	Type	Description	Default
`pixel_values`	`Tensor`	(torch.FloatTensor of shape `(batch_size, num_channels, time, height, width)`): Pixel values.	required

Returns:

Type	Description
`Tensor`	torch.FloatTensor of shape `(batch_size, num_patches, patch_size[0]patch_size[1]patch_size[2] * num_channels)`: Patchified pixel values.

Source code in rslearn/models/prithvi.py

def patchify(self, pixel_values: torch.Tensor) -> torch.Tensor:
    """Patchify.

    Args:
        pixel_values: (torch.FloatTensor of shape `(batch_size, num_channels, time, height, width)`):
            Pixel values.

    Returns:
        torch.FloatTensor of shape
            `(batch_size, num_patches, patch_size[0]*patch_size[1]*patch_size[2] * num_channels)`:
            Patchified pixel values.
    """
    patch_size_t, patch_size_h, patch_size_w = self.encoder.patch_embed.patch_size
    num_channels = self.encoder.in_chans

    # patchify
    patchified_pixel_values = rearrange(
        pixel_values,
        "b c (t s) (h p) (w q) -> b (t h w) (s p q c)",
        c=num_channels,
        s=patch_size_t,
        p=patch_size_h,
        q=patch_size_w,
    )

    return patchified_pixel_values

unpatchify ¶

unpatchify(patchified_pixel_values: Tensor, image_size: tuple[int, int] | None = None) -> Tensor

Unpatchify.

Parameters:

Name	Type	Description	Default
`patchified_pixel_values`	`Tensor`	(`torch.FloatTensor` of shape `(batch_size, num_patches, patch_size[0]patch_size[1]patch_size[2] * num_channels))`: Patchified pixel values.	required
`image_size`	`tuple[int, int] \| None`	(`tuple[int, int]`, optional): Original image size.	`None`

Returns:

Type	Description
`Tensor`	`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`: Pixel values.

Source code in rslearn/models/prithvi.py

def unpatchify(
    self,
    patchified_pixel_values: torch.Tensor,
    image_size: tuple[int, int] | None = None,
) -> torch.Tensor:
    """Unpatchify.

    Args:
        patchified_pixel_values: (`torch.FloatTensor` of shape
            `(batch_size, num_patches, patch_size[0]*patch_size[1]*patch_size[2] * num_channels))`:
            Patchified pixel values.
        image_size: (`tuple[int, int]`, *optional*):
            Original image size.

    Returns:
        `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`:
            Pixel values.
    """
    patch_size_t, patch_size_h, patch_size_w = self.encoder.patch_embed.patch_size
    image_size = (
        to_2tuple(image_size) if image_size is not None else self.encoder.img_size
    )
    original_height, original_width = image_size
    num_patches_h = original_height // patch_size_h
    num_patches_w = original_width // patch_size_w
    num_channels = self.encoder.in_chans

    pixel_values = rearrange(
        patchified_pixel_values,
        "b (t h w) (s p q c) -> b c (t s) (h p) (w q)",
        c=num_channels,
        h=num_patches_h,
        w=num_patches_w,
        s=patch_size_t,
        p=patch_size_h,
        q=patch_size_w,
    )
    return pixel_values

forward_loss ¶

forward_loss(pixel_values: Tensor, pred: Tensor, mask: Tensor) -> Tensor

forward_loss.

Parameters:

Name	Type	Description	Default
`pixel_values`	`Tensor`	(`torch.FloatTensor` of shape `(batch_size, num_channels, time, height, width)`): Pixel values.	required
`pred`	`Tensor`	(`torch.FloatTensor` of shape `(batch_size, num_patches, patch_size[0]patch_size[1]patch_size[2] * num_channels)`: Predicted pixel values.	required
`mask`	`Tensor`	(`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Tensor indicating which patches are masked (1) and which are not (0).	required

Returns:

Type	Description
`Tensor`	`torch.FloatTensor`: Pixel reconstruction loss.

Source code in rslearn/models/prithvi.py

def forward_loss(
    self, pixel_values: torch.Tensor, pred: torch.Tensor, mask: torch.Tensor
) -> torch.Tensor:
    """forward_loss.

    Args:
        pixel_values: (`torch.FloatTensor` of shape `(batch_size, num_channels, time, height, width)`):
            Pixel values.
        pred: (`torch.FloatTensor` of shape
            `(batch_size, num_patches, patch_size[0]*patch_size[1]*patch_size[2] * num_channels)`:
            Predicted pixel values.
        mask: (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
            Tensor indicating which patches are masked (1) and which are not (0).

    Returns:
        `torch.FloatTensor`: Pixel reconstruction loss.
    """
    target = self.patchify(pixel_values)
    if self.norm_pix_loss:
        mean = target.mean(dim=-1, keepdim=True)
        var = target.var(dim=-1, keepdim=True)
        target = (target - mean) / (var + 1.0e-6) ** 0.5

    loss = (pred - target) ** 2
    loss = loss.mean(dim=-1)  # [N, L], mean loss per patch
    loss = (loss * mask).sum() / mask.sum()  # mean loss on removed patches
    return loss

forward ¶

forward(pixel_values: Tensor, temporal_coords: None | Tensor = None, location_coords: None | Tensor = None, mask_ratio: float | None = None) -> tuple[Tensor, Tensor, Tensor]

Forward.

Source code in rslearn/models/prithvi.py

def forward(
    self,
    pixel_values: torch.Tensor,
    temporal_coords: None | torch.Tensor = None,
    location_coords: None | torch.Tensor = None,
    mask_ratio: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Forward."""
    if len(pixel_values.shape) == 4 and self.encoder.patch_embed.input_size[0] == 1:
        # add time dim
        pixel_values = pixel_values.unsqueeze(2)

    mask_ratio = mask_ratio or self.mask_ratio
    latent, mask, ids_restore = self.encoder(
        pixel_values, temporal_coords, location_coords, mask_ratio
    )
    pred = self.decoder(
        latent,
        ids_restore,
        temporal_coords,
        location_coords,
        input_size=pixel_values.shape,
    )
    loss = self.forward_loss(pixel_values, pred, mask)
    return loss, pred, mask

forward_features ¶

forward_features(x: Tensor, temporal_coords: None | Tensor = None, location_coords: None | Tensor = None) -> list[Tensor]

forward_features.

Source code in rslearn/models/prithvi.py

def forward_features(
    self,
    x: torch.Tensor,
    temporal_coords: None | torch.Tensor = None,
    location_coords: None | torch.Tensor = None,
) -> list[torch.Tensor]:
    """forward_features."""
    return self.encoder.forward_features(x, temporal_coords, location_coords)

get_config ¶

get_config(cache_dir: Path, hf_hub_id: str, hf_hub_revision: str) -> dict[str, Any]

Get the JSON config dict.

Parameters:

Name	Type	Description	Default
`cache_dir`	`Path`	the directory to cache the config.json file, which will be downloaded from HF Hub.	required
`hf_hub_id`	`str`	the HF Hub ID from which to download the config.	required
`hf_hub_revision`	`str`	The revision (commit) to download the config from.	required

Source code in rslearn/models/prithvi.py

def get_config(cache_dir: Path, hf_hub_id: str, hf_hub_revision: str) -> dict[str, Any]:
    """Get the JSON config dict.

    Args:
        cache_dir: the directory to cache the config.json file, which will be
            downloaded from HF Hub.
        hf_hub_id: the HF Hub ID from which to download the config.
        hf_hub_revision: The revision (commit) to download the config from.
    """
    cache_fname = cache_dir / HF_HUB_CONFIG_FNAME
    if not cache_fname.exists():
        _ = hf_hub_download(
            local_dir=cache_dir,
            repo_id=hf_hub_id,
            filename=HF_HUB_CONFIG_FNAME,
            revision=hf_hub_revision,
        )  # nosec
    with cache_fname.open() as f:
        return json.load(f)["pretrained_cfg"]

get_3d_sincos_pos_embed ¶

get_3d_sincos_pos_embed(embed_dim: int, grid_size: tuple[int, int, int] | list[int], add_cls_token: bool = False) -> Tensor

Create 3D sin/cos positional embeddings.

Parameters:

Name	Type	Description	Default
`embed_dim`	`int`	Embedding dimension.	required
`grid_size`	`tuple[int, int, int] \| list[int]`	The grid depth, height and width.	required
`add_cls_token`	`bool, optional, defaults to False`	Whether or not to add a classification (CLS) token.	`False`

Returns:

Type	Description
`Tensor`	(`torch.FloatTensor` of shape (grid_size[0]grid_size[1]grid_size[2], embed_dim) or
`(1 + grid_size[0] * grid_size[1] * grid_size[2], embed_dim)`	the position embeddings (with or without cls token)

Source code in rslearn/models/prithvi.py

def get_3d_sincos_pos_embed(
    embed_dim: int,
    grid_size: tuple[int, int, int] | list[int],
    add_cls_token: bool = False,
) -> torch.Tensor:
    """Create 3D sin/cos positional embeddings.

    Args:
        embed_dim (int):
            Embedding dimension.
        grid_size (tuple[int, int, int] | list[int]):
            The grid depth, height and width.
        add_cls_token (bool, *optional*, defaults to False):
            Whether or not to add a classification (CLS) token.

    Returns:
        (`torch.FloatTensor` of shape (grid_size[0]*grid_size[1]*grid_size[2], embed_dim) or
        (1+grid_size[0]*grid_size[1]*grid_size[2], embed_dim): the position embeddings (with or without cls token)
    """
    assert embed_dim % 16 == 0

    t_size, h_size, w_size = grid_size

    w_embed_dim = embed_dim // 16 * 6
    h_embed_dim = embed_dim // 16 * 6
    t_embed_dim = embed_dim // 16 * 4

    w_pos_embed = get_1d_sincos_pos_embed_from_grid(w_embed_dim, np.arange(w_size))
    h_pos_embed = get_1d_sincos_pos_embed_from_grid(h_embed_dim, np.arange(h_size))
    t_pos_embed = get_1d_sincos_pos_embed_from_grid(t_embed_dim, np.arange(t_size))

    w_pos_embed = np.tile(w_pos_embed, (t_size * h_size, 1))
    h_pos_embed = np.tile(np.repeat(h_pos_embed, w_size, axis=0), (t_size, 1))
    t_pos_embed = np.repeat(t_pos_embed, h_size * w_size, axis=0)

    pos_embed = np.concatenate((w_pos_embed, h_pos_embed, t_pos_embed), axis=1)

    if add_cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed

get_1d_sincos_pos_embed_from_grid ¶

get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: Tensor) -> Tensor

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D).

Source code in rslearn/models/prithvi.py

def get_1d_sincos_pos_embed_from_grid(
    embed_dim: int, pos: torch.Tensor
) -> torch.Tensor:
    """embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)."""
    if embed_dim % 2 != 0:
        raise ValueError("embed_dim must be even")

    omega = np.arange(embed_dim // 2, dtype=float)
    omega /= embed_dim / 2.0
    omega = 1.0 / 10000**omega  # (D/2,)

    pos = pos.reshape(-1)  # (M,)
    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product

    emb_sin = np.sin(out)  # (M, D/2)
    emb_cos = np.cos(out)  # (M, D/2)

    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
    return emb

rslearn.models.prithvi¶

prithvi ¶

PrithviV2Models ¶

PrithviV2 ¶

forward ¶

get_backbone_channels ¶

PrithviNormalize ¶

forward ¶

PatchEmbed ¶

forward ¶

TemporalEncoder ¶

forward ¶

LocationEncoder ¶

forward ¶

PrithviViT ¶

initialize_weights ¶

random_masking ¶

interpolate_pos_encoding ¶

forward ¶

forward_features ¶

prepare_features_for_image_model ¶

MAEDecoder ¶

initialize_weights ¶

interpolate_pos_encoding ¶

forward ¶

PrithviMAE ¶

patchify ¶

unpatchify ¶

forward_loss ¶

forward ¶

forward_features ¶

get_config ¶

get_3d_sincos_pos_embed ¶

get_1d_sincos_pos_embed_from_grid ¶

`rslearn.models.prithvi`¶