| """Paper-aligned SatMAE model components. |
| |
| This is an original implementation of the architecture described in SatMAE. |
| The upstream repository was used only as a behavioral reference; no upstream |
| source text is incorporated here. |
| """ |
|
|
| import math |
| from functools import partial |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| def _sincos_1d(values, dim): |
| """Return a fixed sine-cosine embedding for arbitrary scalar positions.""" |
| if dim <= 0: |
| return values.new_zeros((*values.shape, 0)) |
| pairs = (dim + 1) // 2 |
| omega = torch.arange(pairs, device=values.device, dtype=torch.float32) |
| omega = torch.exp(-math.log(10000.0) * omega / max(pairs - 1, 1)) |
| phase = values.to(torch.float32).unsqueeze(-1) * omega |
| return torch.cat((phase.sin(), phase.cos()), dim=-1)[..., :dim] |
|
|
|
|
| def _sincos_2d(grid_size, dim): |
| """Return a fixed row-major 2D sine-cosine position embedding.""" |
| rows, cols = torch.meshgrid( |
| torch.arange(grid_size, dtype=torch.float32), |
| torch.arange(grid_size, dtype=torch.float32), |
| indexing="ij", |
| ) |
| row_dim = dim // 2 |
| return torch.cat( |
| (_sincos_1d(rows.reshape(-1), row_dim), |
| _sincos_1d(cols.reshape(-1), dim - row_dim)), |
| dim=-1, |
| ) |
|
|
|
|
| def _timestamp_embedding(timestamps, dim): |
| """Encode either scalar times or fMoW ``[year, month, hour]`` tuples.""" |
| if timestamps.ndim == 2: |
| return _sincos_1d(timestamps, dim) |
| if timestamps.ndim != 3 or timestamps.shape[-1] != 3: |
| raise ValueError("timestamps must have shape [B, T] or [B, T, 3]") |
| field_dims = [dim // 3] * 3 |
| for index in range(dim % 3): |
| field_dims[index] += 1 |
| return torch.cat( |
| [_sincos_1d(timestamps[..., index], field_dim) |
| for index, field_dim in enumerate(field_dims)], |
| dim=-1, |
| ) |
|
|
|
|
| class PatchEmbed(nn.Module): |
| def __init__(self, image_size, patch_size, in_channels, embed_dim): |
| super().__init__() |
| self.image_size = image_size |
| self.patch_size = patch_size |
| self.num_patches = (image_size // patch_size) ** 2 |
| self.proj = nn.Conv2d( |
| in_channels, embed_dim, kernel_size=patch_size, stride=patch_size |
| ) |
|
|
| def forward(self, images): |
| if images.shape[-2:] != (self.image_size, self.image_size): |
| raise ValueError( |
| f"expected {self.image_size}x{self.image_size} images, " |
| f"got {tuple(images.shape[-2:])}" |
| ) |
| return self.proj(images).flatten(2).transpose(1, 2) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, dim, num_heads, mlp_ratio=4.0, norm_layer=nn.LayerNorm): |
| super().__init__() |
| self.norm1 = norm_layer(dim) |
| self.attention = nn.MultiheadAttention( |
| dim, num_heads, dropout=0.0, bias=True, batch_first=True |
| ) |
| self.norm2 = norm_layer(dim) |
| hidden_dim = int(dim * mlp_ratio) |
| self.mlp = nn.Sequential( |
| nn.Linear(dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, dim) |
| ) |
|
|
| def forward(self, tokens): |
| normalized = self.norm1(tokens) |
| tokens = tokens + self.attention( |
| normalized, normalized, normalized, need_weights=False |
| )[0] |
| return tokens + self.mlp(self.norm2(tokens)) |
|
|
|
|
| class SatMAE(nn.Module): |
| """Masked autoencoder for temporal or grouped multispectral imagery. |
| |
| Temporal inputs use shape ``[B, T, C, H, W]`` and optional timestamps |
| ``[B, T]``. Multispectral inputs use shape ``[B, C, H, W]``. |
| """ |
|
|
| def __init__( |
| self, |
| image_size=224, |
| patch_size=16, |
| in_channels=3, |
| frames=3, |
| embed_dim=1024, |
| encoder_depth=24, |
| encoder_heads=16, |
| decoder_dim=512, |
| decoder_depth=8, |
| decoder_heads=16, |
| mlp_ratio=4.0, |
| mode="temporal", |
| spectral_groups=None, |
| mask_ratio=0.75, |
| norm_pix_loss=False, |
| same_mask=False, |
| spatial_mask=False, |
| temporal_embed_dim=None, |
| decoder_temporal_embed_dim=None, |
| channel_embed_dim=None, |
| decoder_channel_embed_dim=None, |
| norm_layer=None, |
| ): |
| super().__init__() |
| if image_size % patch_size: |
| raise ValueError("image_size must be divisible by patch_size") |
| if not 0.0 <= mask_ratio < 1.0: |
| raise ValueError("mask_ratio must be in [0, 1)") |
| if mode not in {"temporal", "multispectral"}: |
| raise ValueError("mode must be temporal or multispectral") |
| if embed_dim % encoder_heads or decoder_dim % decoder_heads: |
| raise ValueError("embedding dimensions must be divisible by head counts") |
|
|
| norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) |
| self.image_size = image_size |
| self.patch_size = patch_size |
| self.in_channels = in_channels |
| self.frames = frames |
| self.embed_dim = embed_dim |
| self.decoder_dim = decoder_dim |
| self.mode = mode |
| self.mask_ratio = mask_ratio |
| self.norm_pix_loss = norm_pix_loss |
| self.same_mask = same_mask |
| self.spatial_mask = spatial_mask |
| self.grid_size = image_size // patch_size |
| self.num_patches = self.grid_size ** 2 |
|
|
| if mode == "temporal": |
| self.spectral_groups = None |
| self.patch_embed = PatchEmbed( |
| image_size, patch_size, in_channels, embed_dim |
| ) |
| self.token_groups = frames |
| semantic_dim = temporal_embed_dim |
| if semantic_dim is None: |
| semantic_dim = min(128, max(2, embed_dim // 4)) |
| decoder_semantic_dim = decoder_temporal_embed_dim |
| if decoder_semantic_dim is None: |
| decoder_semantic_dim = min(64, max(2, decoder_dim // 4)) |
| prediction_dims = [patch_size ** 2 * in_channels] |
| else: |
| groups = spectral_groups or [list(range(in_channels))] |
| flattened = [channel for group in groups for channel in group] |
| if sorted(flattened) != list(range(in_channels)): |
| raise ValueError("spectral_groups must partition all input channels") |
| self.spectral_groups = tuple(tuple(group) for group in groups) |
| self.patch_embed = nn.ModuleList( |
| PatchEmbed(image_size, patch_size, len(group), embed_dim) |
| for group in self.spectral_groups |
| ) |
| self.token_groups = len(self.spectral_groups) |
| semantic_dim = channel_embed_dim |
| if semantic_dim is None: |
| semantic_dim = min(256, max(2, embed_dim // 4)) |
| decoder_semantic_dim = decoder_channel_embed_dim |
| if decoder_semantic_dim is None: |
| decoder_semantic_dim = min(128, max(2, decoder_dim // 4)) |
| prediction_dims = [patch_size ** 2 * len(g) for g in self.spectral_groups] |
|
|
| if not 0 < semantic_dim < embed_dim: |
| raise ValueError("encoder semantic embedding dimension is invalid") |
| if not 0 < decoder_semantic_dim < decoder_dim: |
| raise ValueError("decoder semantic embedding dimension is invalid") |
| self.semantic_dim = semantic_dim |
| self.decoder_semantic_dim = decoder_semantic_dim |
|
|
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim)) |
| self.register_buffer( |
| "spatial_pos_embed", |
| _sincos_2d(self.grid_size, embed_dim - semantic_dim), |
| persistent=True, |
| ) |
| self.register_buffer( |
| "decoder_spatial_pos_embed", |
| _sincos_2d(self.grid_size, decoder_dim - decoder_semantic_dim), |
| persistent=True, |
| ) |
| if mode == "multispectral": |
| group_ids = torch.arange(self.token_groups, dtype=torch.float32) |
| self.register_buffer( |
| "group_embed", _sincos_1d(group_ids, semantic_dim), persistent=True |
| ) |
| self.register_buffer( |
| "decoder_group_embed", |
| _sincos_1d(group_ids, decoder_semantic_dim), |
| persistent=True, |
| ) |
|
|
| self.blocks = nn.ModuleList( |
| TransformerBlock(embed_dim, encoder_heads, mlp_ratio, norm_layer) |
| for _ in range(encoder_depth) |
| ) |
| self.norm = norm_layer(embed_dim) |
| self.decoder_embed = nn.Linear(embed_dim, decoder_dim) |
| self.decoder_blocks = nn.ModuleList( |
| TransformerBlock(decoder_dim, decoder_heads, mlp_ratio, norm_layer) |
| for _ in range(decoder_depth) |
| ) |
| self.decoder_norm = norm_layer(decoder_dim) |
| self.decoder_pred = nn.ModuleList( |
| nn.Linear(decoder_dim, output_dim) for output_dim in prediction_dims |
| ) |
| self.initialize_weights() |
|
|
| def initialize_weights(self): |
| patch_embeds = ( |
| [self.patch_embed] |
| if isinstance(self.patch_embed, PatchEmbed) |
| else self.patch_embed |
| ) |
| for patch_embed in patch_embeds: |
| nn.init.xavier_uniform_(patch_embed.proj.weight.flatten(1)) |
| if patch_embed.proj.bias is not None: |
| nn.init.zeros_(patch_embed.proj.bias) |
| nn.init.normal_(self.cls_token, std=0.02) |
| nn.init.normal_(self.mask_token, std=0.02) |
| for module in self.modules(): |
| if isinstance(module, nn.Linear): |
| nn.init.xavier_uniform_(module.weight) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.LayerNorm): |
| nn.init.ones_(module.weight) |
| nn.init.zeros_(module.bias) |
|
|
| def patchify(self, images): |
| if images.ndim != 4: |
| raise ValueError("patchify expects [B, C, H, W]") |
| batch, channels, height, width = images.shape |
| patch = self.patch_size |
| if height != width or height != self.image_size: |
| raise ValueError(f"expected square images of size {self.image_size}") |
| patches = images.reshape( |
| batch, channels, height // patch, patch, width // patch, patch |
| ) |
| patches = patches.permute(0, 2, 4, 1, 3, 5) |
| return patches.reshape(batch, self.num_patches, channels * patch ** 2) |
|
|
| def unpatchify(self, patches, channels=None): |
| channels = channels or self.in_channels |
| batch = patches.shape[0] |
| patch = self.patch_size |
| expected = channels * patch ** 2 |
| if patches.shape[1:] != (self.num_patches, expected): |
| raise ValueError("patch tensor has incompatible shape") |
| images = patches.reshape( |
| batch, self.grid_size, self.grid_size, channels, patch, patch |
| ) |
| images = images.permute(0, 3, 1, 4, 2, 5) |
| return images.reshape(batch, channels, self.image_size, self.image_size) |
|
|
| def _random_masking(self, tokens, mask_ratio, share_spatial_mask): |
| batch, length, dim = tokens.shape |
| if share_spatial_mask: |
| units = self.num_patches |
| len_keep_units = int(units * (1.0 - mask_ratio)) |
| noise = torch.rand(batch, units, device=tokens.device) |
| spatial_order = noise.argsort(dim=1) |
| kept = [spatial_order[:, :len_keep_units] + g * units |
| for g in range(self.token_groups)] |
| removed = [spatial_order[:, len_keep_units:] + g * units |
| for g in range(self.token_groups)] |
| ids_shuffle = torch.cat(kept + removed, dim=1) |
| len_keep = len_keep_units * self.token_groups |
| else: |
| len_keep = int(length * (1.0 - mask_ratio)) |
| ids_shuffle = torch.rand(batch, length, device=tokens.device).argsort(dim=1) |
| ids_restore = ids_shuffle.argsort(dim=1) |
| ids_keep = ids_shuffle[:, :len_keep] |
| visible = torch.gather(tokens, 1, ids_keep.unsqueeze(-1).expand(-1, -1, dim)) |
| mask = torch.ones(batch, length, device=tokens.device) |
| mask[:, :len_keep] = 0 |
| mask = torch.gather(mask, 1, ids_restore) |
| return visible, mask, ids_restore |
|
|
| def _temporal_tokens(self, images, timestamps): |
| if images.ndim != 5: |
| raise ValueError("temporal mode expects images shaped [B, T, C, H, W]") |
| batch, frames, channels, _, _ = images.shape |
| if frames != self.frames or channels != self.in_channels: |
| raise ValueError( |
| f"expected T={self.frames}, C={self.in_channels}; got T={frames}, C={channels}" |
| ) |
| if timestamps is None: |
| timestamps = torch.arange(frames, device=images.device).expand(batch, -1) |
| if timestamps.shape[:2] != (batch, frames): |
| raise ValueError( |
| f"timestamps must start with shape {(batch, frames)}, " |
| f"got {tuple(timestamps.shape)}" |
| ) |
| spatial = self.spatial_pos_embed.to(dtype=images.dtype) |
| time = _timestamp_embedding(timestamps, self.semantic_dim).to(dtype=images.dtype) |
| position = torch.cat( |
| (spatial.view(1, 1, self.num_patches, -1).expand(batch, frames, -1, -1), |
| time.unsqueeze(2).expand(-1, -1, self.num_patches, -1)), |
| dim=-1, |
| ).reshape(batch, frames * self.num_patches, self.embed_dim) |
| tokens = torch.stack( |
| [self.patch_embed(images[:, frame]) for frame in range(frames)], dim=1 |
| ).reshape(batch, frames * self.num_patches, self.embed_dim) |
| return tokens + position, timestamps |
|
|
| def _multispectral_tokens(self, images): |
| if images.ndim != 4 or images.shape[1] != self.in_channels: |
| raise ValueError( |
| f"multispectral mode expects images shaped [B, {self.in_channels}, H, W]" |
| ) |
| spatial = self.spatial_pos_embed.to(dtype=images.dtype) |
| group = self.group_embed.to(dtype=images.dtype) |
| positions = torch.cat( |
| (spatial.view(1, self.num_patches, -1).expand(self.token_groups, -1, -1), |
| group.view(self.token_groups, 1, -1).expand(-1, self.num_patches, -1)), |
| dim=-1, |
| ).reshape(1, self.token_groups * self.num_patches, self.embed_dim) |
| tokens = torch.cat( |
| [embed(images[:, channels]) |
| for embed, channels in zip(self.patch_embed, self.spectral_groups)], |
| dim=1, |
| ) |
| return tokens + positions |
|
|
| def forward_encoder(self, images, timestamps=None, mask_ratio=None): |
| ratio = self.mask_ratio if mask_ratio is None else mask_ratio |
| if not 0.0 <= ratio < 1.0: |
| raise ValueError("mask_ratio must be in [0, 1)") |
| if self.mode == "temporal": |
| tokens, timestamps = self._temporal_tokens(images, timestamps) |
| shared = self.same_mask |
| else: |
| tokens = self._multispectral_tokens(images) |
| shared = self.spatial_mask |
| tokens, mask, ids_restore = self._random_masking(tokens, ratio, shared) |
| cls = self.cls_token.expand(tokens.shape[0], -1, -1) |
| tokens = torch.cat((cls, tokens), dim=1) |
| for block in self.blocks: |
| tokens = block(tokens) |
| return self.norm(tokens), mask, ids_restore, timestamps |
|
|
| def _decoder_positions(self, batch, timestamps, dtype, device): |
| spatial = self.decoder_spatial_pos_embed.to(device=device, dtype=dtype) |
| if self.mode == "temporal": |
| semantic = _timestamp_embedding(timestamps, self.decoder_semantic_dim).to(dtype=dtype) |
| else: |
| semantic = self.decoder_group_embed.to(device=device, dtype=dtype) |
| semantic = semantic.unsqueeze(0).expand(batch, -1, -1) |
| position = torch.cat( |
| (spatial.view(1, 1, self.num_patches, -1).expand(batch, self.token_groups, -1, -1), |
| semantic.unsqueeze(2).expand(-1, -1, self.num_patches, -1)), |
| dim=-1, |
| ) |
| return position.reshape(batch, self.token_groups * self.num_patches, self.decoder_dim) |
|
|
| def forward_decoder(self, latent, ids_restore, timestamps=None): |
| tokens = self.decoder_embed(latent) |
| mask_tokens = self.mask_token.expand( |
| tokens.shape[0], ids_restore.shape[1] + 1 - tokens.shape[1], -1 |
| ) |
| restored = torch.cat((tokens[:, 1:], mask_tokens), dim=1) |
| restored = torch.gather( |
| restored, 1, ids_restore.unsqueeze(-1).expand(-1, -1, self.decoder_dim) |
| ) |
| positions = self._decoder_positions( |
| tokens.shape[0], timestamps, tokens.dtype, tokens.device |
| ) |
| tokens = torch.cat((tokens[:, :1], restored + positions), dim=1) |
| for block in self.decoder_blocks: |
| tokens = block(tokens) |
| decoded = self.decoder_norm(tokens)[:, 1:] |
|
|
| if self.mode == "temporal": |
| return [self.decoder_pred[0](decoded)] |
| decoded = decoded.reshape( |
| decoded.shape[0], self.token_groups, self.num_patches, self.decoder_dim |
| ) |
| return [head(decoded[:, index]) for index, head in enumerate(self.decoder_pred)] |
|
|
| def _targets(self, images): |
| if self.mode == "temporal": |
| return [torch.cat( |
| [self.patchify(images[:, frame]) for frame in range(self.frames)], dim=1 |
| )] |
| return [self.patchify(images[:, group]) for group in self.spectral_groups] |
|
|
| def forward_loss(self, targets, predictions, mask): |
| losses = [] |
| if self.mode == "temporal": |
| pairs = [(targets[0], predictions[0], mask)] |
| else: |
| group_mask = mask.reshape(mask.shape[0], self.token_groups, self.num_patches) |
| pairs = [ |
| (target, prediction, group_mask[:, index]) |
| for index, (target, prediction) in enumerate(zip(targets, predictions)) |
| ] |
| removed = mask.new_zeros(()) |
| total = mask.new_zeros(()) |
| for target, prediction, patch_mask in pairs: |
| patch_loss = (prediction - target).square().mean(dim=-1) |
| total = total + (patch_loss * patch_mask).sum() |
| removed = removed + patch_mask.sum() |
| losses.append(patch_loss) |
| return total / removed.clamp_min(1), losses |
|
|
| def _normalize_targets(self, targets): |
| if not self.norm_pix_loss: |
| return targets |
| normalized = [] |
| for target in targets: |
| mean = target.mean(dim=-1, keepdim=True) |
| variance = target.var(dim=-1, keepdim=True, unbiased=False) |
| normalized.append((target - mean) / torch.sqrt(variance + 1e-6)) |
| return normalized |
|
|
| def _padded_outputs(self, tensors): |
| if self.mode == "temporal": |
| return tensors[0] |
| width = max(tensor.shape[-1] for tensor in tensors) |
| padded = [] |
| for tensor in tensors: |
| if tensor.shape[-1] < width: |
| tensor = torch.nn.functional.pad(tensor, (0, width - tensor.shape[-1])) |
| padded.append(tensor) |
| return torch.cat(padded, dim=1) |
|
|
| def forward(self, images, timestamps=None, mask_ratio=None): |
| latent, mask, ids_restore, timestamps = self.forward_encoder( |
| images, timestamps, mask_ratio |
| ) |
| predictions = self.forward_decoder(latent, ids_restore, timestamps) |
| targets = self._normalize_targets(self._targets(images)) |
| loss, patch_losses = self.forward_loss(targets, predictions, mask) |
| return { |
| "loss": loss, |
| "prediction": self._padded_outputs(predictions), |
| "target": self._padded_outputs(targets), |
| "mask": mask.bool(), |
| "features": latent, |
| "ids_restore": ids_restore, |
| "group_predictions": predictions, |
| "group_targets": targets, |
| "patch_losses": patch_losses, |
| } |
|
|
|
|
| def satmae_vit_base_patch16(**kwargs): |
| return SatMAE( |
| patch_size=16, embed_dim=768, encoder_depth=12, encoder_heads=12, |
| decoder_dim=512, decoder_depth=8, decoder_heads=16, |
| temporal_embed_dim=128, decoder_temporal_embed_dim=64, |
| channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs |
| ) |
|
|
|
|
| def satmae_vit_large_patch16(**kwargs): |
| return SatMAE( |
| patch_size=16, embed_dim=1024, encoder_depth=24, encoder_heads=16, |
| decoder_dim=512, decoder_depth=8, decoder_heads=16, |
| temporal_embed_dim=128, decoder_temporal_embed_dim=64, |
| channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs |
| ) |
|
|
|
|
| def satmae_vit_huge_patch14(**kwargs): |
| return SatMAE( |
| patch_size=14, embed_dim=1280, encoder_depth=32, encoder_heads=16, |
| decoder_dim=512, decoder_depth=8, decoder_heads=16, |
| temporal_embed_dim=128, decoder_temporal_embed_dim=64, |
| channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs |
| ) |
|
|
