Skip to content

rslearn.models.use_croma

use_croma

CROMA code, copied from https://github.com/antofuller/CROMA.

The code is released under:

MIT License Copyright (c) 2023 Anthony Fuller

PretrainedCROMA

Bases: Module

Pre-trained CROMA model.

Source code in rslearn/models/use_croma.py 
 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
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
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
class PretrainedCROMA(nn.Module):
    """Pre-trained CROMA model."""

    def __init__(
        self,
        pretrained_path: str = "CROMA_base.pt",
        size: str = "base",
        modality: str = "both",
        image_resolution: int = 120,
    ):
        """Create a new PretrainedCROMA.

        NOTE: image_resolution is not the spatial, spectral, or temporal resolution. It is the height and width of the image, in pixels.
        E.g., CROMA was pretrained on 120x120px images, hence image_resolution is 120 by default
        """
        super().__init__()
        # check types
        assert isinstance(pretrained_path, str), (
            f"pretrained_path must be a string, not {type(pretrained_path)}"
        )
        assert isinstance(size, str), f"size must be a string, not {type(size)}"
        assert isinstance(modality, str), (
            f"modality must be a string, not {type(modality)}"
        )
        assert isinstance(image_resolution, int), (
            f"image_resolution must be an int, not {type(image_resolution)}"
        )

        # check values
        assert size in [
            "base",
            "large",
        ], f"size must be either base or large, not {size}"
        assert image_resolution % 8 == 0, (
            f"image_resolution must be a multiple of 8, not {image_resolution}"
        )
        assert modality in [
            "both",
            "SAR",
            "optical",
        ], f"modality must be either both, SAR, or optical, not {modality}"

        # warn the user if the path contains a different size than the size parameter
        if size == "base" and "large" in pretrained_path:
            warnings.warn(
                "The size is set to base, but the word large appears in the pretrained path!"
            )
        elif size == "large" and "base" in pretrained_path:
            warnings.warn(
                "The size is set to large, but the word base appears in the pretrained path!"
            )

        if size == "base":
            self.encoder_dim = 768
            self.encoder_depth = 12
            self.num_heads = 16
            self.patch_size = 8
        else:
            # large by default
            self.encoder_dim = 1024
            self.encoder_depth = 24
            self.num_heads = 16
            self.patch_size = 8

        self.modality = modality
        self.num_patches = int((image_resolution / 8) ** 2)
        self.s1_channels = 2  # fixed at 2 SAR backscatter channels
        self.s2_channels = 12  # fixed at 12 multispectral optical channels
        self.attn_bias = get_2dalibi(
            num_heads=self.num_heads, num_patches=self.num_patches
        )

        if modality in ["SAR", "both"]:
            print("Initializing SAR encoder")
            self.s1_encoder = ViT(
                dim=self.encoder_dim,
                depth=int(self.encoder_depth / 2),
                in_channels=self.s1_channels,
            )
            self.GAP_FFN_s1 = nn.Sequential(
                nn.LayerNorm(self.encoder_dim),
                nn.Linear(
                    self.encoder_dim, int(4 * self.encoder_dim)
                ),  # (BSZ, num_patches, inner_dim)
                nn.GELU(),  # (BSZ, num_patches, inner_dim)
                nn.Linear(
                    int(4 * self.encoder_dim), self.encoder_dim
                ),  # (BSZ, num_patches, dim)
            )

            # load weights
            self.s1_encoder.load_state_dict(
                torch.load(pretrained_path, weights_only=True)["s1_encoder"]
            )
            self.GAP_FFN_s1.load_state_dict(
                torch.load(pretrained_path, weights_only=True)["s1_GAP_FFN"]
            )

        if modality in ["optical", "both"]:
            print("Initializing optical encoder")
            self.s2_encoder = ViT(
                dim=self.encoder_dim,
                depth=self.encoder_depth,
                in_channels=self.s2_channels,
            )
            self.GAP_FFN_s2 = nn.Sequential(
                nn.LayerNorm(self.encoder_dim),
                nn.Linear(
                    self.encoder_dim, int(4 * self.encoder_dim)
                ),  # (BSZ, num_patches, inner_dim)
                nn.GELU(),  # (BSZ, num_patches, inner_dim)
                nn.Linear(
                    int(4 * self.encoder_dim), self.encoder_dim
                ),  # (BSZ, num_patches, dim)
            )

            # load weights
            self.s2_encoder.load_state_dict(
                torch.load(pretrained_path, weights_only=True)["s2_encoder"]
            )
            self.GAP_FFN_s2.load_state_dict(
                torch.load(pretrained_path, weights_only=True)["s2_GAP_FFN"]
            )

        if modality == "both":
            print("Initializing joint SAR-optical encoder")
            self.cross_encoder = BaseTransformerCrossAttn(
                dim=self.encoder_dim,
                depth=int(self.encoder_depth / 2),
                num_heads=self.num_heads,
            )

            # load weights
            self.cross_encoder.load_state_dict(
                torch.load(pretrained_path, weights_only=True)["joint_encoder"]
            )

    def forward(
        self,
        SAR_images: torch.Tensor | None = None,
        optical_images: torch.Tensor | None = None,
    ) -> dict[str, torch.Tensor]:
        """Forward pass through PretrainedCROMA."""
        return_dict = {}
        if self.modality in ["SAR", "both"]:
            assert SAR_images is not None, (
                f"Modality is set to {self.modality}, but SAR_images are None"
            )
            SAR_encodings = self.s1_encoder(
                imgs=SAR_images, attn_bias=self.attn_bias.to(SAR_images.device)
            )  # (bsz, num_patches, encoder_dim)
            SAR_GAP = self.GAP_FFN_s1(SAR_encodings.mean(dim=1))  # (bsz, encoder_dim)
            return_dict["SAR_encodings"] = SAR_encodings
            return_dict["SAR_GAP"] = SAR_GAP

        if self.modality in ["optical", "both"]:
            assert optical_images is not None, (
                f"Modality is set to {self.modality}, but optical_images are None"
            )
            optical_encodings = self.s2_encoder(
                imgs=optical_images, attn_bias=self.attn_bias.to(optical_images.device)
            )  # (bsz, num_patches, encoder_dim)
            optical_GAP = self.GAP_FFN_s2(
                optical_encodings.mean(dim=1)
            )  # (bsz, encoder_dim)
            return_dict["optical_encodings"] = optical_encodings
            return_dict["optical_GAP"] = optical_GAP

        if self.modality == "both":
            assert SAR_images is not None
            assert optical_images is not None
            joint_encodings = self.cross_encoder(
                x=SAR_encodings,
                context=optical_encodings,
                relative_position_bias=self.attn_bias.to(optical_images.device),
            )  # (bsz, num_patches, encoder_dim)
            joint_GAP = joint_encodings.mean(dim=1)  # (bsz, encoder_dim)
            return_dict["joint_encodings"] = joint_encodings
            return_dict["joint_GAP"] = joint_GAP

        return return_dict

forward 

forward(SAR_images: Tensor | None = None, optical_images: Tensor | None = None) -> dict[str, Tensor]

Forward pass through PretrainedCROMA.

Source code in rslearn/models/use_croma.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def forward(
    self,
    SAR_images: torch.Tensor | None = None,
    optical_images: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
    """Forward pass through PretrainedCROMA."""
    return_dict = {}
    if self.modality in ["SAR", "both"]:
        assert SAR_images is not None, (
            f"Modality is set to {self.modality}, but SAR_images are None"
        )
        SAR_encodings = self.s1_encoder(
            imgs=SAR_images, attn_bias=self.attn_bias.to(SAR_images.device)
        )  # (bsz, num_patches, encoder_dim)
        SAR_GAP = self.GAP_FFN_s1(SAR_encodings.mean(dim=1))  # (bsz, encoder_dim)
        return_dict["SAR_encodings"] = SAR_encodings
        return_dict["SAR_GAP"] = SAR_GAP

    if self.modality in ["optical", "both"]:
        assert optical_images is not None, (
            f"Modality is set to {self.modality}, but optical_images are None"
        )
        optical_encodings = self.s2_encoder(
            imgs=optical_images, attn_bias=self.attn_bias.to(optical_images.device)
        )  # (bsz, num_patches, encoder_dim)
        optical_GAP = self.GAP_FFN_s2(
            optical_encodings.mean(dim=1)
        )  # (bsz, encoder_dim)
        return_dict["optical_encodings"] = optical_encodings
        return_dict["optical_GAP"] = optical_GAP

    if self.modality == "both":
        assert SAR_images is not None
        assert optical_images is not None
        joint_encodings = self.cross_encoder(
            x=SAR_encodings,
            context=optical_encodings,
            relative_position_bias=self.attn_bias.to(optical_images.device),
        )  # (bsz, num_patches, encoder_dim)
        joint_GAP = joint_encodings.mean(dim=1)  # (bsz, encoder_dim)
        return_dict["joint_encodings"] = joint_encodings
        return_dict["joint_GAP"] = joint_GAP

    return return_dict

FFN

Bases: Module

Feed-forward network block.

Source code in rslearn/models/use_croma.py 
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
class FFN(nn.Module):
    """Feed-forward network block."""

    def __init__(
        self,
        dim: int,
        mult: int = 4,
        dropout: float = 0.0,
    ):
        """Create a new FFN.

        Args:
            dim: the input dimension.
            mult: the MLP factor (how much larger the hidden dimension should be).
            dropout: the dropout rate.
        """
        super().__init__()
        inner_dim = int(dim * mult)

        self.net = nn.Sequential(
            nn.Linear(dim, inner_dim),  # (BSZ, num_patches, inner_dim)
            nn.GELU(),  # (BSZ, num_patches, inner_dim)
            nn.Dropout(dropout),  # (BSZ, num_patches, inner_dim)
            nn.Linear(inner_dim, dim),  # (BSZ, num_patches, dim)
        )
        self.input_norm = nn.LayerNorm(dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the FFN."""
        x = self.input_norm(x)  # (BSZ, num_patches, dim)
        return self.net(x)  # (BSZ, num_patches, dim)

forward 

forward(x: Tensor) -> Tensor

Forward pass through the FFN.

Source code in rslearn/models/use_croma.py 
267
268
269
270
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the FFN."""
    x = self.input_norm(x)  # (BSZ, num_patches, dim)
    return self.net(x)  # (BSZ, num_patches, dim)

Attention

Bases: Module

Attention block.

Source code in rslearn/models/use_croma.py 
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
class Attention(nn.Module):
    """Attention block."""

    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        dropout: float = 0.0,
    ):
        """Create a new Attention."""
        super().__init__()
        self.num_heads = num_heads
        assert dim % num_heads == 0, "dim must be evenly divisible by num_heads"
        dim_head = int(dim / num_heads)
        self.scale = dim_head**-0.5

        self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
        self.to_out = nn.Linear(dim, dim)
        self.input_norm = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self, x: torch.Tensor, relative_position_bias: torch.Tensor
    ) -> torch.Tensor:
        """Forward pass through the Attention."""
        x = self.input_norm(x)  # (BSZ, num_patches, dim)
        q, k, v = self.to_qkv(x).chunk(3, dim=-1)  # (BSZ, num_patches, dim)
        q, k, v = map(
            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v)
        )  # (BSZ, num_heads, num_patches, dim_head)

        attention_scores = (
            einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
        )  # (BSZ, num_heads, num_patches, num_patches)
        attention_scores = (
            attention_scores + relative_position_bias
        )  # (BSZ, num_heads, num_patches, num_patches)

        attn = attention_scores.softmax(
            dim=-1
        )  # (BSZ, num_heads, num_patches, num_patches)
        attn = self.dropout(attn)  # (BSZ, num_heads, num_patches, num_patches)

        out = einsum(
            "b h i j, b h j d -> b h i d", attn, v
        )  # (BSZ, num_heads, num_patches, dim_head)
        out = rearrange(out, "b h n d -> b n (h d)")  # (BSZ, num_patches, dim)
        return self.to_out(out)  # (BSZ, num_patches, dim)

forward 

forward(x: Tensor, relative_position_bias: Tensor) -> Tensor

Forward pass through the Attention.

Source code in rslearn/models/use_croma.py 
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
def forward(
    self, x: torch.Tensor, relative_position_bias: torch.Tensor
) -> torch.Tensor:
    """Forward pass through the Attention."""
    x = self.input_norm(x)  # (BSZ, num_patches, dim)
    q, k, v = self.to_qkv(x).chunk(3, dim=-1)  # (BSZ, num_patches, dim)
    q, k, v = map(
        lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v)
    )  # (BSZ, num_heads, num_patches, dim_head)

    attention_scores = (
        einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
    )  # (BSZ, num_heads, num_patches, num_patches)
    attention_scores = (
        attention_scores + relative_position_bias
    )  # (BSZ, num_heads, num_patches, num_patches)

    attn = attention_scores.softmax(
        dim=-1
    )  # (BSZ, num_heads, num_patches, num_patches)
    attn = self.dropout(attn)  # (BSZ, num_heads, num_patches, num_patches)

    out = einsum(
        "b h i j, b h j d -> b h i d", attn, v
    )  # (BSZ, num_heads, num_patches, dim_head)
    out = rearrange(out, "b h n d -> b n (h d)")  # (BSZ, num_patches, dim)
    return self.to_out(out)  # (BSZ, num_patches, dim)

CrossAttention

Bases: Module

Cross-attention block.

Source code in rslearn/models/use_croma.py 
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
class CrossAttention(nn.Module):
    """Cross-attention block."""

    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        dropout: float = 0.0,
    ):
        """Create a new CrossAttention."""
        super().__init__()
        self.num_heads = num_heads
        assert dim % num_heads == 0, "dim must be evenly divisible by num_heads"
        dim_head = int(dim / num_heads)
        self.scale = dim_head**-0.5

        self.to_q = nn.Linear(dim, dim, bias=False)
        self.to_k = nn.Linear(dim, dim, bias=False)
        self.to_v = nn.Linear(dim, dim, bias=False)

        self.to_out = nn.Linear(dim, dim)
        self.input_norm = nn.LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor,
        relative_position_bias: torch.Tensor,
    ) -> torch.Tensor:
        """Forward pass through the CrossAttention."""
        x = self.input_norm(x)  # (BSZ, num_patches, dim)
        context = self.input_norm(context)  # (BSZ, num_patches, dim)

        q = self.to_q(x)  # (BSZ, num_patches, dim)
        k = self.to_k(context)  # (BSZ, num_patches, dim)
        v = self.to_v(context)  # (BSZ, num_patches, dim)

        q, k, v = map(
            lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v)
        )  # (BSZ, num_heads, num_patches, dim_head)

        attention_scores = (
            einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
        )  # (BSZ, num_heads, num_patches, num_patches)
        attention_scores = (
            attention_scores + relative_position_bias
        )  # (BSZ, num_heads, num_patches, num_patches)

        attn = attention_scores.softmax(
            dim=-1
        )  # (BSZ, num_heads, num_patches, num_patches)
        attn = self.dropout(attn)  # (BSZ, num_heads, num_patches, num_patches)

        out = einsum(
            "b h i j, b h j d -> b h i d", attn, v
        )  # (BSZ, num_heads, num_patches, dim_head)
        out = rearrange(out, "b h n d -> b n (h d)")  # (BSZ, num_patches, dim)
        return self.to_out(out)  # (BSZ, num_patches, dim)

forward 

forward(x: Tensor, context: Tensor, relative_position_bias: Tensor) -> Tensor

Forward pass through the CrossAttention.

Source code in rslearn/models/use_croma.py 
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
def forward(
    self,
    x: torch.Tensor,
    context: torch.Tensor,
    relative_position_bias: torch.Tensor,
) -> torch.Tensor:
    """Forward pass through the CrossAttention."""
    x = self.input_norm(x)  # (BSZ, num_patches, dim)
    context = self.input_norm(context)  # (BSZ, num_patches, dim)

    q = self.to_q(x)  # (BSZ, num_patches, dim)
    k = self.to_k(context)  # (BSZ, num_patches, dim)
    v = self.to_v(context)  # (BSZ, num_patches, dim)

    q, k, v = map(
        lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), (q, k, v)
    )  # (BSZ, num_heads, num_patches, dim_head)

    attention_scores = (
        einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
    )  # (BSZ, num_heads, num_patches, num_patches)
    attention_scores = (
        attention_scores + relative_position_bias
    )  # (BSZ, num_heads, num_patches, num_patches)

    attn = attention_scores.softmax(
        dim=-1
    )  # (BSZ, num_heads, num_patches, num_patches)
    attn = self.dropout(attn)  # (BSZ, num_heads, num_patches, num_patches)

    out = einsum(
        "b h i j, b h j d -> b h i d", attn, v
    )  # (BSZ, num_heads, num_patches, dim_head)
    out = rearrange(out, "b h n d -> b n (h d)")  # (BSZ, num_patches, dim)
    return self.to_out(out)  # (BSZ, num_patches, dim)

BaseTransformer

Bases: Module

Base transformer class.

Source code in rslearn/models/use_croma.py 
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
class BaseTransformer(nn.Module):
    """Base transformer class."""

    def __init__(
        self,
        dim: int,
        depth: int,
        num_heads: int = 8,
        attn_dropout: float = 0.0,
        ff_dropout: float = 0.0,
        ff_mult: int = 4,
        final_norm: bool = True,
    ):
        """Create a new BaseTransformer."""
        super().__init__()
        self.final_norm = final_norm
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        Attention(dim=dim, num_heads=num_heads, dropout=attn_dropout),
                        FFN(dim=dim, mult=ff_mult, dropout=ff_dropout),
                    ]
                )
            )

        if self.final_norm:
            self.norm_out = nn.LayerNorm(dim)

    def forward(
        self, x: torch.Tensor, relative_position_bias: torch.Tensor
    ) -> torch.Tensor:
        """Forward pass through the BaseTransformer."""
        for self_attn, ffn in self.layers:
            x = self_attn(x, relative_position_bias) + x  # (BSZ, num_patches, dim)
            x = ffn(x) + x  # (BSZ, num_patches, dim)

        if self.final_norm:
            return self.norm_out(x)
        else:
            return x

forward 

forward(x: Tensor, relative_position_bias: Tensor) -> Tensor

Forward pass through the BaseTransformer.

Source code in rslearn/models/use_croma.py 
414
415
416
417
418
419
420
421
422
423
424
425
def forward(
    self, x: torch.Tensor, relative_position_bias: torch.Tensor
) -> torch.Tensor:
    """Forward pass through the BaseTransformer."""
    for self_attn, ffn in self.layers:
        x = self_attn(x, relative_position_bias) + x  # (BSZ, num_patches, dim)
        x = ffn(x) + x  # (BSZ, num_patches, dim)

    if self.final_norm:
        return self.norm_out(x)
    else:
        return x

BaseTransformerCrossAttn

Bases: Module

Base transformer class for cross attention.

Source code in rslearn/models/use_croma.py 
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
class BaseTransformerCrossAttn(nn.Module):
    """Base transformer class for cross attention."""

    def __init__(
        self,
        dim: int,
        depth: int,
        num_heads: int = 8,
        attn_dropout: float = 0.0,
        ff_dropout: float = 0.0,
        ff_mult: int = 4,
    ):
        """Create a new BaseTransformerCrossAttn."""
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(
                nn.ModuleList(
                    [
                        Attention(dim=dim, num_heads=num_heads, dropout=attn_dropout),
                        CrossAttention(
                            dim=dim, num_heads=num_heads, dropout=attn_dropout
                        ),
                        FFN(dim=dim, mult=ff_mult, dropout=ff_dropout),
                    ]
                )
            )

        self.norm_out = nn.LayerNorm(dim)

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor,
        relative_position_bias: torch.Tensor,
    ) -> torch.Tensor:
        """Forward pass through the BaseTransformerCrossAttn."""
        for self_attn, cross_attn, ffn in self.layers:
            x = self_attn(x, relative_position_bias) + x  # (BSZ, num_patches, dim)
            x = (
                cross_attn(x, context, relative_position_bias) + x
            )  # (BSZ, num_patches, dim)
            x = ffn(x) + x  # (BSZ, num_patches, dim)

        x = self.norm_out(x)
        return x  # (BSZ, num_patches, dim)

forward 

forward(x: Tensor, context: Tensor, relative_position_bias: Tensor) -> Tensor

Forward pass through the BaseTransformerCrossAttn.

Source code in rslearn/models/use_croma.py 
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
def forward(
    self,
    x: torch.Tensor,
    context: torch.Tensor,
    relative_position_bias: torch.Tensor,
) -> torch.Tensor:
    """Forward pass through the BaseTransformerCrossAttn."""
    for self_attn, cross_attn, ffn in self.layers:
        x = self_attn(x, relative_position_bias) + x  # (BSZ, num_patches, dim)
        x = (
            cross_attn(x, context, relative_position_bias) + x
        )  # (BSZ, num_patches, dim)
        x = ffn(x) + x  # (BSZ, num_patches, dim)

    x = self.norm_out(x)
    return x  # (BSZ, num_patches, dim)

ViT

Bases: Module

ViT model.

Source code in rslearn/models/use_croma.py 
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
class ViT(nn.Module):
    """ViT model."""

    def __init__(self, dim: int, depth: int, in_channels: int):
        """Create a new ViT."""
        super().__init__()
        self.depth = depth
        self.in_channels = in_channels
        self.dim = dim
        self.num_heads = 16  # always 16, for base and large models
        self.patch_size = 8  # always 8, for base and large models

        pixels_per_patch = int(self.patch_size * self.patch_size * in_channels)
        self.linear_input = nn.Linear(pixels_per_patch, self.dim)
        self.transformer = BaseTransformer(
            dim=self.dim,
            depth=self.depth,
            num_heads=self.num_heads,
        )

    def forward(self, imgs: torch.Tensor, attn_bias: torch.Tensor) -> torch.Tensor:
        """Forward pass through the ViT."""
        x = rearrange(
            imgs,
            "b c (h i) (w j) -> b (h w) (c i j)",
            i=self.patch_size,
            j=self.patch_size,
        )
        # x is shape -> (bsz, num_patches, self.channels*self.patch_size*self.patch_size)

        x = self.linear_input(x)  # (bsz, num_patches, dim)
        x = self.transformer(x, relative_position_bias=attn_bias)
        return x

forward 

forward(imgs: Tensor, attn_bias: Tensor) -> Tensor

Forward pass through the ViT.

Source code in rslearn/models/use_croma.py 
496
497
498
499
500
501
502
503
504
505
506
507
508
def forward(self, imgs: torch.Tensor, attn_bias: torch.Tensor) -> torch.Tensor:
    """Forward pass through the ViT."""
    x = rearrange(
        imgs,
        "b c (h i) (w j) -> b (h w) (c i j)",
        i=self.patch_size,
        j=self.patch_size,
    )
    # x is shape -> (bsz, num_patches, self.channels*self.patch_size*self.patch_size)

    x = self.linear_input(x)  # (bsz, num_patches, dim)
    x = self.transformer(x, relative_position_bias=attn_bias)
    return x

get_2dalibi 

get_2dalibi(num_heads: int, num_patches: int) -> Tensor

Get 2D bias initialization for attention layer.

Parameters:

Name Type Description Default
num_heads int

the number of heads in the attention layer.

 required
num_patches int

the total number of patches, which should be a square.

 required
Source code in rslearn/models/use_croma.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def get_2dalibi(num_heads: int, num_patches: int) -> torch.Tensor:
    """Get 2D bias initialization for attention layer.

    Args:
        num_heads: the number of heads in the attention layer.
        num_patches: the total number of patches, which should be a square.
    """
    # inspired by: https://github.com/ofirpress/attention_with_linear_biases
    points = list(
        itertools.product(
            range(int(math.sqrt(num_patches))), range(int(math.sqrt(num_patches)))
        )
    )

    def get_slopes(n: int) -> list[float]:
        def get_slopes_power_of_2(n: int) -> list[float]:
            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
            ratio = start
            return [start * ratio**i for i in range(n)]

        if math.log2(n).is_integer():
            return get_slopes_power_of_2(n)
        else:
            closest_power_of_2 = 2 ** math.floor(math.log2(n))
            return (
                get_slopes_power_of_2(closest_power_of_2)
                + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
            )

    slopes = torch.Tensor(get_slopes(num_heads)).unsqueeze(1)
    idxs = []
    for p1 in points:
        for p2 in points:
            dist = math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
            idxs.append(dist * slopes * -1)
    all_bias = torch.cat(idxs, dim=1)
    return all_bias.view(1, num_heads, num_patches, num_patches)