Skip to content

rslearn.models.swin

swin

Swin Transformer.

Swin

Bases: FeatureExtractor

A Swin Transformer model.

It can either be used stand-alone for classification, or as a feature extractor in combination with other decoder modules.

Source code in rslearn/models/swin.py 
 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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
class Swin(FeatureExtractor):
    """A Swin Transformer model.

    It can either be used stand-alone for classification, or as a feature extractor in
    combination with other decoder modules.
    """

    def __init__(
        self,
        arch: str = "swin_v2_b",
        pretrained: bool = False,
        input_channels: int = 3,
        output_layers: list[int] | None = None,
        num_outputs: int = 1000,
    ) -> None:
        """Instantiate a new Swin instance.

        Args:
            arch: the architecture, e.g. "swin_v2_b" (default) or "swin_t"
            pretrained: set True to use ImageNet pre-trained weights
            input_channels: number of input channels (default 3). If not 3, the first
                layer is updated and will be randomly initialized even if pretrained is
                set.
            output_layers: list of layers to output, default use as classification
                model (output FeatureVector). For feature extraction, [1, 3, 5, 7] is
                recommended.
            num_outputs: number of output logits, defaults to 1000 which matches the
                pretrained models.
        """
        super().__init__()
        self.arch = arch
        self.output_layers = output_layers

        kwargs = {}

        if arch == "swin_t":
            if pretrained:
                kwargs["weights"] = Swin_T_Weights.IMAGENET1K_V1
            self.model = torchvision.models.swin_t(**kwargs)
        elif arch == "swin_s":
            if pretrained:
                kwargs["weights"] = Swin_S_Weights.IMAGENET1K_V1
            self.model = torchvision.models.swin_s(**kwargs)
        elif arch == "swin_b":
            if pretrained:
                kwargs["weights"] = Swin_B_Weights.IMAGENET1K_V1
            self.model = torchvision.models.swin_b(**kwargs)
        elif arch == "swin_v2_t":
            if pretrained:
                kwargs["weights"] = Swin_V2_T_Weights.IMAGENET1K_V1
            self.model = torchvision.models.swin_v2_t(**kwargs)
        elif arch == "swin_v2_s":
            if pretrained:
                kwargs["weights"] = Swin_V2_S_Weights.IMAGENET1K_V1
            self.model = torchvision.models.swin_v2_s(**kwargs)
        elif arch == "swin_v2_b":
            if pretrained:
                kwargs["weights"] = Swin_V2_B_Weights.IMAGENET1K_V1
            self.model = torchvision.models.swin_v2_b(**kwargs)
        else:
            raise ValueError(f"unknown swin architecture {arch}")

        # Adjust first layer to accommodate different number of input channels
        # if needed.
        if input_channels != 3:
            self.model.features[0][0] = torch.nn.Conv2d(
                input_channels,
                self.model.features[0][0].out_channels,
                kernel_size=(4, 4),
                stride=(4, 4),
            )

        # Similarly adjust last layer.
        # We do this rather than passing num_classes since passing it would result in
        # incompatibility with loading the pretrained weights.
        if num_outputs != self.model.head.out_features:
            self.model.head = torch.nn.Linear(self.model.head.in_features, num_outputs)

    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.
        """
        assert self.output_layers

        if self.arch in ["swin_b", "swin_v2_b"]:
            all_out_channels = [
                (4, 128),
                (4, 128),
                (4, 128),
                (8, 256),
                (8, 256),
                (16, 512),
                (16, 512),
                (32, 1024),
                (32, 1024),
            ]
        elif self.arch in ["swin_s", "swin_v2_s", "swin_t", "swin_v2_t"]:
            all_out_channels = [
                (4, 96),
                (4, 96),
                (8, 192),
                (8, 192),
                (16, 384),
                (16, 384),
                (32, 768),
                (32, 768),
            ]
        return [all_out_channels[idx] for idx in self.output_layers]

    def forward(
        self,
        context: ModelContext,
    ) -> FeatureVector | 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 image to process.

        Returns:
            a FeatureVector if the configured output_layers is None, or a FeatureMaps
                otherwise containing one feature map per configured output layer.
        """
        images = torch.stack(
            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
        )

        if self.output_layers:
            layer_features = []
            x = images
            for layer in self.model.features:
                x = layer(x)
                layer_features.append(x.permute(0, 3, 1, 2))
            return FeatureMaps([layer_features[idx] for idx in self.output_layers])

        else:
            return FeatureVector(self.model(images))

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/swin.py 
 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
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.
    """
    assert self.output_layers

    if self.arch in ["swin_b", "swin_v2_b"]:
        all_out_channels = [
            (4, 128),
            (4, 128),
            (4, 128),
            (8, 256),
            (8, 256),
            (16, 512),
            (16, 512),
            (32, 1024),
            (32, 1024),
        ]
    elif self.arch in ["swin_s", "swin_v2_s", "swin_t", "swin_v2_t"]:
        all_out_channels = [
            (4, 96),
            (4, 96),
            (8, 192),
            (8, 192),
            (16, 384),
            (16, 384),
            (32, 768),
            (32, 768),
        ]
    return [all_out_channels[idx] for idx in self.output_layers]

forward 

forward(context: ModelContext) -> FeatureVector | 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 image to process.

 required

Returns:

Type Description
FeatureVector | FeatureMaps

a FeatureVector if the configured output_layers is None, or a FeatureMaps otherwise containing one feature map per configured output layer.
Source code in rslearn/models/swin.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def forward(
    self,
    context: ModelContext,
) -> FeatureVector | 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 image to process.

    Returns:
        a FeatureVector if the configured output_layers is None, or a FeatureMaps
            otherwise containing one feature map per configured output layer.
    """
    images = torch.stack(
        [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs], dim=0
    )

    if self.output_layers:
        layer_features = []
        x = images
        for layer in self.model.features:
            x = layer(x)
            layer_features.append(x.permute(0, 3, 1, 2))
        return FeatureMaps([layer_features[idx] for idx in self.output_layers])

    else:
        return FeatureVector(self.model(images))