# Copyright 2024 SatMAE++ Authors and The HuggingFace Inc. team. """Self-contained SatMAE++ model and config for trust_remote_code loading.""" from __future__ import annotations from functools import partial from typing import Optional import numpy as np import torch from timm.models.vision_transformer import Block, PatchEmbed from torch import nn from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput from transformers.modeling_utils import PreTrainedModel from transformers.processing_utils import Unpack from transformers.utils import TransformersKwargs, logging logger = logging.get_logger(__name__) FMOW_RGB_MEAN = [0.4182007312774658, 0.4214799106121063, 0.3991275727748871] FMOW_RGB_STD = [0.28774282336235046, 0.27541765570640564, 0.2764017581939697] FMOW_SENTINEL_MEAN_10 = [ 1184.3824625, 1120.77120066, 1136.26026392, 1263.73947144, 1645.40315151, 1846.87040806, 1762.59530783, 1972.62420416, 1732.16362238, 1247.91870117, ] FMOW_SENTINEL_STD_10 = [ 650.2842772, 712.12507725, 965.23119807, 948.9819932, 1108.06650639, 1258.36394548, 1233.1492281, 1364.38688993, 1310.36996126, 1087.6020813, ] DEFAULT_CHANNEL_GROUPS = [[0, 1, 2, 6], [3, 4, 5, 7], [8, 9]] def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray: grid_h = np.arange(grid_size, dtype=np.float32) grid_w = np.arange(grid_size, dtype=np.float32) grid = np.meshgrid(grid_w, grid_h) grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size]) pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) if cls_token: pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) return pos_embed def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray: emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) return np.concatenate([emb_h, emb_w], axis=1) def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray: omega = np.arange(embed_dim // 2, dtype=np.float32) omega /= embed_dim / 2.0 omega = 1.0 / 10000**omega pos = pos.reshape(-1) out = np.einsum("m,d->md", pos, omega) return np.concatenate([np.sin(out), np.cos(out)], axis=1) class SatMAEppConfig(PreTrainedConfig): model_type = "satmae_pp" def __init__( self, hidden_size: int = 1024, num_hidden_layers: int = 24, num_attention_heads: int = 16, intermediate_size: int | None = None, hidden_act: str = "gelu", hidden_dropout_prob: float = 0.0, attention_probs_dropout_prob: float = 0.0, initializer_range: float = 0.02, layer_norm_eps: float = 1e-6, image_size: int = 224, patch_size: int = 16, num_channels: int = 3, qkv_bias: bool = True, mlp_ratio: float = 4.0, global_pool: bool = True, encoder_type: str = "vanilla", channel_embed_dim: int = 256, channel_groups: list[list[int]] | None = None, channel_order: str = "bgr", dataset: str = "fmow_rgb", checkpoint_stage: str = "finetune", image_mean: list[float] | None = None, image_std: list[float] | None = None, num_labels: int = 0, **kwargs, ): super().__init__(**kwargs) self.hidden_size = hidden_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.hidden_act = hidden_act self.hidden_dropout_prob = hidden_dropout_prob self.attention_probs_dropout_prob = attention_probs_dropout_prob self.initializer_range = initializer_range self.layer_norm_eps = layer_norm_eps self.image_size = image_size self.patch_size = patch_size self.num_channels = num_channels self.qkv_bias = qkv_bias self.mlp_ratio = mlp_ratio self.global_pool = global_pool self.encoder_type = encoder_type self.channel_embed_dim = channel_embed_dim self.channel_groups = channel_groups if channel_groups is not None else list(DEFAULT_CHANNEL_GROUPS) self.channel_order = channel_order self.dataset = dataset self.checkpoint_stage = checkpoint_stage self.num_labels = num_labels self.intermediate_size = int(hidden_size * mlp_ratio) if intermediate_size is None else intermediate_size if image_mean is None or image_std is None: if encoder_type == "group_channel": self.image_mean = FMOW_SENTINEL_MEAN_10 self.image_std = FMOW_SENTINEL_STD_10 else: self.image_mean = FMOW_RGB_MEAN self.image_std = FMOW_RGB_STD else: self.image_mean = image_mean self.image_std = image_std class SatMAEppPreTrainedModel(PreTrainedModel): config_class = SatMAEppConfig config: SatMAEppConfig base_model_prefix = "satmae_pp" main_input_name = "pixel_values" input_modalities = ("image",) supports_gradient_checkpointing = True _no_split_modules = ["Block"] class SatMAEppModel(SatMAEppPreTrainedModel): def __init__(self, config: SatMAEppConfig, add_pooling_layer: bool = True): super().__init__(config) self.config = config self.add_pooling_layer = add_pooling_layer if config.encoder_type == "group_channel": self._init_group_channel_encoder(config) else: self._init_vanilla_encoder(config) self.post_init() def _init_vanilla_encoder(self, config: SatMAEppConfig) -> None: image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0] norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) self.patch_embed = PatchEmbed(image_size, config.patch_size, config.num_channels, config.hidden_size) self.num_patches = self.patch_embed.num_patches self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.hidden_size)) pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=True) self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) self.blocks = nn.ModuleList([ Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer) for _ in range(config.num_hidden_layers) ]) self.global_pool = config.global_pool if self.global_pool: self.fc_norm = norm_layer(config.hidden_size) self.norm = None else: self.fc_norm = None self.norm = norm_layer(config.hidden_size) def _init_group_channel_encoder(self, config: SatMAEppConfig) -> None: image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0] norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) self.channel_groups = tuple(tuple(group) for group in config.channel_groups) self.patch_embed = nn.ModuleList([ PatchEmbed(image_size, config.patch_size, len(group), config.hidden_size) for group in self.channel_groups ]) self.num_patches = self.patch_embed[0].num_patches pos_dim = config.hidden_size - config.channel_embed_dim self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, pos_dim)) pos_embed = get_2d_sincos_pos_embed(pos_dim, int(self.num_patches**0.5), cls_token=True) self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0)) num_groups = len(self.channel_groups) self.channel_embed = nn.Parameter(torch.zeros(1, num_groups, config.channel_embed_dim)) chan_embed = get_1d_sincos_pos_embed_from_grid(self.channel_embed.shape[-1], np.arange(num_groups)) self.channel_embed.data.copy_(torch.from_numpy(chan_embed).float().unsqueeze(0)) self.channel_cls_embed = nn.Parameter(torch.zeros(1, 1, config.channel_embed_dim)) self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) self.blocks = nn.ModuleList([ Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer) for _ in range(config.num_hidden_layers) ]) self.global_pool = config.global_pool if self.global_pool: self.fc_norm = norm_layer(config.hidden_size) self.norm = None else: self.fc_norm = None self.norm = norm_layer(config.hidden_size) def _forward_vanilla(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: batch_size = pixel_values.shape[0] patch_tokens = self.patch_embed(pixel_values) cls_tokens = self.cls_token.expand(batch_size, -1, -1) hidden_states = torch.cat((cls_tokens, patch_tokens), dim=1) + self.pos_embed for block in self.blocks: hidden_states = block(hidden_states) if self.global_pool: pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1)) else: hidden_states = self.norm(hidden_states) pooled_output = hidden_states[:, 0] return hidden_states, pooled_output def _forward_group_channel(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: batch_size = pixel_values.shape[0] group_tokens = [self.patch_embed[i](pixel_values[:, group, :, :]) for i, group in enumerate(self.channel_groups)] hidden_states = torch.stack(group_tokens, dim=1) channel_embed = self.channel_embed.unsqueeze(2).expand(-1, -1, self.pos_embed[:, 1:, :].shape[1], -1) pos_embed = self.pos_embed[:, 1:, :].unsqueeze(1).expand(-1, channel_embed.shape[1], -1, -1) hidden_states = (hidden_states + torch.cat((pos_embed, channel_embed), dim=-1)).view(batch_size, -1, hidden_states.shape[-1]) cls_pos_channel = torch.cat((self.pos_embed[:, :1, :], self.channel_cls_embed), dim=-1) hidden_states = torch.cat((cls_pos_channel + self.cls_token.expand(batch_size, -1, -1), hidden_states), dim=1) for block in self.blocks: hidden_states = block(hidden_states) if self.global_pool: pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1)) else: hidden_states = self.norm(hidden_states) pooled_output = hidden_states[:, 0] return hidden_states, pooled_output def forward( self, pixel_values: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: if pixel_values is None: raise ValueError("You must specify `pixel_values`") pixel_values = pixel_values.to(dtype=self.dtype) if return_dict is None: return_dict = self.config.use_return_dict if self.config.encoder_type == "group_channel": last_hidden_state, pooled_output = self._forward_group_channel(pixel_values) else: last_hidden_state, pooled_output = self._forward_vanilla(pixel_values) if not self.add_pooling_layer: pooled_output = None if not return_dict: return (last_hidden_state, pooled_output) return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooled_output) class SatMAEppForImageClassification(SatMAEppPreTrainedModel): def __init__(self, config: SatMAEppConfig): super().__init__(config) self.satmae_pp = SatMAEppModel(config, add_pooling_layer=True) self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity() self.post_init() def forward( self, pixel_values: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, return_dict: Optional[bool] = None, **kwargs: Unpack[TransformersKwargs], ) -> ImageClassifierOutput: outputs = self.satmae_pp(pixel_values=pixel_values, return_dict=True, **kwargs) logits = self.classifier(outputs.pooler_output) loss = None if labels is not None: loss = self.loss_function(labels, logits, self.config, **kwargs) if not return_dict: output = (logits,) + outputs[1:] return ((loss,) + output) if loss is not None else output return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions) __all__ = [ "SatMAEppConfig", "SatMAEppForImageClassification", "SatMAEppModel", "SatMAEppPreTrainedModel", ]