Skip to content

rslearn.models.module_wrapper

module_wrapper

Module wrapper provided for backwards compatibility.

EncoderModuleWrapper

Bases: FeatureExtractor

Wraps one or more IntermediateComponents to function as the feature extractor.

The first component should input a FeatureMaps, which will be computed from the overall inputs by stacking the "image" key from each input dict.

Source code in rslearn/models/module_wrapper.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
class EncoderModuleWrapper(FeatureExtractor):
    """Wraps one or more IntermediateComponents to function as the feature extractor.

    The first component should input a FeatureMaps, which will be computed from the
    overall inputs by stacking the "image" key from each input dict.
    """

    def __init__(
        self,
        module: IntermediateComponent | None = None,
        modules: list[IntermediateComponent] = [],
    ):
        """Initialize an EncoderModuleWrapper.

        Args:
            module: the IntermediateComponent to wrap for use as a FeatureExtractor.
                Exactly one of module or modules must be set.
            modules: list of modules to wrap
        """
        super().__init__()
        if module is not None and len(modules) > 0:
            raise ValueError("only one of module or modules should be set")
        if module is not None:
            self.encoder_modules = torch.nn.ModuleList([module])
        elif len(modules) > 0:
            self.encoder_modules = torch.nn.ModuleList(modules)
        else:
            raise ValueError("one of module or modules must be set")

    def forward(self, context: ModelContext) -> Any:
        """Compute outputs from the wrapped module.

        Args:
            context: the model context. Input dicts must include "image" key containing
                the image to convert to a FeatureMaps, which will be passed to the
                first wrapped module.

        Returns:
            the output from the last wrapped module.
        """
        # take the first and only timestep. Currently no intermediate
        # components support multi temporal inputs, so if the input is
        # multitemporal it should be wrapped in a simple time series wrapper.
        images = torch.stack(
            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
        )
        cur: Any = FeatureMaps([images])
        for m in self.encoder_modules:
            cur = m(cur, context)
        return cur

forward 

forward(context: ModelContext) -> Any

Compute outputs from the wrapped module.

Parameters:

Name Type Description Default
context ModelContext

the model context. Input dicts must include "image" key containing the image to convert to a FeatureMaps, which will be passed to the first wrapped module.

 required

Returns:

Type Description
Any

the output from the last wrapped module.
Source code in rslearn/models/module_wrapper.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def forward(self, context: ModelContext) -> Any:
    """Compute outputs from the wrapped module.

    Args:
        context: the model context. Input dicts must include "image" key containing
            the image to convert to a FeatureMaps, which will be passed to the
            first wrapped module.

    Returns:
        the output from the last wrapped module.
    """
    # take the first and only timestep. Currently no intermediate
    # components support multi temporal inputs, so if the input is
    # multitemporal it should be wrapped in a simple time series wrapper.
    images = torch.stack(
        [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
    )
    cur: Any = FeatureMaps([images])
    for m in self.encoder_modules:
        cur = m(cur, context)
    return cur