# Copyright 2024 The SSL4EO-S12 Authors and The HuggingFace Inc. team. """Self-contained SSL4EO ViT model and config for trust_remote_code loading.""" from functools import partial from typing import Optional import torch import torch.nn as nn from timm.models.vision_transformer import Block, PatchEmbed 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__) class SSL4EOViTConfig(PreTrainedConfig): model_type = "ssl4eo_vit" def __init__( self, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, intermediate_size=None, hidden_act="gelu", hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0, initializer_range=0.02, layer_norm_eps=1e-6, image_size=224, patch_size=16, num_channels=13, qkv_bias=True, mlp_ratio=4.0, global_pool=False, ssl_method="mae", modality="s2c", num_labels=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.ssl_method = ssl_method self.modality = modality self.num_labels = num_labels self.intermediate_size = int(hidden_size * mlp_ratio) if intermediate_size is None else intermediate_size class SSL4EOViTPreTrainedModel(PreTrainedModel): config: SSL4EOViTConfig base_model_prefix = "ssl4eo_vit" main_input_name = "pixel_values" input_modalities = ("image",) supports_gradient_checkpointing = True _no_split_modules = ["Block"] _supports_sdpa = False def _init_weights(self, module): super()._init_weights(module) if isinstance(module, SSL4EOViTModel): if hasattr(module, "pos_embed"): nn.init.trunc_normal_(module.pos_embed, std=self.config.initializer_range) if hasattr(module, "cls_token"): nn.init.trunc_normal_(module.cls_token, std=self.config.initializer_range) class SSL4EOViTModel(SSL4EOViTPreTrainedModel): def __init__(self, config: SSL4EOViTConfig, add_pooling_layer: bool = True): super().__init__(config) self.config = config image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0] self.patch_embed = PatchEmbed( img_size=image_size, patch_size=config.patch_size, in_chans=config.num_channels, embed_dim=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), requires_grad=True) norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps) 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) self.add_pooling_layer = add_pooling_layer self.post_init() def forward_features(self, pixel_values: 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) hidden_states = hidden_states + 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( 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 last_hidden_state, pooled_output = self.forward_features(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 SSL4EOViTForImageClassification(SSL4EOViTPreTrainedModel): def __init__(self, config: SSL4EOViTConfig): super().__init__(config) self.num_labels = config.num_labels self.ssl4eo_vit = SSL4EOViTModel(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.ssl4eo_vit(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__ = ["SSL4EOViTConfig", "SSL4EOViTForImageClassification", "SSL4EOViTModel", "SSL4EOViTPreTrainedModel"]