| """SatMAE++ masked autoencoder with convolutional multiscale decoding. |
| |
| Adapted for a self-contained package from the Apache-2.0 SatMAE++ reference |
| implementation at commit bf02548ab2bf5123761cf059491ac5a631fbb428. |
| """ |
|
|
| import math |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def sincos_1d(values, dim): |
| half = (dim + 1) // 2 |
| omega = torch.exp( |
| -math.log(10000.0) * torch.arange(half, dtype=torch.float32) |
| / max(half - 1, 1) |
| ) |
| phase = values.float().unsqueeze(-1) * omega |
| return torch.cat((phase.sin(), phase.cos()), dim=-1)[..., :dim] |
|
|
|
|
| def sincos_2d(side, dim): |
| y, x = torch.meshgrid( |
| torch.arange(side, dtype=torch.float32), |
| torch.arange(side, dtype=torch.float32), indexing="ij" |
| ) |
| split = dim // 2 |
| return torch.cat((sincos_1d(y.flatten(), split), |
| sincos_1d(x.flatten(), dim - split)), dim=-1) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, dim, heads, mlp_ratio=4.0): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(dim, eps=1e-6) |
| self.attn = nn.MultiheadAttention(dim, heads, batch_first=True) |
| self.norm2 = nn.LayerNorm(dim, eps=1e-6) |
| hidden = int(dim * mlp_ratio) |
| self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim)) |
|
|
| def forward(self, x): |
| value = self.norm1(x) |
| x = x + self.attn(value, value, value, need_weights=False)[0] |
| return x + self.mlp(self.norm2(x)) |
|
|
|
|
| class ChannelFirstNorm(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(channels)) |
| self.bias = nn.Parameter(torch.zeros(channels)) |
|
|
| def forward(self, x): |
| mean = x.mean(1, keepdim=True) |
| variance = (x - mean).square().mean(1, keepdim=True) |
| x = (x - mean) / torch.sqrt(variance + 1e-6) |
| return x * self.weight[:, None, None] + self.bias[:, None, None] |
|
|
|
|
| class ResidualBlock(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.conv1 = nn.Conv2d(channels, channels, 3, padding=1) |
| self.conv2 = nn.Conv2d(channels, channels, 3, padding=1) |
|
|
| def forward(self, x): |
| return x + 0.5 * self.conv2(F.relu(self.conv1(x))) |
|
|
|
|
| class UpsampleBlock(nn.Module): |
| def __init__(self, hidden_channels, output_channels): |
| super().__init__() |
| self.up = nn.ConvTranspose2d(hidden_channels, hidden_channels, 4, 2, 1) |
| self.up_norm = ChannelFirstNorm(hidden_channels) |
| self.residual = ResidualBlock(hidden_channels) |
| self.residual_norm = ChannelFirstNorm(hidden_channels) |
| self.output = nn.Conv2d(hidden_channels, output_channels, 3, padding=1) |
|
|
| def forward(self, x): |
| hidden = F.leaky_relu(self.up_norm(self.up(x))) |
| hidden = self.residual_norm(self.residual(hidden)) |
| return hidden, self.output(hidden) |
|
|
|
|
| class SatMAEPP(nn.Module): |
| def __init__(self, image_size=224, patch_size=16, in_channels=3, |
| embed_dim=1024, encoder_depth=24, encoder_heads=16, |
| decoder_dim=512, decoder_depth=8, decoder_heads=16, |
| mask_ratio=0.75, scales=None, mode="rgb", spectral_groups=None, |
| spatial_mask=False, norm_pix_loss=False, proj_ratio=4, |
| channel_embed_dim=None, decoder_channel_embed_dim=None): |
| super().__init__() |
| if image_size <= 0 or patch_size <= 0 or image_size % patch_size: |
| raise ValueError("image_size must be divisible by patch_size") |
| if in_channels <= 0 or embed_dim <= 0 or decoder_dim <= 0: |
| raise ValueError("channel and embedding dimensions must be positive") |
| if encoder_depth <= 0 or decoder_depth <= 0 or encoder_heads <= 0 or decoder_heads <= 0: |
| raise ValueError("transformer depths and head counts must be positive") |
| if embed_dim % encoder_heads or decoder_dim % decoder_heads: |
| raise ValueError("embedding dimensions must be divisible by head counts") |
| if not 0 <= mask_ratio < 1: |
| raise ValueError("mask_ratio must be in [0, 1)") |
| if mode not in {"rgb", "multispectral"}: |
| raise ValueError("mode must be rgb or multispectral") |
| self.image_size = image_size |
| self.patch_size = patch_size |
| self.in_channels = in_channels |
| self.mask_ratio = mask_ratio |
| self.mode = mode |
| self.spatial_mask = spatial_mask |
| self.norm_pix_loss = norm_pix_loss |
| self.grid = image_size // patch_size |
| self.num_patches = self.grid ** 2 |
| self.scales = tuple(scales or ([1, 2] if mode == "rgb" else [1, 2, 4])) |
| expected_scales = (1, 2) if mode == "rgb" else (1, 2, 4) |
| if self.scales != expected_scales: |
| raise ValueError("supported scales are [1, 2] or [1, 2, 4]") |
|
|
| if mode == "rgb": |
| self.groups = (tuple(range(in_channels)),) |
| channel_embed_dim = 0 |
| decoder_channel_embed_dim = 0 |
| else: |
| groups = spectral_groups or [[0, 1, 2, 6], [3, 4, 5, 7], [8, 9]] |
| if sorted(channel for group in groups for channel in group) != list(range(in_channels)): |
| raise ValueError("spectral_groups must partition all channels") |
| self.groups = tuple(tuple(group) for group in groups) |
| channel_embed_dim = channel_embed_dim or min(256, embed_dim // 4) |
| decoder_channel_embed_dim = decoder_channel_embed_dim or min(128, decoder_dim // 4) |
| self.group_count = len(self.groups) |
| self.channel_embed_dim = channel_embed_dim |
| self.decoder_channel_embed_dim = decoder_channel_embed_dim |
|
|
| self.patch_embeds = nn.ModuleList([ |
| nn.Conv2d(len(group), embed_dim, patch_size, patch_size) for group in self.groups |
| ]) |
| spatial_dim = embed_dim - channel_embed_dim |
| decoder_spatial_dim = decoder_dim - decoder_channel_embed_dim |
| self.register_buffer("position", sincos_2d(self.grid, spatial_dim)) |
| self.register_buffer("decoder_position", sincos_2d(self.grid, decoder_spatial_dim)) |
| if self.group_count > 1: |
| ids = torch.arange(self.group_count, dtype=torch.float32) |
| self.register_buffer("group_position", sincos_1d(ids, channel_embed_dim)) |
| self.register_buffer("decoder_group_position", sincos_1d(ids, decoder_channel_embed_dim)) |
|
|
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.blocks = nn.ModuleList([ |
| TransformerBlock(embed_dim, encoder_heads) for _ in range(encoder_depth) |
| ]) |
| self.norm = nn.LayerNorm(embed_dim, eps=1e-6) |
| self.decoder_embed = nn.Linear(embed_dim, decoder_dim) |
| self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim)) |
| self.decoder_blocks = nn.ModuleList([ |
| TransformerBlock(decoder_dim, decoder_heads) for _ in range(decoder_depth) |
| ]) |
| self.decoder_norm = nn.LayerNorm(decoder_dim, eps=1e-6) |
| self.decoder_heads = nn.ModuleList([ |
| nn.Linear(decoder_dim, len(group) * patch_size ** 2) for group in self.groups |
| ]) |
|
|
| hidden_channels = in_channels * proj_ratio |
| self.multiscale_projection = nn.Conv2d(in_channels, hidden_channels, 1) |
| self.multiscale_norm = ChannelFirstNorm(hidden_channels) |
| self.upsample_blocks = nn.ModuleList([ |
| UpsampleBlock(hidden_channels, in_channels) for _ in self.scales[1:] |
| ]) |
| self.initialize_weights() |
|
|
| def initialize_weights(self): |
| 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.Conv2d, nn.ConvTranspose2d)): |
| nn.init.xavier_uniform_(module.weight.flatten(1)) |
| 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): |
| batch, channels, height, width = images.shape |
| p = self.patch_size |
| if channels != self.in_channels or height != self.image_size or width != self.image_size: |
| raise ValueError("images must have shape [B, in_channels, image_size, image_size]") |
| x = images.reshape(batch, channels, height // p, p, width // p, p) |
| return x.permute(0, 2, 4, 1, 3, 5).reshape(batch, -1, channels * p ** 2) |
|
|
| def unpatchify(self, patches): |
| batch = patches.shape[0] |
| p = self.patch_size |
| x = patches.reshape(batch, self.grid, self.grid, self.in_channels, p, p) |
| return x.permute(0, 3, 1, 4, 2, 5).reshape( |
| batch, self.in_channels, self.image_size, self.image_size |
| ) |
|
|
| def _positions(self, decoder=False): |
| spatial = self.decoder_position if decoder else self.position |
| if self.group_count == 1: |
| return spatial.unsqueeze(0) |
| group = self.decoder_group_position if decoder else self.group_position |
| value = torch.cat(( |
| spatial.unsqueeze(0).expand(self.group_count, -1, -1), |
| group.unsqueeze(1).expand(-1, self.num_patches, -1), |
| ), dim=-1) |
| return value.reshape(1, self.group_count * self.num_patches, -1) |
|
|
| def _mask(self, tokens, ratio): |
| batch, length, dim = tokens.shape |
| shared = self.group_count > 1 and self.spatial_mask |
| if shared: |
| kept_spatial = int(self.num_patches * (1 - ratio)) |
| order = torch.rand(batch, self.num_patches, device=tokens.device).argsort(1) |
| kept = [order[:, :kept_spatial] + index * self.num_patches |
| for index in range(self.group_count)] |
| removed = [order[:, kept_spatial:] + index * self.num_patches |
| for index in range(self.group_count)] |
| shuffle = torch.cat(kept + removed, dim=1) |
| keep = kept_spatial * self.group_count |
| else: |
| keep = int(length * (1 - ratio)) |
| shuffle = torch.rand(batch, length, device=tokens.device).argsort(1) |
| restore = shuffle.argsort(1) |
| visible = torch.gather(tokens, 1, shuffle[:, :keep, None].expand(-1, -1, dim)) |
| mask = torch.ones(batch, length, device=tokens.device) |
| mask[:, :keep] = 0 |
| return visible, torch.gather(mask, 1, restore), restore |
|
|
| def forward_encoder(self, images, ratio): |
| pieces = [embed(images[:, group]).flatten(2).transpose(1, 2) |
| for embed, group in zip(self.patch_embeds, self.groups)] |
| tokens = torch.cat(pieces, dim=1) + self._positions(False) |
| tokens, mask, restore = self._mask(tokens, ratio) |
| tokens = torch.cat((self.cls_token.expand(images.shape[0], -1, -1), tokens), 1) |
| for block in self.blocks: |
| tokens = block(tokens) |
| return self.norm(tokens), mask, restore |
|
|
| def forward_decoder(self, latent, restore): |
| tokens = self.decoder_embed(latent) |
| missing = restore.shape[1] + 1 - tokens.shape[1] |
| restored = torch.cat((tokens[:, 1:], self.mask_token.expand(tokens.shape[0], missing, -1)), 1) |
| restored = torch.gather(restored, 1, restore[:, :, None].expand(-1, -1, tokens.shape[-1])) |
| tokens = torch.cat((tokens[:, :1], restored + self._positions(True)), 1) |
| for block in self.decoder_blocks: |
| tokens = block(tokens) |
| decoded = self.decoder_norm(tokens)[:, 1:].reshape( |
| tokens.shape[0], self.group_count, self.num_patches, -1 |
| ) |
| group_predictions = [head(decoded[:, index]) for index, head in enumerate(self.decoder_heads)] |
| patch_channels = [] |
| for prediction, group in zip(group_predictions, self.groups): |
| patch_channels.append(prediction.reshape( |
| prediction.shape[0], self.num_patches, len(group), self.patch_size ** 2 |
| )) |
| patches = torch.empty( |
| tokens.shape[0], self.num_patches, self.in_channels, self.patch_size ** 2, |
| device=tokens.device, dtype=group_predictions[0].dtype |
| ) |
| for values, group in zip(patch_channels, self.groups): |
| patches[:, :, list(group)] = values |
| return patches.flatten(2), group_predictions |
|
|
| def forward_multiscale(self, reconstruction): |
| hidden = self.multiscale_norm(F.gelu(self.multiscale_projection(reconstruction))) |
| predictions = {"1": reconstruction} |
| for scale, block in zip(self.scales[1:], self.upsample_blocks): |
| hidden, predictions[str(scale)] = block(hidden) |
| return predictions |
|
|
| def forward(self, images, high_resolution_targets=None, mask_ratio=None): |
| ratio = self.mask_ratio if mask_ratio is None else mask_ratio |
| if images.ndim != 4 or images.shape[1:] != (self.in_channels, self.image_size, self.image_size): |
| raise ValueError("images must have shape [B, in_channels, image_size, image_size]") |
| if not 0 <= ratio < 1: |
| raise ValueError("mask_ratio must be in [0, 1)") |
| high_resolution_targets = high_resolution_targets or {} |
| latent, mask, restore = self.forward_encoder(images, ratio) |
| patch_prediction, group_predictions = self.forward_decoder(latent, restore) |
| target_patches = self.patchify(images) |
| loss_target = target_patches |
| if self.norm_pix_loss: |
| mean = loss_target.mean(-1, keepdim=True) |
| variance = loss_target.var(-1, keepdim=True, unbiased=False) |
| loss_target = (loss_target - mean) / torch.sqrt(variance + 1e-6) |
| grouped_mask = mask.reshape(mask.shape[0], self.group_count, self.num_patches) |
| prediction_channels = patch_prediction.reshape( |
| patch_prediction.shape[0], self.num_patches, |
| self.in_channels, self.patch_size ** 2 |
| ) |
| target_channels = loss_target.reshape( |
| loss_target.shape[0], self.num_patches, |
| self.in_channels, self.patch_size ** 2 |
| ) |
| base_total = mask.new_zeros(()) |
| removed = mask.new_zeros(()) |
| for index, group in enumerate(self.groups): |
| group_error = ( |
| prediction_channels[:, :, list(group)] |
| - target_channels[:, :, list(group)] |
| ).square().mean(dim=(-1, -2)) |
| base_total = base_total + (group_error * grouped_mask[:, index]).sum() |
| removed = removed + grouped_mask[:, index].sum() |
| base_mse = base_total / removed.clamp_min(1) |
| base_l1_total = mask.new_zeros(()) |
| for index, group in enumerate(self.groups): |
| group_l1 = (prediction_channels[:, :, list(group)] - |
| target_channels[:, :, list(group)]).abs().mean(dim=(-1, -2)) |
| base_l1_total = base_l1_total + (group_l1 * grouped_mask[:, index]).sum() |
| base_l1 = base_l1_total / removed.clamp_min(1) |
| base_loss = base_mse + base_l1 |
| reconstruction = self.unpatchify(patch_prediction) |
| predictions = self.forward_multiscale(reconstruction) |
| targets = {"1": images} |
| multiscale_losses = {} |
| for scale in self.scales[1:]: |
| key = str(scale) |
| target = None if high_resolution_targets is None else high_resolution_targets.get(key) |
| if target is None: |
| raise ValueError(f"native target for scale x{scale} is required") |
| expected_shape = (images.shape[0], self.in_channels, |
| self.image_size * scale, self.image_size * scale) |
| if tuple(target.shape) != expected_shape: |
| raise ValueError(f"target x{scale} must have shape {expected_shape}") |
| targets[key] = target |
| multiscale_losses[key] = F.mse_loss(predictions[key], target) + F.l1_loss( |
| predictions[key], target |
| ) |
| multiscale_loss = (sum(multiscale_losses.values()) / len(multiscale_losses) |
| if multiscale_losses else base_loss.new_zeros(())) |
| return { |
| "loss": base_loss + multiscale_loss, |
| "reconstruction_loss": base_loss, |
| "multiscale_loss": multiscale_loss, |
| "reconstruction": reconstruction, |
| "patch_prediction": patch_prediction, |
| "target_patches": loss_target, |
| "predictions": predictions, |
| "targets": targets, |
| "mask": mask.bool(), |
| "features": latent, |
| "ids_restore": restore, |
| "group_predictions": group_predictions, |
| "scale_losses": multiscale_losses, |
| } |
|
|
|
|
| def satmaepp_vit_base(**kwargs): |
| return SatMAEPP(embed_dim=768, encoder_depth=12, encoder_heads=12, |
| decoder_dim=512, decoder_depth=8, decoder_heads=16, **kwargs) |
|
|
|
|
| def satmaepp_vit_large(**kwargs): |
| return SatMAEPP(embed_dim=1024, encoder_depth=24, encoder_heads=16, |
| decoder_dim=512, decoder_depth=8, decoder_heads=16, **kwargs) |
|
|
