OneScience-Group/SatMAE

1"""Paper-aligned SatMAE model components.2 3This is an original implementation of the architecture described in SatMAE.4The upstream repository was used only as a behavioral reference; no upstream5source text is incorporated here.6"""7 8import math9from functools import partial10 11import torch12from torch import nn13 14 15def _sincos_1d(values, dim):16    """Return a fixed sine-cosine embedding for arbitrary scalar positions."""17    if dim <= 0:18        return values.new_zeros((*values.shape, 0))19    pairs = (dim + 1) // 220    omega = torch.arange(pairs, device=values.device, dtype=torch.float32)21    omega = torch.exp(-math.log(10000.0) * omega / max(pairs - 1, 1))22    phase = values.to(torch.float32).unsqueeze(-1) * omega23    return torch.cat((phase.sin(), phase.cos()), dim=-1)[..., :dim]24 25 26def _sincos_2d(grid_size, dim):27    """Return a fixed row-major 2D sine-cosine position embedding."""28    rows, cols = torch.meshgrid(29        torch.arange(grid_size, dtype=torch.float32),30        torch.arange(grid_size, dtype=torch.float32),31        indexing="ij",32    )33    row_dim = dim // 234    return torch.cat(35        (_sincos_1d(rows.reshape(-1), row_dim),36         _sincos_1d(cols.reshape(-1), dim - row_dim)),37        dim=-1,38    )39 40 41def _timestamp_embedding(timestamps, dim):42    """Encode either scalar times or fMoW ``[year, month, hour]`` tuples."""43    if timestamps.ndim == 2:44        return _sincos_1d(timestamps, dim)45    if timestamps.ndim != 3 or timestamps.shape[-1] != 3:46        raise ValueError("timestamps must have shape [B, T] or [B, T, 3]")47    field_dims = [dim // 3] * 348    for index in range(dim % 3):49        field_dims[index] += 150    return torch.cat(51        [_sincos_1d(timestamps[..., index], field_dim)52         for index, field_dim in enumerate(field_dims)],53        dim=-1,54    )55 56 57class PatchEmbed(nn.Module):58    def __init__(self, image_size, patch_size, in_channels, embed_dim):59        super().__init__()60        self.image_size = image_size61        self.patch_size = patch_size62        self.num_patches = (image_size // patch_size) ** 263        self.proj = nn.Conv2d(64            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size65        )66 67    def forward(self, images):68        if images.shape[-2:] != (self.image_size, self.image_size):69            raise ValueError(70                f"expected {self.image_size}x{self.image_size} images, "71                f"got {tuple(images.shape[-2:])}"72            )73        return self.proj(images).flatten(2).transpose(1, 2)74 75 76class TransformerBlock(nn.Module):77    def __init__(self, dim, num_heads, mlp_ratio=4.0, norm_layer=nn.LayerNorm):78        super().__init__()79        self.norm1 = norm_layer(dim)80        self.attention = nn.MultiheadAttention(81            dim, num_heads, dropout=0.0, bias=True, batch_first=True82        )83        self.norm2 = norm_layer(dim)84        hidden_dim = int(dim * mlp_ratio)85        self.mlp = nn.Sequential(86            nn.Linear(dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, dim)87        )88 89    def forward(self, tokens):90        normalized = self.norm1(tokens)91        tokens = tokens + self.attention(92            normalized, normalized, normalized, need_weights=False93        )[0]94        return tokens + self.mlp(self.norm2(tokens))95 96 97class SatMAE(nn.Module):98    """Masked autoencoder for temporal or grouped multispectral imagery.99 100    Temporal inputs use shape ``[B, T, C, H, W]`` and optional timestamps101    ``[B, T]``. Multispectral inputs use shape ``[B, C, H, W]``.102    """103 104    def __init__(105        self,106        image_size=224,107        patch_size=16,108        in_channels=3,109        frames=3,110        embed_dim=1024,111        encoder_depth=24,112        encoder_heads=16,113        decoder_dim=512,114        decoder_depth=8,115        decoder_heads=16,116        mlp_ratio=4.0,117        mode="temporal",118        spectral_groups=None,119        mask_ratio=0.75,120        norm_pix_loss=False,121        same_mask=False,122        spatial_mask=False,123        temporal_embed_dim=None,124        decoder_temporal_embed_dim=None,125        channel_embed_dim=None,126        decoder_channel_embed_dim=None,127        norm_layer=None,128    ):129        super().__init__()130        if image_size % patch_size:131            raise ValueError("image_size must be divisible by patch_size")132        if not 0.0 <= mask_ratio < 1.0:133            raise ValueError("mask_ratio must be in [0, 1)")134        if mode not in {"temporal", "multispectral"}:135            raise ValueError("mode must be temporal or multispectral")136        if embed_dim % encoder_heads or decoder_dim % decoder_heads:137            raise ValueError("embedding dimensions must be divisible by head counts")138 139        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)140        self.image_size = image_size141        self.patch_size = patch_size142        self.in_channels = in_channels143        self.frames = frames144        self.embed_dim = embed_dim145        self.decoder_dim = decoder_dim146        self.mode = mode147        self.mask_ratio = mask_ratio148        self.norm_pix_loss = norm_pix_loss149        self.same_mask = same_mask150        self.spatial_mask = spatial_mask151        self.grid_size = image_size // patch_size152        self.num_patches = self.grid_size ** 2153 154        if mode == "temporal":155            self.spectral_groups = None156            self.patch_embed = PatchEmbed(157                image_size, patch_size, in_channels, embed_dim158            )159            self.token_groups = frames160            semantic_dim = temporal_embed_dim161            if semantic_dim is None:162                semantic_dim = min(128, max(2, embed_dim // 4))163            decoder_semantic_dim = decoder_temporal_embed_dim164            if decoder_semantic_dim is None:165                decoder_semantic_dim = min(64, max(2, decoder_dim // 4))166            prediction_dims = [patch_size ** 2 * in_channels]167        else:168            groups = spectral_groups or [list(range(in_channels))]169            flattened = [channel for group in groups for channel in group]170            if sorted(flattened) != list(range(in_channels)):171                raise ValueError("spectral_groups must partition all input channels")172            self.spectral_groups = tuple(tuple(group) for group in groups)173            self.patch_embed = nn.ModuleList(174                PatchEmbed(image_size, patch_size, len(group), embed_dim)175                for group in self.spectral_groups176            )177            self.token_groups = len(self.spectral_groups)178            semantic_dim = channel_embed_dim179            if semantic_dim is None:180                semantic_dim = min(256, max(2, embed_dim // 4))181            decoder_semantic_dim = decoder_channel_embed_dim182            if decoder_semantic_dim is None:183                decoder_semantic_dim = min(128, max(2, decoder_dim // 4))184            prediction_dims = [patch_size ** 2 * len(g) for g in self.spectral_groups]185 186        if not 0 < semantic_dim < embed_dim:187            raise ValueError("encoder semantic embedding dimension is invalid")188        if not 0 < decoder_semantic_dim < decoder_dim:189            raise ValueError("decoder semantic embedding dimension is invalid")190        self.semantic_dim = semantic_dim191        self.decoder_semantic_dim = decoder_semantic_dim192 193        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))194        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))195        self.register_buffer(196            "spatial_pos_embed",197            _sincos_2d(self.grid_size, embed_dim - semantic_dim),198            persistent=True,199        )200        self.register_buffer(201            "decoder_spatial_pos_embed",202            _sincos_2d(self.grid_size, decoder_dim - decoder_semantic_dim),203            persistent=True,204        )205        if mode == "multispectral":206            group_ids = torch.arange(self.token_groups, dtype=torch.float32)207            self.register_buffer(208                "group_embed", _sincos_1d(group_ids, semantic_dim), persistent=True209            )210            self.register_buffer(211                "decoder_group_embed",212                _sincos_1d(group_ids, decoder_semantic_dim),213                persistent=True,214            )215 216        self.blocks = nn.ModuleList(217            TransformerBlock(embed_dim, encoder_heads, mlp_ratio, norm_layer)218            for _ in range(encoder_depth)219        )220        self.norm = norm_layer(embed_dim)221        self.decoder_embed = nn.Linear(embed_dim, decoder_dim)222        self.decoder_blocks = nn.ModuleList(223            TransformerBlock(decoder_dim, decoder_heads, mlp_ratio, norm_layer)224            for _ in range(decoder_depth)225        )226        self.decoder_norm = norm_layer(decoder_dim)227        self.decoder_pred = nn.ModuleList(228            nn.Linear(decoder_dim, output_dim) for output_dim in prediction_dims229        )230        self.initialize_weights()231 232    def initialize_weights(self):233        patch_embeds = (234            [self.patch_embed]235            if isinstance(self.patch_embed, PatchEmbed)236            else self.patch_embed237        )238        for patch_embed in patch_embeds:239            nn.init.xavier_uniform_(patch_embed.proj.weight.flatten(1))240            if patch_embed.proj.bias is not None:241                nn.init.zeros_(patch_embed.proj.bias)242        nn.init.normal_(self.cls_token, std=0.02)243        nn.init.normal_(self.mask_token, std=0.02)244        for module in self.modules():245            if isinstance(module, nn.Linear):246                nn.init.xavier_uniform_(module.weight)247                if module.bias is not None:248                    nn.init.zeros_(module.bias)249            elif isinstance(module, nn.LayerNorm):250                nn.init.ones_(module.weight)251                nn.init.zeros_(module.bias)252 253    def patchify(self, images):254        if images.ndim != 4:255            raise ValueError("patchify expects [B, C, H, W]")256        batch, channels, height, width = images.shape257        patch = self.patch_size258        if height != width or height != self.image_size:259            raise ValueError(f"expected square images of size {self.image_size}")260        patches = images.reshape(261            batch, channels, height // patch, patch, width // patch, patch262        )263        patches = patches.permute(0, 2, 4, 1, 3, 5)264        return patches.reshape(batch, self.num_patches, channels * patch ** 2)265 266    def unpatchify(self, patches, channels=None):267        channels = channels or self.in_channels268        batch = patches.shape[0]269        patch = self.patch_size270        expected = channels * patch ** 2271        if patches.shape[1:] != (self.num_patches, expected):272            raise ValueError("patch tensor has incompatible shape")273        images = patches.reshape(274            batch, self.grid_size, self.grid_size, channels, patch, patch275        )276        images = images.permute(0, 3, 1, 4, 2, 5)277        return images.reshape(batch, channels, self.image_size, self.image_size)278 279    def _random_masking(self, tokens, mask_ratio, share_spatial_mask):280        batch, length, dim = tokens.shape281        if share_spatial_mask:282            units = self.num_patches283            len_keep_units = int(units * (1.0 - mask_ratio))284            noise = torch.rand(batch, units, device=tokens.device)285            spatial_order = noise.argsort(dim=1)286            kept = [spatial_order[:, :len_keep_units] + g * units287                    for g in range(self.token_groups)]288            removed = [spatial_order[:, len_keep_units:] + g * units289                       for g in range(self.token_groups)]290            ids_shuffle = torch.cat(kept + removed, dim=1)291            len_keep = len_keep_units * self.token_groups292        else:293            len_keep = int(length * (1.0 - mask_ratio))294            ids_shuffle = torch.rand(batch, length, device=tokens.device).argsort(dim=1)295        ids_restore = ids_shuffle.argsort(dim=1)296        ids_keep = ids_shuffle[:, :len_keep]297        visible = torch.gather(tokens, 1, ids_keep.unsqueeze(-1).expand(-1, -1, dim))298        mask = torch.ones(batch, length, device=tokens.device)299        mask[:, :len_keep] = 0300        mask = torch.gather(mask, 1, ids_restore)301        return visible, mask, ids_restore302 303    def _temporal_tokens(self, images, timestamps):304        if images.ndim != 5:305            raise ValueError("temporal mode expects images shaped [B, T, C, H, W]")306        batch, frames, channels, _, _ = images.shape307        if frames != self.frames or channels != self.in_channels:308            raise ValueError(309                f"expected T={self.frames}, C={self.in_channels}; got T={frames}, C={channels}"310            )311        if timestamps is None:312            timestamps = torch.arange(frames, device=images.device).expand(batch, -1)313        if timestamps.shape[:2] != (batch, frames):314            raise ValueError(315                f"timestamps must start with shape {(batch, frames)}, "316                f"got {tuple(timestamps.shape)}"317            )318        spatial = self.spatial_pos_embed.to(dtype=images.dtype)319        time = _timestamp_embedding(timestamps, self.semantic_dim).to(dtype=images.dtype)320        position = torch.cat(321            (spatial.view(1, 1, self.num_patches, -1).expand(batch, frames, -1, -1),322             time.unsqueeze(2).expand(-1, -1, self.num_patches, -1)),323            dim=-1,324        ).reshape(batch, frames * self.num_patches, self.embed_dim)325        tokens = torch.stack(326            [self.patch_embed(images[:, frame]) for frame in range(frames)], dim=1327        ).reshape(batch, frames * self.num_patches, self.embed_dim)328        return tokens + position, timestamps329 330    def _multispectral_tokens(self, images):331        if images.ndim != 4 or images.shape[1] != self.in_channels:332            raise ValueError(333                f"multispectral mode expects images shaped [B, {self.in_channels}, H, W]"334            )335        spatial = self.spatial_pos_embed.to(dtype=images.dtype)336        group = self.group_embed.to(dtype=images.dtype)337        positions = torch.cat(338            (spatial.view(1, self.num_patches, -1).expand(self.token_groups, -1, -1),339             group.view(self.token_groups, 1, -1).expand(-1, self.num_patches, -1)),340            dim=-1,341        ).reshape(1, self.token_groups * self.num_patches, self.embed_dim)342        tokens = torch.cat(343            [embed(images[:, channels])344             for embed, channels in zip(self.patch_embed, self.spectral_groups)],345            dim=1,346        )347        return tokens + positions348 349    def forward_encoder(self, images, timestamps=None, mask_ratio=None):350        ratio = self.mask_ratio if mask_ratio is None else mask_ratio351        if not 0.0 <= ratio < 1.0:352            raise ValueError("mask_ratio must be in [0, 1)")353        if self.mode == "temporal":354            tokens, timestamps = self._temporal_tokens(images, timestamps)355            shared = self.same_mask356        else:357            tokens = self._multispectral_tokens(images)358            shared = self.spatial_mask359        tokens, mask, ids_restore = self._random_masking(tokens, ratio, shared)360        cls = self.cls_token.expand(tokens.shape[0], -1, -1)361        tokens = torch.cat((cls, tokens), dim=1)362        for block in self.blocks:363            tokens = block(tokens)364        return self.norm(tokens), mask, ids_restore, timestamps365 366    def _decoder_positions(self, batch, timestamps, dtype, device):367        spatial = self.decoder_spatial_pos_embed.to(device=device, dtype=dtype)368        if self.mode == "temporal":369            semantic = _timestamp_embedding(timestamps, self.decoder_semantic_dim).to(dtype=dtype)370        else:371            semantic = self.decoder_group_embed.to(device=device, dtype=dtype)372            semantic = semantic.unsqueeze(0).expand(batch, -1, -1)373        position = torch.cat(374            (spatial.view(1, 1, self.num_patches, -1).expand(batch, self.token_groups, -1, -1),375             semantic.unsqueeze(2).expand(-1, -1, self.num_patches, -1)),376            dim=-1,377        )378        return position.reshape(batch, self.token_groups * self.num_patches, self.decoder_dim)379 380    def forward_decoder(self, latent, ids_restore, timestamps=None):381        tokens = self.decoder_embed(latent)382        mask_tokens = self.mask_token.expand(383            tokens.shape[0], ids_restore.shape[1] + 1 - tokens.shape[1], -1384        )385        restored = torch.cat((tokens[:, 1:], mask_tokens), dim=1)386        restored = torch.gather(387            restored, 1, ids_restore.unsqueeze(-1).expand(-1, -1, self.decoder_dim)388        )389        positions = self._decoder_positions(390            tokens.shape[0], timestamps, tokens.dtype, tokens.device391        )392        tokens = torch.cat((tokens[:, :1], restored + positions), dim=1)393        for block in self.decoder_blocks:394            tokens = block(tokens)395        decoded = self.decoder_norm(tokens)[:, 1:]396 397        if self.mode == "temporal":398            return [self.decoder_pred[0](decoded)]399        decoded = decoded.reshape(400            decoded.shape[0], self.token_groups, self.num_patches, self.decoder_dim401        )402        return [head(decoded[:, index]) for index, head in enumerate(self.decoder_pred)]403 404    def _targets(self, images):405        if self.mode == "temporal":406            return [torch.cat(407                [self.patchify(images[:, frame]) for frame in range(self.frames)], dim=1408            )]409        return [self.patchify(images[:, group]) for group in self.spectral_groups]410 411    def forward_loss(self, targets, predictions, mask):412        losses = []413        if self.mode == "temporal":414            pairs = [(targets[0], predictions[0], mask)]415        else:416            group_mask = mask.reshape(mask.shape[0], self.token_groups, self.num_patches)417            pairs = [418                (target, prediction, group_mask[:, index])419                for index, (target, prediction) in enumerate(zip(targets, predictions))420            ]421        removed = mask.new_zeros(())422        total = mask.new_zeros(())423        for target, prediction, patch_mask in pairs:424            patch_loss = (prediction - target).square().mean(dim=-1)425            total = total + (patch_loss * patch_mask).sum()426            removed = removed + patch_mask.sum()427            losses.append(patch_loss)428        return total / removed.clamp_min(1), losses429 430    def _normalize_targets(self, targets):431        if not self.norm_pix_loss:432            return targets433        normalized = []434        for target in targets:435            mean = target.mean(dim=-1, keepdim=True)436            variance = target.var(dim=-1, keepdim=True, unbiased=False)437            normalized.append((target - mean) / torch.sqrt(variance + 1e-6))438        return normalized439 440    def _padded_outputs(self, tensors):441        if self.mode == "temporal":442            return tensors[0]443        width = max(tensor.shape[-1] for tensor in tensors)444        padded = []445        for tensor in tensors:446            if tensor.shape[-1] < width:447                tensor = torch.nn.functional.pad(tensor, (0, width - tensor.shape[-1]))448            padded.append(tensor)449        return torch.cat(padded, dim=1)450 451    def forward(self, images, timestamps=None, mask_ratio=None):452        latent, mask, ids_restore, timestamps = self.forward_encoder(453            images, timestamps, mask_ratio454        )455        predictions = self.forward_decoder(latent, ids_restore, timestamps)456        targets = self._normalize_targets(self._targets(images))457        loss, patch_losses = self.forward_loss(targets, predictions, mask)458        return {459            "loss": loss,460            "prediction": self._padded_outputs(predictions),461            "target": self._padded_outputs(targets),462            "mask": mask.bool(),463            "features": latent,464            "ids_restore": ids_restore,465            "group_predictions": predictions,466            "group_targets": targets,467            "patch_losses": patch_losses,468        }469 470 471def satmae_vit_base_patch16(**kwargs):472    return SatMAE(473        patch_size=16, embed_dim=768, encoder_depth=12, encoder_heads=12,474        decoder_dim=512, decoder_depth=8, decoder_heads=16,475        temporal_embed_dim=128, decoder_temporal_embed_dim=64,476        channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs477    )478 479 480def satmae_vit_large_patch16(**kwargs):481    return SatMAE(482        patch_size=16, embed_dim=1024, encoder_depth=24, encoder_heads=16,483        decoder_dim=512, decoder_depth=8, decoder_heads=16,484        temporal_embed_dim=128, decoder_temporal_embed_dim=64,485        channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs486    )487 488 489def satmae_vit_huge_patch14(**kwargs):490    return SatMAE(491        patch_size=14, embed_dim=1280, encoder_depth=32, encoder_heads=16,492        decoder_dim=512, decoder_depth=8, decoder_heads=16,493        temporal_embed_dim=128, decoder_temporal_embed_dim=64,494        channel_embed_dim=256, decoder_channel_embed_dim=128, **kwargs495    )496