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

rslearn/models/ssl4eo_s12.py

@@ -1,20 +1,22 @@
 """SSL4EO-S12 models."""
 
-from typing import Any
-
 import torch
 import torchvision
 
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps
+
 
-class Ssl4eoS12(torch.nn.Module):
+class Ssl4eoS12(FeatureExtractor):
     """The SSL4EO-S12 family of pretrained models."""
 
     def __init__(
         self,
-        backbone_ckpt_path: str,
+        backbone_ckpt_path: str | None,
         arch: str = "resnet50",
         output_layers: list[int] = [0, 1, 2, 3],
-    ):
+    ) -> None:
         """Instantiate a new Swin instance.
 
         Args:
@@ -37,21 +39,24 @@ class Ssl4eoS12(torch.nn.Module):
         else:
             raise ValueError(f"unknown SSL4EO-S12 architecture {arch}")
 
-        state_dict = torch.load(backbone_ckpt_path)
-        state_dict = state_dict["teacher"]
-        prefix = "module.backbone."
-        state_dict = {
-            k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)
-        }
-        missing_keys, unexpected_keys = self.model.load_state_dict(
-            state_dict, strict=False
-        )
-        if missing_keys or unexpected_keys:
-            print(
-                f"warning: got missing_keys={missing_keys}, unexpected_keys={unexpected_keys} when loading SSL4EO-S12 state dict"
+        if backbone_ckpt_path is not None:
+            state_dict = torch.load(backbone_ckpt_path, weights_only=True)
+            state_dict = state_dict["teacher"]
+            prefix = "module.backbone."
+            state_dict = {
+                k[len(prefix) :]: v
+                for k, v in state_dict.items()
+                if k.startswith(prefix)
+            }
+            missing_keys, unexpected_keys = self.model.load_state_dict(
+                state_dict, strict=False
             )
+            if missing_keys or unexpected_keys:
+                print(
+                    f"warning: got missing_keys={missing_keys}, unexpected_keys={unexpected_keys} when loading SSL4EO-S12 state dict"
+                )
 
-    def get_backbone_channels(self):
+    def get_backbone_channels(self) -> list[tuple[int, int]]:
         """Returns the output channels of this model when used as a backbone.
 
         The output channels is a list of (downsample_factor, depth) that corresponds
@@ -65,28 +70,33 @@ class Ssl4eoS12(torch.nn.Module):
         """
         if self.arch == "resnet50":
             all_out_channels = [
-                [4, 256],
-                [8, 512],
-                [16, 1024],
-                [32, 2048],
+                (4, 256),
+                (8, 512),
+                (16, 1024),
+                (32, 2048),
             ]
         return [all_out_channels[idx] for idx in self.output_layers]
 
     def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+        self,
+        context: ModelContext,
+    ) -> FeatureMaps:
         """Compute outputs from the backbone.
 
         If output_layers is set, then the outputs are multi-scale feature maps;
         otherwise, the model is being used for classification so the outputs are class
         probabilities and the loss.
 
-        Inputs:
-            inputs: input dicts that must include "image" key containing the image to
-                process.
-            targets: target dicts that are ignored unless
+        Args:
+            context: the model context. Input dicts must include "image" key containing
+                the images to process.
+
+        Returns:
+            feature maps computed by the pre-trained model.
         """
-        x = torch.stack([inp["image"] for inp in inputs], dim=0)
+        x = torch.stack(
+            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
+        )
         x = self.model.conv1(x)
         x = self.model.bn1(x)
         x = self.model.relu(x)
@@ -97,4 +107,4 @@ class Ssl4eoS12(torch.nn.Module):
         layer3 = self.model.layer3(layer2)
         layer4 = self.model.layer4(layer3)
         all_features = [layer1, layer2, layer3, layer4]
-        return [all_features[idx] for idx in self.output_layers]
+        return FeatureMaps([all_features[idx] for idx in self.output_layers])

rslearn/models/swin.py

@@ -1,7 +1,5 @@
 """Swin Transformer."""
 
-from typing import Any
-
 import torch
 import torchvision
 from torchvision.models.swin_transformer import (
@@ -13,8 +11,12 @@ from torchvision.models.swin_transformer import (
     Swin_V2_T_Weights,
 )
 
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureExtractor, FeatureMaps, FeatureVector
 
-class Swin(torch.nn.Module):
+
+class Swin(FeatureExtractor):
     """A Swin Transformer model.
 
     It can either be used stand-alone for classification, or as a feature extractor in
@@ -28,15 +30,18 @@ class Swin(torch.nn.Module):
         input_channels: int = 3,
         output_layers: list[int] | None = None,
         num_outputs: int = 1000,
-    ):
+    ) -> None:
         """Instantiate a new Swin instance.
 
         Args:
             arch: the architecture, e.g. "swin_v2_b" (default) or "swin_t"
             pretrained: set True to use ImageNet pre-trained weights
-            input_channels: number of input channels (default 3)
+            input_channels: number of input channels (default 3). If not 3, the first
+                layer is updated and will be randomly initialized even if pretrained is
+                set.
             output_layers: list of layers to output, default use as classification
-                model. For feature extraction, [1, 3, 5, 7] is recommended.
+                model (output FeatureVector). For feature extraction, [1, 3, 5, 7] is
+                recommended.
             num_outputs: number of output logits, defaults to 1000 which matches the
                 pretrained models.
         """
@@ -89,7 +94,7 @@ class Swin(torch.nn.Module):
         if num_outputs != self.model.head.out_features:
             self.model.head = torch.nn.Linear(self.model.head.in_features, num_outputs)
 
-    def get_backbone_channels(self):
+    def get_backbone_channels(self) -> list[tuple[int, int]]:
         """Returns the output channels of this model when used as a backbone.
 
         The output channels is a list of (downsample_factor, depth) that corresponds
@@ -105,43 +110,50 @@ class Swin(torch.nn.Module):
 
         if self.arch in ["swin_b", "swin_v2_b"]:
             all_out_channels = [
-                [4, 128],
-                [4, 128],
-                [8, 256],
-                [8, 256],
-                [16, 512],
-                [16, 512],
-                [32, 1024],
-                [32, 1024],
+                (4, 128),
+                (4, 128),
+                (4, 128),
+                (8, 256),
+                (8, 256),
+                (16, 512),
+                (16, 512),
+                (32, 1024),
+                (32, 1024),
             ]
         elif self.arch in ["swin_s", "swin_v2_s", "swin_t", "swin_v2_t"]:
             all_out_channels = [
-                [4, 96],
-                [4, 96],
-                [8, 192],
-                [8, 192],
-                [16, 384],
-                [16, 384],
-                [32, 768],
-                [32, 768],
+                (4, 96),
+                (4, 96),
+                (8, 192),
+                (8, 192),
+                (16, 384),
+                (16, 384),
+                (32, 768),
+                (32, 768),
             ]
         return [all_out_channels[idx] for idx in self.output_layers]
 
     def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+        self,
+        context: ModelContext,
+    ) -> FeatureVector | FeatureMaps:
         """Compute outputs from the backbone.
 
         If output_layers is set, then the outputs are multi-scale feature maps;
         otherwise, the model is being used for classification so the outputs are class
         probabilities and the loss.
 
-        Inputs:
-            inputs: input dicts that must include "image" key containing the image to
-                process.
-            targets: target dicts that are ignored unless
+        Args:
+            context: the model context. Input dicts must include "image" key containing
+                the image to process.
+
+        Returns:
+            a FeatureVector if the configured output_layers is None, or a FeatureMaps
+                otherwise containing one feature map per configured output layer.
         """
-        images = torch.stack([inp["image"] for inp in inputs], dim=0)
+        images = torch.stack(
+            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
+        )
 
         if self.output_layers:
             layer_features = []
@@ -149,7 +161,7 @@ class Swin(torch.nn.Module):
             for layer in self.model.features:
                 x = layer(x)
                 layer_features.append(x.permute(0, 3, 1, 2))
-            return [layer_features[idx] for idx in self.output_layers]
+            return FeatureMaps([layer_features[idx] for idx in self.output_layers])
 
         else:
-            return self.model(images)
+            return FeatureVector(self.model(images))

rslearn/models/task_embedding.py

@@ -0,0 +1,250 @@
+"""Task embedding modules."""
+
+import math
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class PositionalEncoding(nn.Module):
+    """Simple sinusoidal positional encoding for the task embedding. From torch docs."""
+
+    def __init__(self, d_model: int, dropout: float = 0.0, max_len: int = 1024):
+        """Initialize the positional encoding module.
+
+        Args:
+            d_model: The dimension of the model.
+            dropout: The dropout rate.
+            max_len: The maximum length of the sequence.
+        """
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
+        )
+        pe = torch.zeros(max_len, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply positional encoding to the input tensor.
+
+        Args:
+            x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``
+        """
+        x = x + self.pe[: x.size(0)]
+        return self.dropout(x)
+
+
+class BaseTaskEmbedding(torch.nn.Module):
+    """Base class for task embedding modules."""
+
+    def __init__(self, encoder_embedding_size: int) -> None:
+        """Initialize the base task embedding module.
+
+        Args:
+            encoder_embedding_size: The size of the encoder embedding.
+        """
+        super().__init__()
+        self.encoder_embedding_size = encoder_embedding_size
+
+    def register_tasks(self, task_names: list[str]) -> None:
+        """Register the tasks.
+
+        This must happen post-init so that we can dynamically determine
+        the tasks to use, so it doesn't have to be specified in the config.
+
+        Args:
+            task_names: The names of the tasks.
+        """
+        raise NotImplementedError
+
+    def compute_embeds(
+        self,
+        features: list[torch.tensor],
+        inputs: list[dict[str, Any]],
+    ) -> torch.Tensor:
+        """Compute the task-specific embeddings.
+
+        Args:
+            features: The encoder features.
+            inputs: The inputs to the model.
+
+        Returns:
+            The task-specific embeddings.
+        """
+        raise NotImplementedError
+
+
+class TaskChannelEmbedding(BaseTaskEmbedding):
+    """Registers task-specific 'tokens', i.e. embeddings.
+
+    Each embedding is learned per-channel and copied over the full spatial dimensions.
+    Optionally, add a spatial sinusoidal positional embedding to the task embedding.
+    """
+
+    def __init__(
+        self,
+        encoder_embedding_size: int,
+        default_idx: int = 0,
+        add_spatial_embed: bool = False,
+    ) -> None:
+        """Initialize the task channel embedding module.
+
+        Args:
+            encoder_embedding_size: The size of the encoder embedding.
+            default_idx: The index of the default task, useful if loading a merged model.
+            add_spatial_embed: if true, add a spatial sinusoidal positional embedding to the task embedding
+        """
+        super().__init__(encoder_embedding_size)
+        self.default_idx = default_idx
+        self.add_spatial_embed = add_spatial_embed
+        if add_spatial_embed:
+            self.pos_embed = PositionalEncoding(encoder_embedding_size)
+
+    def register_tasks(self, task_names: list[str]) -> None:
+        """Register the tasks.
+
+        This must happen post-init so that we can dynamically determine
+        the tasks to use, so it doesn't have to be specified in the config.
+
+        Args:
+            task_names: The names of the tasks.
+        """
+        self.embed = torch.nn.Embedding(len(task_names), self.encoder_embedding_size)
+        self.target_to_embed_idx = {name: i for i, name in enumerate(task_names)}
+
+    def compute_embeds(
+        self,
+        features: list[torch.tensor],
+        inputs: list[dict[str, Any]],
+    ) -> torch.Tensor:
+        """Compute the task-specific embeddings.
+
+        Args:
+            inputs: The inputs to the model.
+            features: computed encoder features
+            device: The device to compute the embeddings on.
+
+        Returns:
+            The task-specific embeddings, shape (B, T, C), T = HW
+            The embeddings are repeated over the spatial dimensions, and optionally
+            a sinusoidal positional embedding is added.
+        """
+        try:
+            idx = [self.target_to_embed_idx[inp["dataset_source"]] for inp in inputs]
+        except KeyError:
+            idx = [self.default_idx] * len(inputs)
+        embeds = self.embed(torch.tensor(idx).to(features[0].device))
+        seq_len = features[0].shape[-1] * features[0].shape[-2]  # T = HW
+        embeds = embeds.unsqueeze(0).repeat(seq_len, 1, 1)  # T x B x C
+        if self.add_spatial_embed:
+            embeds = self.pos_embed(embeds)
+        embeds = torch.einsum("tbc->btc", embeds)  # B x T x C
+        return embeds
+
+    def forward(
+        self,
+        features: list[torch.tensor],
+        inputs: list[dict[str, Any]],
+        embeds: torch.Tensor | None = None,
+    ) -> list[torch.tensor]:
+        """Compute and apply task-specific embeddings to encoder features.
+
+        Optionally, add a spatial sinusoidal positional embedding to the task embedding.
+        Otherwise, the task embedding is repeated over the spatial dimensions.
+
+        Args:
+            features: The encoder features, a 1-list of B x C x H x W features.
+            inputs: The inputs to the model.
+            embeds: Already-computed task embeddings, if provided, skip the computation.
+
+        Returns:
+            The encoder features with the task-specific embeddings added.
+        """
+        height, width = features[0].shape[-2:]
+        assert all(f.shape[-2:] == (height, width) for f in features), (
+            "features must have the same spatial dimensions"
+        )
+        if embeds is None:
+            embeds = self.compute_embeds(features, inputs)  # B x HW x C
+        embeds = embeds.unflatten(dim=1, sizes=(height, width))  # B x H x W x C
+        for i in range(len(features)):
+            features[i] += torch.einsum("bhwc->bchw", embeds)  # B x C x H x W
+        return features
+
+
+class TaskMHAEmbedding(TaskChannelEmbedding):
+    """Multi-headed cross-attention over the spatial dimensions.
+
+    The task embedding is the query and the features are the key and value.
+    We copy the task embedding over the spatial dimensions, and optionally
+    add a sinusoidal positional embedding before the MHA layer.
+    """
+
+    def __init__(
+        self,
+        encoder_embedding_size: int,
+        num_heads: int,
+        default_idx: int = 0,
+        add_spatial_embed: bool = True,
+    ) -> None:
+        """Initialize the task MHA embedding module.
+
+        Args:
+            encoder_embedding_size: The size of the encoder embedding.
+            num_heads: The number of attention heads.
+            default_idx: The index of the default task, useful if loading a merged model.
+            add_spatial_embed: if true, add a spatial sinusoidal positional embedding to the task embedding
+        """
+        super().__init__(encoder_embedding_size, default_idx, add_spatial_embed)
+        self.mha = torch.nn.MultiheadAttention(
+            encoder_embedding_size, num_heads, batch_first=True
+        )
+
+    def register_tasks(self, task_names: list[str]) -> None:
+        """Register the tasks.
+
+        This must happen post-init so that we can dynamically determine
+        the tasks to use, so it doesn't have to be specified in the config.
+
+        Args:
+            task_names: The names of the tasks.
+        """
+        super().register_tasks(task_names)
+
+    def forward(
+        self,
+        features: list[torch.tensor],
+        inputs: list[dict[str, Any]],
+        embeds: torch.Tensor | None = None,
+    ) -> list[torch.tensor]:
+        """Compute and apply task-specific embeddings to encoder features.
+
+        Also apply the MHA layer across the spatial dimension, with the task embedding
+        as the query and the features as the key and value.
+
+        Args:
+            features: The encoder features, a 1-list of B x C x H x W features.
+            inputs: The inputs to the model.
+            embeds: Already-computed task embeddings, if provided, skip the computation.
+
+        Returns:
+            The encoder features with the task-specific embeddings added.
+        """
+        assert len(features) == 1, "TaskMHAEmbedding only supports one feature"
+        x = torch.flatten(features[0], start_dim=2)  # B x C x T, T = HW
+        if embeds is None:
+            embeds = self.compute_embeds(features, inputs)  # B x T x C
+        out = self.mha(
+            embeds,  # B x T x C
+            torch.einsum("bct->btc", x),
+            torch.einsum("bct->btc", x),
+        )[0]  # B x T x C
+        out = torch.einsum("btc->bct", out)
+        out = out.view(*features[0].shape)  # B x C x H x W
+        return [out]