Fix cnn and add SatCLIP

configs/model/predictive_cnn_s2.yaml

@@ -3,7 +3,6 @@ _target_: src.models.predictive_model.PredictiveModel
 geo_encoder:
   _target_: src.models.components.geo_encoders.cnn_encoder.CNNEncoder
   resnet_version: 18
-  freezing_strategy: none
   geo_data_name: s2
 
 prediction_head:

pyproject.toml

@@ -36,6 +36,8 @@ dependencies = [
     "geotessera>=0.8.0",
     "ipykernel>=7.2.0",
     "xgboost>=3.2.0",
+    "albumentations",
+    "satclip @ git+https://github.com/gabrieletijunaityte/satclip",
 ]
 
 [project.optional-dependencies]

src/models/components/geo_encoders/cnn_encoder.py

@@ -3,9 +3,16 @@
 import torch
 import torchvision.models as models
 from torch import nn
+from torchgeo.models import resnet50, ResNet50_Weights, ResNet18_Weights, resnet18
 
 from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder
+from src.utils.errors import IllegalArgumentCombination
 
+RN_DIM = {
+    18 : 512,
+    34: 512,
+    50: 2048
+}
 
 class CNNEncoder(BaseGeoEncoder):
     """Convolutional neural network EO encoder. Adapted from PECL.
@@ -23,27 +30,29 @@ def __init__(
         backbone="resnet",
         pretrained_cnn="imagenet",
         resnet_version=18,
-        freezing_strategy="all",
         geo_data_name="s2",
         input_n_bands: int | None = None,
-        output_dim=512,
     ) -> None:
         super().__init__()
 
+        # Backbone configurations
         self.backbone = backbone
-        self.pretrained_cnn = pretrained_cnn
-        self.resnet_version = resnet_version
-        self.freezing_strategy = freezing_strategy
+        if self.backbone == "resnet":
+            assert resnet_version in [18, 34, 50], f"Unsupported resnet version: {resnet_version}"
+            self.resnet_version = resnet_version
+
+            assert pretrained_cnn in ["imagenet", "IMAGENET1K_V1", 'SSL4EO_RGB_MOCO', None], f"Unsupported pretrained_cnn: {pretrained_cnn}"
+            self.pretrained_cnn = pretrained_cnn
+
+            self.output_dim = RN_DIM[resnet_version]
 
+        #  Input modality configurations
         self.allowed_geo_data_names = ["s2", "aef", "tessera"]
         assert geo_data_name in self.allowed_geo_data_names
         self.geo_data_name = geo_data_name
 
         self.set_n_input_bands(input_n_bands)
-        assert (
-            self.input_n_bands >= 3 and type(self.input_n_bands) is int
-        ), f"input_n_bands must be int >=3, got {self.input_n_bands}"
-        self.output_dim = output_dim
+        assert (self.input_n_bands >= 3 and type(self.input_n_bands) is int), f"input_n_bands must be int >=3, got {self.input_n_bands}"
 
     def set_n_input_bands(self, n_bands: int | None = None) -> None:
         """Sets number of input bands based on geo_data_name if n_bands is None.
@@ -65,76 +74,67 @@ def set_n_input_bands(self, n_bands: int | None = None) -> None:
             self.input_n_bands = n_bands
         return None
 
-    def get_backbone(self):
+    @override
+    def _setup(self) -> List[str]:
         """Gets backbone model given configuration stored in self.
-
         :return: backbone model
         """
+        trainable_modules = []
         if self.backbone == "resnet":
-            assert self.resnet_version in [
-                18,
-                34,
-                50,
-            ], f"Unsupported resnet version: {self.resnet_version}"
-            assert self.pretrained_cnn in [
-                "imagenet",
-                "IMAGENET1K_V1",
-                None,
-            ], f"Unsupported pretrained_cnn: {self.pretrained_cnn}"
-            if self.pretrained_cnn == "imagenet":
-                self.pretrained_cnn = "IMAGENET1K_V1"
-            if self.resnet_version == 18:
-                model = models.resnet18(weights=self.pretrained_cnn)
-            elif self.resnet_version == 34:
-                model = models.resnet34(weights=self.pretrained_cnn)
-            elif self.resnet_version == 50:
-                model = models.resnet50(weights=self.pretrained_cnn)
+            # Weights
+            # SSL4EO
+            if self.pretrained_cnn == "SSL4EO_RGB_MOCO":
+                if self.resnet_version == 18:
+                    self.geo_encoder = resnet18(weights=ResNet18_Weights.SENTINEL2_RGB_MOCO)
+                elif self.resnet_version == 34:
+                    raise IllegalArgumentCombination('SSL4EO_RGB_MOCO weights are not available for RN-34')
+                else:
+                    self.geo_encoder = resnet50(weights=ResNet50_Weights.SENTINEL2_RGB_MOCO)
+            # Imagenet
             else:
-                raise ValueError(f"Unsupported resnet version: {self.resnet_version}")
+                if self.pretrained_cnn == "imagenet":
+                    self.pretrained_cnn = "IMAGENET1K_V1"
+                elif self.pretrained_cnn == "imagenet_v2":
+                    self.pretrained_cnn = "IMAGENET1K_V2"
+
+                if self.resnet_version == 18:
+                    self.geo_encoder = models.resnet18(weights=self.pretrained_cnn)
+                elif self.resnet_version == 34:
+                    self.geo_encoder = models.resnet34(weights=self.pretrained_cnn)
+                else:
+                    self.geo_encoder = models.resnet50(weights=self.pretrained_cnn)
 
             # Modify the first conv layer to accept input_n_bands channels
-            if self.pretrained_cnn is not None and self.input_n_bands != 3:
-                weight = model.conv1.weight.clone()
             if self.input_n_bands != 3:
-                model.conv1 = torch.nn.Conv2d(
+
+                # Copy pre-trained weights
+                if self.pretrained_cnn is not None:
+                    weight = self.geo_encoder.conv1.weight.clone()
+
+                # Replace 1st conv layer
+                self.geo_encoder.conv1 = torch.nn.Conv2d(
                     self.input_n_bands, 64, kernel_size=7, stride=2, padding=3, bias=False
                 )
-                if self.pretrained_cnn is not None:  # copy pre-trained RGB bands
-                    for i in range(self.input_n_bands):
-                        model.conv1.weight.data[:, i, :, :] = weight[
-                            :, i % 3, :, :
-                        ]  # ensure this is not frozen
-            model.fc = nn.Linear(model.fc.in_features, self.output_dim)
-
-            assert self.freezing_strategy in [
-                "all",
-                "none",
-            ], f"Unsupported freezing_strategy: {self.freezing_strategy}"
-            layers_resnet = list(model.children())
-            n_layers = len(layers_resnet)
-            for i_c, child in enumerate(layers_resnet):
-                if i_c == 0:  # train first layer if not 3 bands (or no freezing)
-                    train_if = self.freezing_strategy == "none" or self.input_n_bands != 3
-                    for param in child.parameters():
-                        param.requires_grad = train_if
-                elif i_c == n_layers - 1:  # always train last layer
-                    for param in child.parameters():
-                        param.requires_grad = True
-                else:  # train other layers if no freezing
-                    train_if = self.freezing_strategy == "none"
-                    for param in child.parameters():
-                        param.requires_grad = train_if
-
-            return model
+
+                # Copy pre-trained RGB bands
+                if self.pretrained_cnn is not None:
+                    with torch.no_grad():
+                        for i in range(self.input_n_bands):
+                            self.geo_encoder.conv1.weight[:, i, :, :] = weight[:, i % 3, :, :]
+
+                # Ensure replaced layer is not frozen
+                trainable_modules.append('geo_encoder.conv1')
+
+            # I think for features fc often is replaced with identity?
+            self.geo_encoder.fc = nn.Identity()
+
+            # self.geo_encoder.fc = nn.Linear(self.geo_encoder.fc.in_features, self.output_dim)
+            # trainable_modules.append('geo_encoder.fc')
+
+            return trainable_modules
         else:
             raise ValueError(f"Unsupported backbone: {self.backbone}")
 
-    @override
-    def _setup(self) -> List[str]:
-        # TODO: could you make sure new layers are returned here to be added to trainable parts?
-        self.geo_encoder = self.get_backbone()
-        return []
-
     @override
     def forward(
         self,
@@ -145,18 +145,15 @@ def forward(
         :param batch: input batch
         :return: extracted features
         """
-        eo_data = batch.get("eo", {})
-
+        eo_data = batch.get("eo", KeyError(f"Batch must contain batch['eo']"))
+        eo_data = eo_data.get(self.geo_data_name,  KeyError(f"Batch must contain batch['eo']['{self.geo_data_name}']"))
         dtype = self.dtype
+
         if eo_data.dtype != dtype:
             eo_data = eo_data.to(dtype)
-        feats = self.geo_encoder(eo_data[self.geo_data_name])
-        # n_nans = torch.sum(torch.isnan(feats)).item()
-        # assert (
-        #     n_nans == 0
-        # ), f"CNNEncoder output contains {n_nans}/{feats.numel()} NaNs PRIOR to normalization with data min {eo_data[self.geo_data_name].min()} and max {eo_data[self.geo_data_name].max()}."
+        feats = self.geo_encoder(eo_data)
 
         if self.extra_projector:
             feats = self.extra_projector(feats)
 
-        return feats.to(dtype)
+        return feats

src/models/components/geo_encoders/satclip.py

@@ -0,0 +1,66 @@
+from typing import Dict, List, override
+
+import torch
+from huggingface_hub import hf_hub_download
+from satclip.load import get_satclip
+
+from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder
+
+
+class SatClipCoordinateEncoder(BaseGeoEncoder):
+    def __init__(
+        self,
+        geo_data_name="coords",
+        hf_cache_dir: str = "../.cache",
+        accelerator: torch.device = torch.device("cpu"),
+    ) -> None:
+        """SatClip coordinate encoder :param geo_data_name: type of geo data used for this encoder
+        (supports only coordinates) :param hf_cache_dir: hugging face cache directory to store data
+        :param accelerator: where to load model (as it is float64, mps is not supported)"""
+        super().__init__()
+
+        self.allowed_geo_data_names = ["coords"]
+        assert (
+            geo_data_name in self.allowed_geo_data_names
+        ), f"geo_data_name must be one of {self.allowed_geo_data_names}, got {geo_data_name}"
+        self.geo_data_name = geo_data_name
+
+        self.cache_dir = hf_cache_dir
+        assert accelerator != torch.device("mps"), f"accelerator {accelerator} is not supported"
+        self.accelerator = accelerator
+
+    @override
+    def _setup(self) -> List[str]:
+        """Setup satclip encoder from hugging face hub and set output dimension."""
+        self.geo_encoder = get_satclip(
+            hf_hub_download(
+                "microsoft/SatCLIP-ViT16-L40", "satclip-vit16-l40.ckpt", cache_dir=self.cache_dir
+            ),
+            device=self.accelerator,
+        )
+
+        self.output_dim = self.geo_encoder.nnet.last_layer.dim_out
+        return []
+
+    @override
+    def forward(
+        self,
+        batch: Dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        """Forward pass of satclip encoder."""
+
+        coords = batch.get("eo", {}).get("coords")
+
+        # Swap coordinates
+        coords = coords[:, [1, 0]]
+
+        # SatClip needs float64
+        dtype = self.dtype
+        if coords.dtype != dtype:
+            coords = coords.to(dtype)
+
+        feats = self.geo_encoder(coords)
+        if self.extra_projector:
+            feats = self.extra_projector(feats)
+
+        return feats