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