Skip to content

rslearn.models.satlaspretrain

satlaspretrain

SatlasPretrain models.

SatlasPretrain

Bases: FeatureExtractor

SatlasPretrain backbones.

Source code in rslearn/models/satlaspretrain.py 
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
class SatlasPretrain(FeatureExtractor):
    """SatlasPretrain backbones."""

    def __init__(
        self, model_identifier: str, fpn: bool = False, resize_to_pretrain: bool = False
    ) -> None:
        """Instantiate a new SatlasPretrain instance.

        Args:
            model_identifier: the checkpoint name from the table at
                https://github.com/allenai/satlaspretrain_models
            fpn: whether to include the feature pyramid network, otherwise only the
                Swin-v2-Transformer is used.
            resize_to_pretrain: whether to resize inputs to the pretraining input
                size (512 x 512)
        """
        super().__init__()
        weights_manager = satlaspretrain_models.Weights()
        self.model = weights_manager.get_pretrained_model(
            model_identifier=model_identifier, fpn=fpn, device="cpu"
        )

        if "SwinB" in model_identifier:
            self.backbone_channels = [
                [4, 128],
                [8, 256],
                [16, 512],
                [32, 1024],
            ]
        elif "SwinT" in model_identifier:
            self.backbone_channels = [
                [4, 96],
                [8, 192],
                [16, 384],
                [32, 768],
            ]
        elif "Resnet" in model_identifier:
            self.backbone_channels = [
                [4, 256],
                [8, 512],
                [16, 1024],
                [32, 2048],
            ]
        self.resize_to_pretrain = resize_to_pretrain

    def maybe_resize(self, data: torch.Tensor) -> list[torch.Tensor]:
        """Resize to pretraining sizes if resize_to_pretrain == True."""
        if self.resize_to_pretrain:
            return F.interpolate(
                data,
                size=(512, 512),
                mode="bilinear",
                align_corners=False,
            )
        else:
            return data

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

        Args:
            context: the model context. Input dicts must contain an "image" key
                containing the image input to the model.

        Returns:
            multi-resolution feature maps computed by the model.
        """
        # take the first (assumed to be only) timestep
        images = torch.stack(
            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
        )
        feature_maps = self.model(self.maybe_resize(images))
        return FeatureMaps(feature_maps)

    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 (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.
        """
        return self.backbone_channels

maybe_resize 

maybe_resize(data: Tensor) -> list[Tensor]

Resize to pretraining sizes if resize_to_pretrain == True.

Source code in rslearn/models/satlaspretrain.py 
57
58
59
60
61
62
63
64
65
66
67
def maybe_resize(self, data: torch.Tensor) -> list[torch.Tensor]:
    """Resize to pretraining sizes if resize_to_pretrain == True."""
    if self.resize_to_pretrain:
        return F.interpolate(
            data,
            size=(512, 512),
            mode="bilinear",
            align_corners=False,
        )
    else:
        return data

forward 

forward(context: ModelContext) -> FeatureMaps

Compute feature maps from the SatlasPretrain backbone.

Parameters:

Name Type Description Default
context ModelContext

the model context. Input dicts must contain an "image" key containing the image input to the model.

 required

Returns:

Type Description
FeatureMaps

multi-resolution feature maps computed by the model.
Source code in rslearn/models/satlaspretrain.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def forward(self, context: ModelContext) -> FeatureMaps:
    """Compute feature maps from the SatlasPretrain backbone.

    Args:
        context: the model context. Input dicts must contain an "image" key
            containing the image input to the model.

    Returns:
        multi-resolution feature maps computed by the model.
    """
    # take the first (assumed to be only) timestep
    images = torch.stack(
        [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
    )
    feature_maps = self.model(self.maybe_resize(images))
    return FeatureMaps(feature_maps)

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

the output channels of the backbone as a list of (downsample_factor, depth)
list

tuples.
Source code in rslearn/models/satlaspretrain.py 
86
87
88
89
90
91
92
93
94
95
96
97
98
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 (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.
    """
    return self.backbone_channels