BiliSakura commited on
Commit
423bf77
·
verified ·
1 Parent(s): d06f7c9

Upload SARMAE ViT-B transformers-format pretrain checkpoint

Browse files
vit-base-patch16-pretrain/README.md ADDED
@@ -0,0 +1,128 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: cc-by-nc-4.0
3
+ language: en
4
+ tags:
5
+ - vision
6
+ - image-feature-extraction
7
+ - sar
8
+ - remote-sensing
9
+ - synthetic-aperture-radar
10
+ - masked-autoencoder
11
+ - transformers
12
+ library_name: transformers
13
+ pipeline_tag: image-feature-extraction
14
+ datasets:
15
+ - Wenquandan777/SAR-1M
16
+ arxiv: 2512.16635
17
+ base_model: Wenquandan777/SARMAE
18
+ model-index:
19
+ - name: sarmae-vit-base-patch16-pretrain
20
+ results: []
21
+ ---
22
+
23
+ # sarmae-vit-base-patch16-pretrain
24
+
25
+ SARMAE (ViT-B, patch 16) encoder checkpoint converted to native Hugging Face Transformers format.
26
+
27
+ SARMAE is a Noise-Aware Masked Autoencoder for self-supervised SAR representation learning, pretrained on [SAR-1M](https://huggingface.co/datasets/Wenquandan777/SAR-1M) with Speckle-Aware Representation Enhancement (SARE) and Semantic Anchor Representation Constraint (SARC).
28
+
29
+ - **Paper:** [2512.16635](https://arxiv.org/abs/2512.16635)
30
+ - **Legacy weights:** [Wenquandan777/SARMAE](https://huggingface.co/Wenquandan777/SARMAE)
31
+ - **Stage:** `pretrain`
32
+ - **Input:** 3 x 224 x 224 (single-channel SAR is repeated to 3 channels)
33
+ - **Architecture:** 12 layers, hidden size 768, 12 heads
34
+
35
+ ## Model specifications
36
+
37
+ | Property | Value |
38
+ |----------|-------|
39
+ | Model type | `sarmae` |
40
+ | Backbone | ViT-B |
41
+ | Patch size | 16 |
42
+ | Image size | 224 |
43
+ | Hidden size | 768 |
44
+ | Layers | 12 |
45
+ | Attention heads | 12 |
46
+ | Global pooling | `True` |
47
+ | Normalization mean | `[0.485, 0.456, 0.406]` |
48
+ | Normalization std | `[0.229, 0.224, 0.225]` |
49
+
50
+ ## Intended use
51
+
52
+ - SAR image feature extraction for downstream classification, detection, and segmentation
53
+ - Initializing OpenMMLab backbones (`mmrotate`, `mmseg`) after weight porting
54
+ - Research and non-commercial use under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/)
55
+
56
+ ## Quick start
57
+
58
+ Install dependencies:
59
+
60
+ ```bash
61
+ pip install transformers timm torch torchvision safetensors
62
+ ```
63
+
64
+ ### Feature extraction with `transformers.pipeline`
65
+
66
+ ```python
67
+ from transformers import pipeline
68
+
69
+ pipe = pipeline(
70
+ task="image-feature-extraction",
71
+ model="BiliSakura/SARMAE-transformers", subfolder="vit-base-patch16-pretrain",
72
+ trust_remote_code=True,
73
+ )
74
+ features = pipe(sar_image, pool=True, return_tensors=True)
75
+ print(features.shape)
76
+ ```
77
+
78
+ ### Direct model loading
79
+
80
+ ```python
81
+ from transformers import AutoModel, AutoImageProcessor
82
+
83
+ model = AutoModel.from_pretrained("BiliSakura/SARMAE-transformers", subfolder="vit-base-patch16-pretrain", trust_remote_code=True)
84
+ processor = AutoImageProcessor.from_pretrained("BiliSakura/SARMAE-transformers", subfolder="vit-base-patch16-pretrain", trust_remote_code=True)
85
+
86
+ inputs = processor(images=sar_image, return_tensors="pt")
87
+ outputs = model(**inputs)
88
+ pooled_features = outputs.pooler_output
89
+ ```
90
+
91
+ ### Local checkout
92
+
93
+ ```python
94
+ pipe = pipeline(
95
+ task="image-feature-extraction",
96
+ model="./sarmae-vit-base-patch16-pretrain",
97
+ trust_remote_code=True,
98
+ )
99
+ ```
100
+
101
+ ## Preprocessing
102
+
103
+ - Resize to 224x224
104
+ - Scale pixel values to `[0, 1]` (`rescale_factor=1/255`)
105
+ - Repeat grayscale SAR to 3 channels when `repeat_grayscale_channels=true`
106
+ - Normalize with ImageNet mean/std (same as SARMAE fine-tuning code)
107
+
108
+ ## Training data
109
+
110
+ Pretrained on **SAR-1M**, a million-scale SAR dataset with paired optical anchors for a subset of samples.
111
+
112
+ ## Citation
113
+
114
+ ```bibtex
115
+ @misc{liu2025sarmaemaskedautoencodersar,
116
+ title={SARMAE: Masked Autoencoder for SAR Representation Learning},
117
+ author={Danxu Liu and Di Wang and Hebaixu Wang and Haoyang Chen and Wentao Jiang and Yilin Cheng and Haonan Guo and Wei Cui and Jing Zhang},
118
+ year={2025},
119
+ eprint={2512.16635},
120
+ archivePrefix={arXiv},
121
+ primaryClass={cs.CV},
122
+ url={https://arxiv.org/abs/2512.16635},
123
+ }
124
+ ```
125
+
126
+ ## License
127
+
128
+ This model is released under [CC BY-NC 4.0](https://creativecommons.org/licenses/by-nc/4.0/).
vit-base-patch16-pretrain/config.json ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "SarmaeModel"
4
+ ],
5
+ "attention_probs_dropout_prob": 0.0,
6
+ "checkpoint_stage": "pretrain",
7
+ "dtype": "float32",
8
+ "global_pool": true,
9
+ "hidden_act": "gelu",
10
+ "hidden_dropout_prob": 0.0,
11
+ "hidden_size": 768,
12
+ "id2label": {},
13
+ "image_mean": [
14
+ 0.485,
15
+ 0.456,
16
+ 0.406
17
+ ],
18
+ "image_size": 224,
19
+ "image_std": [
20
+ 0.229,
21
+ 0.224,
22
+ 0.225
23
+ ],
24
+ "initializer_range": 0.02,
25
+ "intermediate_size": 3072,
26
+ "label2id": {},
27
+ "layer_norm_eps": 1e-06,
28
+ "mlp_ratio": 4.0,
29
+ "model_type": "sarmae",
30
+ "num_attention_heads": 12,
31
+ "num_channels": 3,
32
+ "num_hidden_layers": 12,
33
+ "patch_size": 16,
34
+ "qkv_bias": true,
35
+ "repeat_grayscale_channels": true,
36
+ "transformers_version": "5.0.0",
37
+ "auto_map": {
38
+ "AutoConfig": "modeling_sarmae.SarmaeConfig",
39
+ "AutoModel": "modeling_sarmae.SarmaeModel",
40
+ "AutoModelForImageClassification": "modeling_sarmae.SarmaeForImageClassification"
41
+ },
42
+ "custom_pipelines": {
43
+ "sarmae-feature-extraction": {
44
+ "impl": "pipeline_sarmae.SarmaeImageFeatureExtractionPipeline",
45
+ "pt": [
46
+ "AutoModel"
47
+ ]
48
+ }
49
+ }
50
+ }
vit-base-patch16-pretrain/image_processing_sarmae.py ADDED
@@ -0,0 +1,164 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
2
+ """Image processor for SARMAE models (self-contained for trust_remote_code)."""
3
+
4
+ from typing import Optional, Union
5
+
6
+ import numpy as np
7
+
8
+ from transformers.image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
9
+ from transformers.image_transforms import resize, to_channel_dimension_format
10
+ from transformers.image_utils import (
11
+ ChannelDimension,
12
+ ImageInput,
13
+ PILImageResampling,
14
+ infer_channel_dimension_format,
15
+ to_numpy_array,
16
+ valid_images,
17
+ validate_preprocess_arguments,
18
+ )
19
+ from transformers.utils import TensorType, filter_out_non_signature_kwargs, logging
20
+
21
+
22
+ logger = logging.get_logger(__name__)
23
+
24
+
25
+ def _repeat_grayscale_channels(image: np.ndarray, target_channels: int, input_data_format: ChannelDimension) -> np.ndarray:
26
+ if input_data_format == ChannelDimension.FIRST:
27
+ num_channels = image.shape[0]
28
+ if num_channels == target_channels:
29
+ return image
30
+ if num_channels == 1:
31
+ return np.repeat(image, target_channels, axis=0)
32
+ return image[:target_channels]
33
+ num_channels = image.shape[-1]
34
+ if num_channels == target_channels:
35
+ return image
36
+ if num_channels == 1:
37
+ return np.repeat(image, target_channels, axis=-1)
38
+ return image[..., :target_channels]
39
+
40
+
41
+ def _prepare_image_batch(images: ImageInput) -> list:
42
+ if isinstance(images, np.ndarray):
43
+ images = [images]
44
+ elif not isinstance(images, (list, tuple)):
45
+ images = [images]
46
+
47
+ prepared = []
48
+ for image in images:
49
+ array = to_numpy_array(image)
50
+ if array.ndim == 2:
51
+ array = np.expand_dims(array, axis=-1)
52
+ prepared.append(array)
53
+ return prepared
54
+
55
+
56
+ class SarmaeImageProcessor(BaseImageProcessor):
57
+ model_input_names = ["pixel_values"]
58
+
59
+ def __init__(
60
+ self,
61
+ do_resize: bool = True,
62
+ size: Optional[dict[str, int]] = None,
63
+ resample: PILImageResampling = PILImageResampling.BILINEAR,
64
+ do_rescale: bool = True,
65
+ rescale_factor: float = 1 / 255.0,
66
+ do_normalize: bool = True,
67
+ image_mean: Optional[Union[float, list[float]]] = None,
68
+ image_std: Optional[Union[float, list[float]]] = None,
69
+ do_convert_rgb: bool = False,
70
+ repeat_grayscale_channels: bool = True,
71
+ **kwargs,
72
+ ):
73
+ super().__init__(**kwargs)
74
+ size = size if size is not None else {"height": 224, "width": 224}
75
+ self.do_resize = do_resize
76
+ self.size = size
77
+ self.resample = resample
78
+ self.do_rescale = do_rescale
79
+ self.rescale_factor = rescale_factor
80
+ self.do_normalize = do_normalize
81
+ self.image_mean = image_mean
82
+ self.image_std = image_std
83
+ self.do_convert_rgb = do_convert_rgb
84
+ self.repeat_grayscale_channels = repeat_grayscale_channels
85
+
86
+ @filter_out_non_signature_kwargs()
87
+ def preprocess(
88
+ self,
89
+ images: ImageInput,
90
+ do_resize: Optional[bool] = None,
91
+ size: Optional[dict[str, int]] = None,
92
+ resample: Optional[PILImageResampling] = None,
93
+ do_rescale: Optional[bool] = None,
94
+ rescale_factor: Optional[float] = None,
95
+ do_normalize: Optional[bool] = None,
96
+ image_mean: Optional[Union[float, list[float]]] = None,
97
+ image_std: Optional[Union[float, list[float]]] = None,
98
+ return_tensors: Optional[Union[str, TensorType]] = None,
99
+ data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
100
+ input_data_format: Optional[Union[str, ChannelDimension]] = None,
101
+ do_convert_rgb: Optional[bool] = None,
102
+ repeat_grayscale_channels: Optional[bool] = None,
103
+ ):
104
+ do_resize = do_resize if do_resize is not None else self.do_resize
105
+ size = get_size_dict(size if size is not None else self.size, default_to_square=True)
106
+ resample = resample if resample is not None else self.resample
107
+ do_rescale = do_rescale if do_rescale is not None else self.do_rescale
108
+ rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
109
+ do_normalize = do_normalize if do_normalize is not None else self.do_normalize
110
+ image_mean = image_mean if image_mean is not None else self.image_mean
111
+ image_std = image_std if image_std is not None else self.image_std
112
+ do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
113
+ repeat_grayscale_channels = (
114
+ repeat_grayscale_channels if repeat_grayscale_channels is not None else self.repeat_grayscale_channels
115
+ )
116
+
117
+ if do_normalize and (image_mean is None or image_std is None):
118
+ raise ValueError("Normalization requires `image_mean` and `image_std` with one value per channel.")
119
+
120
+ images = _prepare_image_batch(images)
121
+ if not valid_images(images):
122
+ raise ValueError("Invalid image type. Must be PIL, numpy, or torch tensor.")
123
+
124
+ validate_preprocess_arguments(
125
+ do_rescale=do_rescale,
126
+ rescale_factor=rescale_factor,
127
+ do_normalize=do_normalize,
128
+ image_mean=image_mean,
129
+ image_std=image_std,
130
+ do_resize=do_resize,
131
+ size=size,
132
+ resample=resample,
133
+ )
134
+
135
+ processed_images = []
136
+ for image in images:
137
+ image = to_numpy_array(image)
138
+ if do_convert_rgb:
139
+ image = self._convert_image_to_rgb(image)
140
+ if input_data_format is None:
141
+ try:
142
+ input_data_format = infer_channel_dimension_format(image)
143
+ except ValueError:
144
+ input_data_format = ChannelDimension.LAST
145
+ if repeat_grayscale_channels:
146
+ image = _repeat_grayscale_channels(image, target_channels=3, input_data_format=input_data_format)
147
+ if do_resize:
148
+ image = resize(
149
+ image,
150
+ size=(size["height"], size["width"]),
151
+ resample=resample,
152
+ input_data_format=input_data_format,
153
+ )
154
+ if do_rescale:
155
+ image = image * rescale_factor
156
+ if do_normalize:
157
+ image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
158
+ image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
159
+ processed_images.append(image)
160
+
161
+ return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
162
+
163
+
164
+ __all__ = ["SarmaeImageProcessor"]
vit-base-patch16-pretrain/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a6208f15468b0526191ca3f4bc80a143cee87280b1360e9d29d623c5533cf890
3
+ size 343208584
vit-base-patch16-pretrain/modeling_sarmae.py ADDED
@@ -0,0 +1,200 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
2
+ """Self-contained SARMAE model and config for trust_remote_code loading."""
3
+
4
+ from __future__ import annotations
5
+
6
+ from functools import partial
7
+ from typing import Optional
8
+
9
+ import numpy as np
10
+ import torch
11
+ from timm.models.vision_transformer import Block, PatchEmbed
12
+ from torch import nn
13
+
14
+ from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
15
+ from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
16
+ from transformers.modeling_utils import PreTrainedModel
17
+ from transformers.processing_utils import Unpack
18
+ from transformers.utils import TransformersKwargs, logging
19
+
20
+
21
+ logger = logging.get_logger(__name__)
22
+
23
+ IMAGENET_MEAN = [0.485, 0.456, 0.406]
24
+ IMAGENET_STD = [0.229, 0.224, 0.225]
25
+
26
+
27
+ def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, cls_token: bool = False) -> np.ndarray:
28
+ grid_h = np.arange(grid_size, dtype=np.float32)
29
+ grid_w = np.arange(grid_size, dtype=np.float32)
30
+ grid = np.meshgrid(grid_w, grid_h)
31
+ grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size])
32
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
33
+ if cls_token:
34
+ pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
35
+ return pos_embed
36
+
37
+
38
+ def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
39
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
40
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
41
+ return np.concatenate([emb_h, emb_w], axis=1)
42
+
43
+
44
+ def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
45
+ omega = np.arange(embed_dim // 2, dtype=np.float32)
46
+ omega /= embed_dim / 2.0
47
+ omega = 1.0 / 10000**omega
48
+ pos = pos.reshape(-1)
49
+ out = np.einsum("m,d->md", pos, omega)
50
+ return np.concatenate([np.sin(out), np.cos(out)], axis=1)
51
+
52
+
53
+ class SarmaeConfig(PreTrainedConfig):
54
+ model_type = "sarmae"
55
+
56
+ def __init__(
57
+ self,
58
+ hidden_size: int = 768,
59
+ num_hidden_layers: int = 12,
60
+ num_attention_heads: int = 12,
61
+ intermediate_size: int | None = None,
62
+ hidden_act: str = "gelu",
63
+ hidden_dropout_prob: float = 0.0,
64
+ attention_probs_dropout_prob: float = 0.0,
65
+ initializer_range: float = 0.02,
66
+ layer_norm_eps: float = 1e-6,
67
+ image_size: int = 224,
68
+ patch_size: int = 16,
69
+ num_channels: int = 3,
70
+ qkv_bias: bool = True,
71
+ mlp_ratio: float = 4.0,
72
+ global_pool: bool = True,
73
+ repeat_grayscale_channels: bool = True,
74
+ checkpoint_stage: str = "pretrain",
75
+ image_mean: list[float] | None = None,
76
+ image_std: list[float] | None = None,
77
+ num_labels: int = 0,
78
+ **kwargs,
79
+ ):
80
+ super().__init__(**kwargs)
81
+ self.hidden_size = hidden_size
82
+ self.num_hidden_layers = num_hidden_layers
83
+ self.num_attention_heads = num_attention_heads
84
+ self.hidden_act = hidden_act
85
+ self.hidden_dropout_prob = hidden_dropout_prob
86
+ self.attention_probs_dropout_prob = attention_probs_dropout_prob
87
+ self.initializer_range = initializer_range
88
+ self.layer_norm_eps = layer_norm_eps
89
+ self.image_size = image_size
90
+ self.patch_size = patch_size
91
+ self.num_channels = num_channels
92
+ self.qkv_bias = qkv_bias
93
+ self.mlp_ratio = mlp_ratio
94
+ self.global_pool = global_pool
95
+ self.repeat_grayscale_channels = repeat_grayscale_channels
96
+ self.checkpoint_stage = checkpoint_stage
97
+ self.num_labels = num_labels
98
+ self.intermediate_size = int(hidden_size * mlp_ratio) if intermediate_size is None else intermediate_size
99
+ self.image_mean = image_mean if image_mean is not None else IMAGENET_MEAN
100
+ self.image_std = image_std if image_std is not None else IMAGENET_STD
101
+
102
+
103
+ class SarmaePreTrainedModel(PreTrainedModel):
104
+ config_class = SarmaeConfig
105
+ config: SarmaeConfig
106
+ base_model_prefix = "sarmae"
107
+ main_input_name = "pixel_values"
108
+ input_modalities = ("image",)
109
+ supports_gradient_checkpointing = True
110
+ _no_split_modules = ["Block"]
111
+
112
+
113
+ class SarmaeModel(SarmaePreTrainedModel):
114
+ def __init__(self, config: SarmaeConfig, add_pooling_layer: bool = True):
115
+ super().__init__(config)
116
+ self.config = config
117
+ self.add_pooling_layer = add_pooling_layer
118
+
119
+ image_size = config.image_size if isinstance(config.image_size, int) else config.image_size[0]
120
+ norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
121
+ self.patch_embed = PatchEmbed(image_size, config.patch_size, config.num_channels, config.hidden_size)
122
+ self.num_patches = self.patch_embed.num_patches
123
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
124
+ self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, config.hidden_size))
125
+ pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.num_patches**0.5), cls_token=True)
126
+ self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
127
+ self.blocks = nn.ModuleList([
128
+ Block(config.hidden_size, config.num_attention_heads, config.mlp_ratio, qkv_bias=config.qkv_bias, norm_layer=norm_layer)
129
+ for _ in range(config.num_hidden_layers)
130
+ ])
131
+ self.global_pool = config.global_pool
132
+ if self.global_pool:
133
+ self.fc_norm = norm_layer(config.hidden_size)
134
+ self.norm = None
135
+ else:
136
+ self.fc_norm = None
137
+ self.norm = norm_layer(config.hidden_size)
138
+ self.post_init()
139
+
140
+ def forward(
141
+ self,
142
+ pixel_values: Optional[torch.Tensor] = None,
143
+ return_dict: Optional[bool] = None,
144
+ **kwargs: Unpack[TransformersKwargs],
145
+ ) -> BaseModelOutputWithPooling:
146
+ if pixel_values is None:
147
+ raise ValueError("You must specify `pixel_values`")
148
+ pixel_values = pixel_values.to(dtype=self.dtype)
149
+ if return_dict is None:
150
+ return_dict = self.config.use_return_dict
151
+
152
+ batch_size = pixel_values.shape[0]
153
+ patch_tokens = self.patch_embed(pixel_values)
154
+ cls_tokens = self.cls_token.expand(batch_size, -1, -1)
155
+ hidden_states = torch.cat((cls_tokens, patch_tokens), dim=1) + self.pos_embed
156
+ for block in self.blocks:
157
+ hidden_states = block(hidden_states)
158
+ if self.global_pool:
159
+ pooled_output = self.fc_norm(hidden_states[:, 1:, :].mean(dim=1))
160
+ else:
161
+ hidden_states = self.norm(hidden_states)
162
+ pooled_output = hidden_states[:, 0]
163
+ if not self.add_pooling_layer:
164
+ pooled_output = None
165
+ if not return_dict:
166
+ return (hidden_states, pooled_output)
167
+ return BaseModelOutputWithPooling(last_hidden_state=hidden_states, pooler_output=pooled_output)
168
+
169
+
170
+ class SarmaeForImageClassification(SarmaePreTrainedModel):
171
+ def __init__(self, config: SarmaeConfig):
172
+ super().__init__(config)
173
+ self.sarmae = SarmaeModel(config, add_pooling_layer=True)
174
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
175
+ self.post_init()
176
+
177
+ def forward(
178
+ self,
179
+ pixel_values: Optional[torch.Tensor] = None,
180
+ labels: Optional[torch.Tensor] = None,
181
+ return_dict: Optional[bool] = None,
182
+ **kwargs: Unpack[TransformersKwargs],
183
+ ) -> ImageClassifierOutput:
184
+ outputs = self.sarmae(pixel_values=pixel_values, return_dict=True, **kwargs)
185
+ logits = self.classifier(outputs.pooler_output)
186
+ loss = None
187
+ if labels is not None:
188
+ loss = self.loss_function(labels, logits, self.config, **kwargs)
189
+ if not return_dict:
190
+ output = (logits,) + outputs[1:]
191
+ return ((loss,) + output) if loss is not None else output
192
+ return ImageClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
193
+
194
+
195
+ __all__ = [
196
+ "SarmaeConfig",
197
+ "SarmaeForImageClassification",
198
+ "SarmaeModel",
199
+ "SarmaePreTrainedModel",
200
+ ]
vit-base-patch16-pretrain/pipeline_sarmae.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright 2026 SARMAE Authors and The HuggingFace Inc. team.
2
+ """SARMAE image feature extraction pipeline (self-contained for trust_remote_code)."""
3
+
4
+ from typing import Any, Union
5
+
6
+ from transformers.pipelines.base import GenericTensor, build_pipeline_init_args
7
+ from transformers.pipelines.image_feature_extraction import ImageFeatureExtractionPipeline
8
+ from transformers.utils import add_end_docstrings, is_vision_available
9
+
10
+
11
+ if is_vision_available():
12
+ from transformers.image_utils import load_image
13
+
14
+
15
+ @add_end_docstrings(
16
+ build_pipeline_init_args(has_image_processor=True),
17
+ """
18
+ pool (`bool`, *optional*, defaults to `False`):
19
+ Whether or not to return the pooled output. If `False`, the model will return the raw hidden states.
20
+ """,
21
+ )
22
+ class SarmaeImageFeatureExtractionPipeline(ImageFeatureExtractionPipeline):
23
+ def _sanitize_parameters(
24
+ self,
25
+ image_processor_kwargs=None,
26
+ return_tensors=None,
27
+ pool=None,
28
+ **kwargs,
29
+ ):
30
+ preprocess_params = {} if image_processor_kwargs is None else dict(image_processor_kwargs)
31
+ if "timeout" in kwargs:
32
+ preprocess_params["timeout"] = kwargs["timeout"]
33
+ postprocess_params = {}
34
+ if pool is not None:
35
+ postprocess_params["pool"] = pool
36
+ if return_tensors is not None:
37
+ postprocess_params["return_tensors"] = return_tensors
38
+ return preprocess_params, {}, postprocess_params
39
+
40
+ def preprocess(self, image, timeout=None, **image_processor_kwargs) -> dict[str, GenericTensor]:
41
+ if not isinstance(image, (list, tuple)) and not hasattr(image, "shape"):
42
+ image = load_image(image, timeout=timeout)
43
+ model_inputs = self.image_processor(image, return_tensors="pt", **image_processor_kwargs)
44
+ model_inputs = model_inputs.to(self.dtype)
45
+ return model_inputs
46
+
47
+ def __call__(self, *args: Union[str, Any, list[Any]], **kwargs: Any) -> list[Any]:
48
+ return super().__call__(*args, **kwargs)
49
+
50
+
51
+ __all__ = ["SarmaeImageFeatureExtractionPipeline"]
vit-base-patch16-pretrain/preprocessor_config.json ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "image_processor_type": "SarmaeImageProcessor",
3
+ "size": {
4
+ "height": 224,
5
+ "width": 224
6
+ },
7
+ "do_resize": true,
8
+ "do_rescale": true,
9
+ "rescale_factor": 0.00392156862745098,
10
+ "do_normalize": true,
11
+ "do_convert_rgb": false,
12
+ "repeat_grayscale_channels": true,
13
+ "image_mean": [
14
+ 0.485,
15
+ 0.456,
16
+ 0.406
17
+ ],
18
+ "image_std": [
19
+ 0.229,
20
+ 0.224,
21
+ 0.225
22
+ ],
23
+ "auto_map": {
24
+ "AutoImageProcessor": "image_processing_sarmae.SarmaeImageProcessor"
25
+ }
26
+ }