Skip to content

rslearn.models.ssl4eo_s12

ssl4eo_s12

SSL4EO-S12 models.

Ssl4eoS12

Bases: FeatureExtractor

The SSL4EO-S12 family of pretrained models.

Source code in rslearn/models/ssl4eo_s12.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
class Ssl4eoS12(FeatureExtractor):
    """The SSL4EO-S12 family of pretrained models."""

    def __init__(
        self,
        backbone_ckpt_path: str | None,
        arch: str = "resnet50",
        output_layers: list[int] = [0, 1, 2, 3],
    ) -> None:
        """Instantiate a new Swin instance.

        Args:
            backbone_ckpt_path: the path to the backbone checkpoint, need to download
                from https://github.com/zhu-xlab/SSL4EO-S12
            arch: model architecture, currently only resnet50 is supported
            output_layers: return the outputs from these layers, defaults to all of the
                four resnet outputs.
        """
        super().__init__()
        self.arch = arch
        self.output_layers = output_layers

        if arch == "resnet50":
            self.model = torchvision.models.resnet50()
            self.model.conv1 = torch.nn.Conv2d(
                13, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
            )

        else:
            raise ValueError(f"unknown SSL4EO-S12 architecture {arch}")

        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) -> 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
        to the feature maps that the backbone returns. For example, an element [2, 32]
        indicates that the corresponding feature map is 1/2 the input resolution and
        has 32 channels.

        Returns:
            the output channels of the backbone as a list of (downsample_factor, depth)
            tuples.
        """
        if self.arch == "resnet50":
            all_out_channels = [
                (4, 256),
                (8, 512),
                (16, 1024),
                (32, 2048),
            ]
        return [all_out_channels[idx] for idx in self.output_layers]

    def forward(
        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.

        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"].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)
        x = self.model.maxpool(x)

        layer1 = self.model.layer1(x)
        layer2 = self.model.layer2(layer1)
        layer3 = self.model.layer3(layer2)
        layer4 = self.model.layer4(layer3)
        all_features = [layer1, layer2, layer3, layer4]
        return FeatureMaps([all_features[idx] for idx in self.output_layers])

get_backbone_channels 

get_backbone_channels() -> 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 to the feature maps that the backbone returns. For example, an element [2, 32] indicates that the corresponding feature map is 1/2 the input resolution and has 32 channels.

Returns:

Type Description
list[tuple[int, int]]

the output channels of the backbone as a list of (downsample_factor, depth)
list[tuple[int, int]]

tuples.
Source code in rslearn/models/ssl4eo_s12.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
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
    to the feature maps that the backbone returns. For example, an element [2, 32]
    indicates that the corresponding feature map is 1/2 the input resolution and
    has 32 channels.

    Returns:
        the output channels of the backbone as a list of (downsample_factor, depth)
        tuples.
    """
    if self.arch == "resnet50":
        all_out_channels = [
            (4, 256),
            (8, 512),
            (16, 1024),
            (32, 2048),
        ]
    return [all_out_channels[idx] for idx in self.output_layers]

forward 

forward(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.

Parameters:

Name Type Description Default
context ModelContext

the model context. Input dicts must include "image" key containing the images to process.

 required

Returns:

Type Description
FeatureMaps

feature maps computed by the pre-trained model.
Source code in rslearn/models/ssl4eo_s12.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def forward(
    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.

    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"].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)
    x = self.model.maxpool(x)

    layer1 = self.model.layer1(x)
    layer2 = self.model.layer2(layer1)
    layer3 = self.model.layer3(layer2)
    layer4 = self.model.layer4(layer3)
    all_features = [layer1, layer2, layer3, layer4]
    return FeatureMaps([all_features[idx] for idx in self.output_layers])