Skip to content

rslearn.models.pooling_decoder

pooling_decoder

Simple pooling decoder.

PoolingDecoder

Bases: IntermediateComponent

Decoder that computes flat vector from a 2D feature map.

It inputs multi-scale features, but only uses the last feature map. Then applies a configurable number of convolutional layers before pooling, and a configurable number of fully connected layers after pooling.

For patch-level regression from a spatial backbone (e.g. OlmoEarth), set out_channels=1 to produce a Bx1 FeatureVector that feeds directly into RegressionHead without any additional projection step.

Source code in rslearn/models/pooling_decoder.py 
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
class PoolingDecoder(IntermediateComponent):
    """Decoder that computes flat vector from a 2D feature map.

    It inputs multi-scale features, but only uses the last feature map. Then applies a
    configurable number of convolutional layers before pooling, and a configurable
    number of fully connected layers after pooling.

    For patch-level regression from a spatial backbone (e.g. OlmoEarth), set
    out_channels=1 to produce a Bx1 FeatureVector that feeds directly into
    RegressionHead without any additional projection step.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_conv_layers: int = 0,
        num_fc_layers: int = 0,
        conv_channels: int = 128,
        fc_channels: int = 512,
    ) -> None:
        """Initialize a PoolingDecoder.

        Args:
            in_channels: input channels (channels in the last feature map passed to
                this module)
            out_channels: channels for the output flat feature vector
            num_conv_layers: number of convolutional layers to apply, default 0
            num_fc_layers: number of fully-connected layers to apply, default 0
            conv_channels: number of channels to use for convolutional layers
            fc_channels: number of channels to use for fully-connected layers
        """
        super().__init__()
        conv_layers = []
        prev_channels = in_channels
        for _ in range(num_conv_layers):
            conv_layer = torch.nn.Sequential(
                torch.nn.Conv2d(prev_channels, conv_channels, 3, padding=1),
                torch.nn.ReLU(inplace=True),
            )
            conv_layers.append(conv_layer)
            prev_channels = conv_channels
        self.conv_layers = torch.nn.Sequential(*conv_layers)

        fc_layers = []
        for _ in range(num_fc_layers):
            fc_layer = torch.nn.Sequential(
                torch.nn.Linear(prev_channels, fc_channels),
                torch.nn.ReLU(inplace=True),
            )
            fc_layers.append(fc_layer)
            prev_channels = fc_channels
        self.fc_layers = torch.nn.Sequential(*fc_layers)

        self.output_layer = torch.nn.Linear(prev_channels, out_channels)

    def forward(self, intermediates: Any, context: ModelContext) -> Any:
        """Compute flat output vector from multi-scale feature map.

        Args:
            intermediates: the output from the previous component, which must be a FeatureMaps.
            context: the model context.

        Returns:
            flat feature vector
        """
        if not isinstance(intermediates, FeatureMaps):
            raise ValueError("input to PoolingDecoder must be a FeatureMaps")

        # Only use last feature map.
        features = intermediates.feature_maps[-1]

        features = self.conv_layers(features)
        features = torch.amax(features, dim=(2, 3))
        features = self.fc_layers(features)
        return FeatureVector(self.output_layer(features))

forward 

forward(intermediates: Any, context: ModelContext) -> Any

Compute flat output vector from multi-scale feature map.

Parameters:

Name Type Description Default
intermediates Any

the output from the previous component, which must be a FeatureMaps.

 required
context ModelContext

the model context.

 required

Returns:

Type Description
Any

flat feature vector
Source code in rslearn/models/pooling_decoder.py 
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def forward(self, intermediates: Any, context: ModelContext) -> Any:
    """Compute flat output vector from multi-scale feature map.

    Args:
        intermediates: the output from the previous component, which must be a FeatureMaps.
        context: the model context.

    Returns:
        flat feature vector
    """
    if not isinstance(intermediates, FeatureMaps):
        raise ValueError("input to PoolingDecoder must be a FeatureMaps")

    # Only use last feature map.
    features = intermediates.feature_maps[-1]

    features = self.conv_layers(features)
    features = torch.amax(features, dim=(2, 3))
    features = self.fc_layers(features)
    return FeatureVector(self.output_layer(features))

SegmentationPoolingDecoder

Bases: PoolingDecoder

Like PoolingDecoder, but copy output to all pixels.

This allows for the model to produce a global output while still being compatible with SegmentationTask. This only makes sense for very small windows, since the output probabilities will be the same at all pixels. The main use case is to train for a classification-like task on small windows, but still produce a raster during inference on large windows.

Source code in rslearn/models/pooling_decoder.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
class SegmentationPoolingDecoder(PoolingDecoder):
    """Like PoolingDecoder, but copy output to all pixels.

    This allows for the model to produce a global output while still being compatible
    with SegmentationTask. This only makes sense for very small windows, since the
    output probabilities will be the same at all pixels. The main use case is to train
    for a classification-like task on small windows, but still produce a raster during
    inference on large windows.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        image_key: str = "image",
        **kwargs: Any,
    ):
        """Create a new SegmentationPoolingDecoder.

        Args:
            in_channels: input channels (channels in the last feature map passed to
                this module)
            out_channels: channels for the output flat feature vector
            image_key: the key in inputs for the image from which the expected width
                and height is derived.
            kwargs: other arguments to pass to PoolingDecoder.
        """
        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
        self.image_key = image_key

    def forward(self, intermediates: Any, context: ModelContext) -> Any:
        """Extend PoolingDecoder forward to upsample the output to a segmentation mask.

        This only works when all of the pixels have the same segmentation target.
        """
        output_probs = super().forward(intermediates, context)
        # get HW from CTHW image
        h, w = context.inputs[0][self.image_key].image.shape[2:4]
        # BC -> BCHW
        feat_map = output_probs.feature_vector[:, :, None, None].repeat([1, 1, h, w])
        return FeatureMaps([feat_map])

forward 

forward(intermediates: Any, context: ModelContext) -> Any

Extend PoolingDecoder forward to upsample the output to a segmentation mask.

This only works when all of the pixels have the same segmentation target.

Source code in rslearn/models/pooling_decoder.py 
124
125
126
127
128
129
130
131
132
133
134
def forward(self, intermediates: Any, context: ModelContext) -> Any:
    """Extend PoolingDecoder forward to upsample the output to a segmentation mask.

    This only works when all of the pixels have the same segmentation target.
    """
    output_probs = super().forward(intermediates, context)
    # get HW from CTHW image
    h, w = context.inputs[0][self.image_key].image.shape[2:4]
    # BC -> BCHW
    feat_map = output_probs.feature_vector[:, :, None, None].repeat([1, 1, h, w])
    return FeatureMaps([feat_map])