satmae-pp-vit-large-patch8-fmow-sentinel-pretrain/modeling_satmae_pp.py

# 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",
]