rslearn 0.0.11__py3-none-any.whl → 0.0.13__py3-none-any.whl

rslearn/models/copernicusfm_src/flexivit/patch_embed.py

@@ -1,260 +0,0 @@
-# type: ignore
-from collections.abc import Sequence
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-from functorch import vmap
-from torch import Tensor
-
-from .utils import to_2tuple
-
-
-def pi_resize_patch_embed(
-    patch_embed: Tensor,
-    new_patch_size: tuple[int, int],
-    interpolation: str = "bicubic",
-    antialias: bool = True,
-):
-    """Resample patch embedding weights to a target resolution via pseudo-inverse
-    resizing.
-
-    Based on:
-        https://github.com/google-research/big_vision/blob/b00544b81f8694488d5f36295aeb7972f3755ffe/big_vision/models/proj/flexi/vit.py
-        https://arxiv.org/abs/2212.08013
-
-    Args:
-        patch_embed: Patch embedding parameters of size [d, c, h, w]
-        new_patch_size: Target [height, width] of embedding
-        interpolation: Resize interpolation type
-        antialias: Whether to apply antialiasing resizing
-    Returns:
-        Resized pos_embed of size [d, c h', w']
-    """
-    assert len(patch_embed.shape) == 4, "Patch embed kernel should be a 4D tensor"
-    assert len(new_patch_size) == 2, "New patch size should only be (height, width)"
-
-    old_patch_size = tuple(patch_embed.shape[2:])
-
-    # Return original kernel if no resize is necessary
-    if old_patch_size == new_patch_size:
-        return patch_embed
-
-    def resize(x: Tensor, shape: tuple[int, int]):
-        x_resized = F.interpolate(
-            x[None, None, ...],
-            shape,
-            mode=interpolation,
-            antialias=antialias,
-        )
-        return x_resized[0, 0, ...]
-
-    def calculate_pinv(old_shape: tuple[int, int], new_shape: tuple[int, int]):
-        mat = []
-        for i in range(np.prod(old_shape)):
-            basis_vec = torch.zeros(old_shape)
-            basis_vec[np.unravel_index(i, old_shape)] = 1.0
-            mat.append(resize(basis_vec, new_shape).reshape(-1))
-        resize_matrix = torch.stack(mat)
-        return torch.linalg.pinv(resize_matrix)
-
-    # Calculate pseudo-inverse of resize matrix
-    resize_matrix_pinv = calculate_pinv(old_patch_size, new_patch_size)
-    resize_matrix_pinv = resize_matrix_pinv.to(patch_embed.device)
-
-    def resample_patch_embed(patch_embed: Tensor):
-        h, w = new_patch_size
-        resampled_kernel = resize_matrix_pinv @ patch_embed.reshape(-1)
-        return rearrange(resampled_kernel, "(h w) -> h w", h=h, w=w)
-
-    v_resample_patch_embed = vmap(vmap(resample_patch_embed, 0, 0), 1, 1)
-
-    return v_resample_patch_embed(patch_embed)
-
-
-def interpolate_resize_patch_embed(
-    patch_embed: Tensor,
-    new_patch_size: tuple[int, int],
-    interpolation: str = "bicubic",
-    antialias: bool = True,
-):
-    """Resample patch embedding weights to a target resolution via interpolation
-
-    Args:
-        patch_embed: Patch embedding parameters of size [d, c, h, w]
-        new_patch_size: Target [height, width] of embedding
-        interpolation: Resize interpolation type
-        antialias: Whether to apply antialiasing resizing
-    Returns:
-        Resized pos_embed of size [d, c h', w']
-    """
-    assert len(patch_embed.shape) == 4, "Patch embed kernel should be a 4D tensor"
-    assert len(new_patch_size) == 2, "New patch size should only be (height, width)"
-
-    patch_embed = F.interpolate(
-        patch_embed, new_patch_size, mode=interpolation, antialias=antialias
-    )
-
-    return patch_embed
-
-
-class FlexiPatchEmbed(nn.Module):
-    def __init__(
-        self,
-        img_size: int | tuple[int, int] = 240,
-        patch_size: int | tuple[int, int] = 32,
-        grid_size: int | tuple[int, int] = 7,
-        in_chans: int = 3,
-        embed_dim: int = 768,
-        norm_layer: nn.Module | None = None,
-        flatten: bool = True,
-        bias: bool = True,
-        patch_size_seq: Sequence[int] = (8, 10, 12, 15, 16, 20, 24, 30, 40, 48),
-        patch_size_probs: Sequence[float] | None = None,
-        interpolation: str = "bicubic",
-        antialias: bool = True,
-    ) -> None:
-        """2D image to patch embedding w/ flexible patch sizes
-        Extended from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/patch_embed.py#L24
-
-        Args:
-            img_size: Input image size
-            patch_size: Base patch size. i.e the size of the parameter buffer
-            grid_size: Size of pos_embed buffer
-            in_chans: Number of input image channels
-            embed_dim: Network embedding dimension size
-            norm_layer: Optional normalization layer
-            flatten: Whether to flatten the spatial dimensions of the output
-            bias: Whether to use bias in convolution
-            patch_size_seq: List of patch sizes to randomly sample from
-            patch_size_probs: Optional list of probabilities to sample corresponding
-                patch_size_seq elements. If None, then uniform distribution is used
-            interpolation: Resize interpolation type
-            antialias: Whether to apply antialiasing resizing
-        """
-        super().__init__()
-
-        self.img_size = to_2tuple(img_size)
-        self.patch_size = to_2tuple(patch_size)
-        self.grid_size = to_2tuple(grid_size)
-        self.num_patches = self.grid_size[0] * self.grid_size[1]
-
-        self.flatten = flatten
-        self.proj = nn.Conv2d(
-            in_chans,
-            embed_dim,
-            kernel_size=self.patch_size,
-            stride=self.patch_size,
-            bias=bias,
-        )
-        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
-
-        # Flexi specific attributes
-        self.interpolation = interpolation
-        self.antialias = antialias
-
-        self.patch_size_seq = patch_size_seq
-
-        if self.patch_size_seq:
-            if not patch_size_probs:
-                n = len(self.patch_size_seq)
-                self.patch_size_probs = [1.0 / n] * n
-            else:
-                self.patch_size_probs = [
-                    p / sum(patch_size_probs) for p in patch_size_probs
-                ]
-        else:
-            self.patch_size_probs = []
-
-        # Pre-calculate pinvs
-        self.pinvs = self._cache_pinvs()
-
-    def _cache_pinvs(self) -> dict:
-        """Pre-calculate all pinv matrices"""
-        pinvs = {}
-        for ps in self.patch_size_seq:
-            ps = to_2tuple(ps)
-            pinvs[ps] = self._calculate_pinv(self.patch_size, ps)
-        return pinvs
-
-    def _resize(self, x: Tensor, shape: tuple[int, int]) -> Tensor:
-        x_resized = F.interpolate(
-            x[None, None, ...],
-            shape,
-            mode=self.interpolation,
-            antialias=self.antialias,
-        )
-        return x_resized[0, 0, ...]
-
-    def _calculate_pinv(
-        self, old_shape: tuple[int, int], new_shape: tuple[int, int]
-    ) -> Tensor:
-        mat = []
-        for i in range(np.prod(old_shape)):
-            basis_vec = torch.zeros(old_shape)
-            basis_vec[np.unravel_index(i, old_shape)] = 1.0
-            mat.append(self._resize(basis_vec, new_shape).reshape(-1))
-        resize_matrix = torch.stack(mat)
-        return torch.linalg.pinv(resize_matrix)
-
-    def resize_patch_embed(self, patch_embed: Tensor, new_patch_size: tuple[int, int]):
-        """Resize patch_embed to target resolution via pseudo-inverse resizing"""
-        # Return original kernel if no resize is necessary
-        if self.patch_size == new_patch_size:
-            return patch_embed
-
-        # Calculate pseudo-inverse of resize matrix
-        if new_patch_size not in self.pinvs:
-            self.pinvs[new_patch_size] = self._calculate_pinv(
-                self.patch_size, new_patch_size
-            )
-        pinv = self.pinvs[new_patch_size]
-        pinv = pinv.to(patch_embed.device)
-
-        def resample_patch_embed(patch_embed: Tensor):
-            h, w = new_patch_size
-            resampled_kernel = pinv @ patch_embed.reshape(-1)
-            return rearrange(resampled_kernel, "(h w) -> h w", h=h, w=w)
-
-        v_resample_patch_embed = vmap(vmap(resample_patch_embed, 0, 0), 1, 1)
-
-        return v_resample_patch_embed(patch_embed)
-
-    def forward(
-        self,
-        x: Tensor,
-        patch_size: int | tuple[int, int] | None = None,
-        return_patch_size: bool = False,
-    ) -> Tensor | tuple[Tensor, tuple[int, int]]:
-        if not patch_size and not self.training:
-            # During evaluation use base patch size if not specified
-            patch_size = self.patch_size
-        elif not patch_size:
-            # During training choose uniformly at random if not specified
-            assert self.patch_size_seq, (
-                "No patch size specified during forward and no patch_size_seq given to FlexiPatchEmbed"
-            )
-            patch_size = np.random.choice(self.patch_size_seq, p=self.patch_size_probs)
-
-        patch_size = to_2tuple(patch_size)
-
-        # Resize conv weights
-        if patch_size == self.patch_size:
-            weight = self.proj.weight
-        else:
-            weight = self.resize_patch_embed(self.proj.weight, patch_size)
-
-        # Apply conv with resized weights
-        x = F.conv2d(x, weight, bias=self.proj.bias, stride=patch_size)
-
-        if self.flatten:
-            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
-
-        x = self.norm(x)
-
-        if return_patch_size:
-            return x, patch_size
-
-        return x

rslearn/models/copernicusfm_src/flexivit/utils.py

@@ -1,69 +0,0 @@
-# type: ignore
-import collections.abc
-import math
-from itertools import repeat
-from typing import Any
-
-import torch
-import torch.nn.functional as F
-
-
-def to_2tuple(x: Any) -> tuple:
-    if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
-        return tuple(x)
-    return tuple(repeat(x, 2))
-
-
-def resize_abs_pos_embed(
-    pos_embed: torch.Tensor,
-    new_size: tuple[int, int],
-    old_size: int | tuple[int, int] | None = None,
-    num_prefix_tokens: int = 1,
-    interpolation: str = "bicubic",
-    antialias: bool = True,
-) -> torch.Tensor:
-    """Resize absolute position embeddings to a target resolution via interpolation
-
-    Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/pos_embed.py
-
-    Args:
-        pos_embed: Position embeddings tensor of size [b, n, d]
-        new_size: Target [height, width] of embedding
-        old_size: Original [height, width] of embedding
-        num_prefix_tokens: Number of non-spatial prefix tokens (eg. cls)
-        interpolation: Resize interpolation type
-        antialias: Whether to apply antialiasing resizing
-    Returns:
-        Resized pos_embed of size [b, n', d]
-    """
-    new_size = to_2tuple(new_size)
-    new_ntok = new_size[0] * new_size[1]
-
-    if not old_size:
-        old_size = int(math.sqrt(pos_embed.shape[1] - num_prefix_tokens))  # type:ignore
-    old_size = to_2tuple(old_size)
-
-    # Return if no resize necessary
-    if new_size == old_size:
-        return pos_embed
-
-    if num_prefix_tokens:
-        posemb_prefix, pos_embed = (
-            pos_embed[:, :num_prefix_tokens],
-            pos_embed[:, num_prefix_tokens:],
-        )
-    else:
-        posemb_prefix, pos_embed = None, pos_embed
-
-    # Interpolate position embedding
-    pos_embed = pos_embed.reshape(1, old_size[0], old_size[1], -1).permute(0, 3, 1, 2)
-    pos_embed = F.interpolate(
-        pos_embed, size=new_size, mode=interpolation, antialias=antialias
-    )
-    pos_embed = pos_embed.permute(0, 2, 3, 1).reshape(1, new_ntok, -1)
-
-    # Add back extra prefix tokens
-    if posemb_prefix is not None:
-        pos_embed = torch.cat([posemb_prefix, pos_embed], dim=1)
-
-    return pos_embed

rslearn/models/copernicusfm_src/model_vit.py

@@ -1,348 +0,0 @@
-# mypy: ignore-errors
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# --------------------------------------------------------
-# References:
-# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
-# DeiT: https://github.com/facebookresearch/deit
-# --------------------------------------------------------
-
-import math
-from functools import partial
-
-import torch
-import torch.nn as nn
-from timm.models.vision_transformer import Block
-
-from .aurora.fourier import FourierExpansion
-from .dynamic_hypernetwork import Dynamic_MLP_OFA_spectral, Dynamic_MLP_OFA_variable
-from .flexivit.utils import resize_abs_pos_embed
-
-
-class CopernicusFMViT(nn.Module):
-    """CopernicusFM: VisionTransformer backbone"""
-
-    def __init__(
-        self,
-        img_size=224,
-        patch_size=16,
-        drop_rate=0.0,
-        embed_dim=1024,
-        depth=24,
-        num_heads=16,
-        wv_planes=128,
-        num_classes=0,
-        global_pool=True,
-        mlp_ratio=4.0,
-        norm_layer=nn.LayerNorm,
-        loc_option="lonlat",
-        return_intermediate=False,
-        intermediate_indices=None,
-    ):
-        super().__init__()
-
-        self.wv_planes = wv_planes
-        self.global_pool = global_pool
-        if self.global_pool:
-            norm_layer = norm_layer
-            embed_dim = embed_dim
-            self.fc_norm = norm_layer(embed_dim)
-        else:
-            self.norm = norm_layer(embed_dim)
-
-        self.patch_embed_spectral = Dynamic_MLP_OFA_spectral(
-            wv_planes=128, inter_dim=128, kernel_size=16, embed_dim=embed_dim
-        )
-        self.patch_embed_variable = Dynamic_MLP_OFA_variable(
-            wv_planes=128, inter_dim=128, kernel_size=16, embed_dim=embed_dim
-        )
-
-        self.num_patches = (img_size // patch_size) ** 2
-
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        # ---------------------------------------------------------------------------
-
-        self.pos_embed = nn.Parameter(
-            torch.zeros(1, self.num_patches + 1, embed_dim), requires_grad=False
-        )  # fixed sin-cos embedding
-
-        self.loc_option = loc_option
-        if loc_option == "cartesian":
-            self.coord_expansion = FourierExpansion(1e-7, 2)
-        elif loc_option == "lonlat":
-            self.coord_expansion = FourierExpansion(0.0001, 720)
-
-        self.scale_expansion = FourierExpansion(0.001, 5.1e8)  # 1m2 to 5.1e8 km2
-        self.time_expansion = FourierExpansion(
-            1, 365.25, assert_range=False
-        )  # 1 to 365.25 days, enable more than 1 year
-        self.coord_fc = nn.Linear(embed_dim, embed_dim)
-        self.scale_fc = nn.Linear(embed_dim, embed_dim)
-        self.time_fc = nn.Linear(embed_dim, embed_dim)
-        # if meta info is not available, set to a learned parameter
-        self.coord_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.scale_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        self.time_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-
-        self.blocks = nn.ModuleList(
-            [
-                Block(
-                    embed_dim,
-                    num_heads,
-                    mlp_ratio,
-                    qkv_bias=True,
-                    norm_layer=norm_layer,
-                )
-                for i in range(depth)
-            ]
-        )
-
-        self.head_drop = nn.Dropout(drop_rate)
-        self.head = (
-            nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
-        )
-
-        self.return_intermediate = return_intermediate
-        self.intermediate_indices = intermediate_indices
-
-    def get_coord_pos_embed(self, lons, lats, embed_dim):
-        if self.loc_option == "cartesian":
-            # convert to spherical coordinates
-            spherical_x = (
-                torch.cos(lons * math.pi / 180) * torch.cos(lats * math.pi / 180)
-                + 1
-                + 1e-7
-            )
-            spherical_y = (
-                torch.sin(lons * math.pi / 180) * torch.cos(lats * math.pi / 180)
-                + 1
-                + 1e-7
-            )
-            spherical_z = torch.sin(lats * math.pi / 180) + 1 + 1e-7
-            coord_embed_spherical_x = self.coord_expansion(spherical_x, embed_dim // 3)
-            coord_embed_spherical_y = self.coord_expansion(spherical_y, embed_dim // 3)
-            coord_embed_spherical_z = self.coord_expansion(spherical_z, embed_dim // 3)
-            coord_embed = torch.cat(
-                [
-                    coord_embed_spherical_x,
-                    coord_embed_spherical_y,
-                    coord_embed_spherical_z,
-                ],
-                dim=-1,
-            )  # [B,D]
-        elif self.loc_option == "lonlat":
-            coord_embed_lon = self.coord_expansion(lons + 180, embed_dim // 2)
-            coord_embed_lat = self.coord_expansion(lats + 90, embed_dim // 2)
-            coord_embed = torch.cat([coord_embed_lon, coord_embed_lat], dim=-1)
-
-        if coord_embed.shape[-1] < embed_dim:
-            # pad zeros
-            coord_embed = torch.cat(
-                (
-                    coord_embed,
-                    torch.zeros(
-                        coord_embed.shape[0],
-                        embed_dim - coord_embed.shape[-1],
-                        device=coord_embed.device,
-                    ),
-                ),
-                dim=-1,
-            )
-
-        return coord_embed.unsqueeze(1)  # [B,1,D]
-
-    def get_area_pos_embed(self, areas, embed_dim):
-        scale_embed = self.scale_expansion(areas, embed_dim)  # B, D
-        return scale_embed.unsqueeze(1)  # [B,1,D]
-
-    def get_time_pos_embed(self, times, embed_dim):
-        time_embed = self.time_expansion(times, embed_dim)  # B, D
-        return time_embed.unsqueeze(1)  # [B,1,D]
-
-    def forward_features(
-        self,
-        x,
-        meta_info,
-        wave_list,
-        bandwidth,
-        language_embed,
-        input_mode,
-        kernel_size=None,
-    ):
-        # embed patches
-        if input_mode == "spectral":
-            wavelist = torch.tensor(wave_list, device=x.device).float()
-            bandwidths = torch.tensor(bandwidth, device=x.device).float()
-            self.waves = wavelist
-            x, _ = self.patch_embed_spectral(x, self.waves, bandwidths, kernel_size)
-        elif input_mode == "variable":
-            x, _ = self.patch_embed_variable(x, language_embed, kernel_size)
-
-        # resize pos embed
-        num_patches = x.size(1)
-        num_patches_sqrt = int(math.sqrt(num_patches))
-        num_patches_sqrt_origin = int(math.sqrt(self.num_patches))
-        pos_embed = resize_abs_pos_embed(
-            self.pos_embed,
-            num_patches_sqrt,
-            (num_patches_sqrt_origin, num_patches_sqrt_origin),
-            num_prefix_tokens=1,
-        )
-
-        # coord, scale and time pos embed
-        lons, lats, times, areas = (
-            meta_info[:, 0],
-            meta_info[:, 1],
-            meta_info[:, 2],
-            meta_info[:, 3],
-        )
-        embed_dim = pos_embed.shape[-1]
-        if torch.isnan(lons).any() or torch.isnan(lats).any():
-            coord_embed = self.coord_token
-        else:
-            coord_embed = self.get_coord_pos_embed(lons, lats, embed_dim)
-        coord_embed = self.coord_fc(coord_embed)
-        if torch.isnan(areas).any():
-            area_embed = self.scale_token
-        else:
-            area_embed = self.get_area_pos_embed(areas, embed_dim)
-        area_embed = self.scale_fc(area_embed)
-        if torch.isnan(times).any():
-            time_embed = self.time_token
-        else:
-            time_embed = self.get_time_pos_embed(times, embed_dim)
-        time_embed = self.time_fc(time_embed)
-        pos_embed = pos_embed + coord_embed + area_embed + time_embed
-
-        # add pos embed w/o cls token
-        x = x + pos_embed[:, 1:, :]
-
-        # 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)
-
-        intermediate_features = []
-        hw = num_patches_sqrt
-        hw_shape = (hw, hw)
-
-        # apply Transformer blocks
-        for i, block in enumerate(self.blocks):
-            x = block(x)
-            if self.return_intermediate and (i in self.intermediate_indices):
-                out = x[:, 1:]
-                B, _, C = out.shape
-                out = (
-                    out.reshape(B, hw_shape[0], hw_shape[1], C)
-                    .permute(0, 3, 1, 2)
-                    .contiguous()
-                )
-                intermediate_features.append(out)
-
-        # if self.global_pool:
-        #     x = x[:, 1:, :].mean(dim=1)  # global pool without cls token
-        #     outcome = self.fc_norm(x)
-        # else:
-        #     x = self.norm(x)
-        #     outcome = x[:, 0]
-
-        # if self.return_intermediate:
-        #     return outcome, intermediate_features
-
-        # for segmentation tasks, ignore the norm
-        # https://github.com/zhu-xlab/Copernicus-FM/blob/main/Copernicus-Bench/src/foundation_models/CopernicusFM/models_dwv_seg.py
-        # for classification, we will apply the fc_norm in the wrapper
-        return x[:, 1:, :]
-
-    def forward_head(self, x, pre_logits=False):
-        x = self.head_drop(x)
-        return x if pre_logits else self.head(x)
-
-    def forward(
-        self,
-        x,
-        meta_info,
-        wave_list,
-        bandwidth,
-        language_embed,
-        input_mode,
-        kernel_size=None,
-    ):
-        if self.return_intermediate:
-            x, intermediate_features = self.forward_features(
-                x,
-                meta_info,
-                wave_list,
-                bandwidth,
-                language_embed,
-                input_mode,
-                kernel_size,
-            )
-            return x, intermediate_features
-        else:
-            fx = self.forward_features(
-                x,
-                meta_info,
-                wave_list,
-                bandwidth,
-                language_embed,
-                input_mode,
-                kernel_size,
-            )
-            x = self.forward_head(fx)
-            return x, fx
-
-
-def vit_small_patch16(**kwargs):
-    model = CopernicusFMViT(
-        patch_size=16,
-        embed_dim=384,
-        depth=12,
-        num_heads=6,
-        mlp_ratio=4,
-        norm_layer=partial(nn.LayerNorm, eps=1e-6),
-        **kwargs,
-    )
-    return model
-
-
-def vit_base_patch16(**kwargs):
-    model = CopernicusFMViT(
-        patch_size=16,
-        embed_dim=768,
-        depth=12,
-        num_heads=12,
-        mlp_ratio=4,
-        norm_layer=partial(nn.LayerNorm, eps=1e-6),
-        **kwargs,
-    )
-    return model
-
-
-def vit_large_patch16(**kwargs):
-    model = CopernicusFMViT(
-        patch_size=16,
-        embed_dim=1024,
-        depth=24,
-        num_heads=16,
-        mlp_ratio=4,
-        norm_layer=partial(nn.LayerNorm, eps=1e-6),
-        **kwargs,
-    )
-    return model
-
-
-def vit_huge_patch14(**kwargs):
-    model = CopernicusFMViT(
-        patch_size=14,
-        embed_dim=1280,
-        depth=32,
-        num_heads=16,
-        mlp_ratio=4,
-        norm_layer=partial(nn.LayerNorm, eps=1e-6),
-        **kwargs,
-    )
-    return model