Hugging Face's logo

BiliSakura
/
SkySensePlusPlus-transformers

Feature Extraction
Safetensors
remote-sensing
earth-observation
skysensepp
Model card Files
xet
 Community
BiliSakura commited on
Commit
8d93654
·
verified ·
1 Parent(s): 7c92590

Upload SkySense++ Transformers checkpoints

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. README.md +140 -0
  2. skysensepp-fewshot-release/config.json +357 -0
  3. skysensepp-fewshot-release/configuration_skysensepp.py +281 -0
  4. skysensepp-fewshot-release/conversion_manifest.json +1595 -0
  5. skysensepp-fewshot-release/model.safetensors +3 -0
  6. skysensepp-fewshot-release/modeling_skysensepp.py +214 -0
  7. skysensepp-fewshot-release/modeling_skysensepp_components.py +238 -0
  8. skysensepp-fewshot-release/modeling_skysensepp_fusion_neck.py +164 -0
  9. skysensepp-fewshot-release/modeling_skysensepp_swinv2_msl.py +343 -0
  10. skysensepp-fewshot-release/modeling_skysensepp_vit_msl.py +265 -0
  11. skysensepp-fewshot-release/modeling_utils.py +557 -0
  12. skysensepp-fewshot-release/pipeline_skysensepp.py +86 -0
  13. skysensepp-fewshot-release/pipeline_skysensepp_fewshot.py +132 -0
  14. skysensepp-fewshot-release/pipeline_skysensepp_fusion.py +53 -0
  15. skysensepp-fusion-neck/config.json +47 -0
  16. skysensepp-fusion-neck/configuration_skysensepp.py +165 -0
  17. skysensepp-fusion-neck/conversion_manifest.json +301 -0
  18. skysensepp-fusion-neck/model.safetensors +3 -0
  19. skysensepp-fusion-neck/modeling_skysensepp_fusion_neck.py +164 -0
  20. skysensepp-fusion-neck/modeling_utils.py +557 -0
  21. skysensepp-fusion-neck/pipeline_skysensepp.py +86 -0
  22. skysensepp-fusion-neck/pipeline_skysensepp_fusion.py +53 -0
  23. skysensepp-swinv2-msl-hr/__init__.py +25 -0
  24. skysensepp-swinv2-msl-hr/config.json +82 -0
  25. skysensepp-swinv2-msl-hr/configuration_skysensepp.py +124 -0
  26. skysensepp-swinv2-msl-hr/conversion_manifest.json +523 -0
  27. skysensepp-swinv2-msl-hr/model.safetensors +3 -0
  28. skysensepp-swinv2-msl-hr/modeling_skysensepp_swinv2_msl.py +343 -0
  29. skysensepp-swinv2-msl-hr/modeling_skysensepp_vit_msl.py +265 -0
  30. skysensepp-swinv2-msl-hr/modeling_utils.py +557 -0
  31. skysensepp-swinv2-msl-hr/pipeline_skysensepp.py +86 -0
  32. skysensepp-swinv2-msl-hr/pipeline_skysensepp_fusion.py +53 -0
  33. skysensepp-vit-msl-s1/__init__.py +25 -0
  34. skysensepp-vit-msl-s1/config.json +68 -0
  35. skysensepp-vit-msl-s1/configuration_skysensepp.py +124 -0
  36. skysensepp-vit-msl-s1/conversion_manifest.json +305 -0
  37. skysensepp-vit-msl-s1/model.safetensors +3 -0
  38. skysensepp-vit-msl-s1/modeling_skysensepp_swinv2_msl.py +343 -0
  39. skysensepp-vit-msl-s1/modeling_skysensepp_vit_msl.py +265 -0
  40. skysensepp-vit-msl-s1/modeling_utils.py +557 -0
  41. skysensepp-vit-msl-s1/pipeline_skysensepp.py +86 -0
  42. skysensepp-vit-msl-s1/pipeline_skysensepp_fusion.py +53 -0
  43. skysensepp-vit-msl-s2/__init__.py +25 -0
  44. skysensepp-vit-msl-s2/config.json +68 -0
  45. skysensepp-vit-msl-s2/configuration_skysensepp.py +124 -0
  46. skysensepp-vit-msl-s2/conversion_manifest.json +305 -0
  47. skysensepp-vit-msl-s2/model.safetensors +3 -0
  48. skysensepp-vit-msl-s2/modeling_skysensepp_swinv2_msl.py +343 -0
  49. skysensepp-vit-msl-s2/modeling_skysensepp_vit_msl.py +265 -0
  50. skysensepp-vit-msl-s2/modeling_utils.py +557 -0
README.md ADDED
@@ -0,0 +1,140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ tags:
4
+ - remote-sensing
5
+ - earth-observation
6
+ - skysensepp
7
+ - feature-extraction
8
+ pipeline_tag: feature-extraction
9
+ ---
10
+
11
+ # SkySense++ Transformers
12
+
13
+ HuggingFace-compatible checkpoints for SkySense++ zero-shot MSL backbones, converted from the official release weights.
14
+
15
+ ## Checkpoints
16
+
17
+ | Directory | Modality | Architecture | Source |
18
+ |-----------|----------|--------------|--------|
19
+ | `skysensepp-swinv2-msl-hr` | High-res optical | SwinV2 Huge + MSL | `skysensepp_release_hr.pth` |
20
+ | `skysensepp-vit-msl-s2` | Sentinel-2 | ViT-Large + MSL | `skysensepp_release_s2.pth` |
21
+ | `skysensepp-vit-msl-s1` | Sentinel-1 | ViT-Large + MSL | `skysensepp_release_s1.pth` |
22
+ | `skysensepp-fusion-neck` | Multi-modal fusion (optional) | TransformerEncoder | `fusion.*` from `skysensepp_release.ckpt` |
23
+ | `skysensepp-fewshot-release` | Full 1-shot segmentation | HR + S2 + S1 + fusion + VAE + UPerHead | `skysensepp_release.ckpt` |
24
+
25
+ Each subdirectory is a self-contained HuggingFace model repo with remote code (`trust_remote_code=True`).
26
+
27
+ The fusion neck is an **optional** component — backbone checkpoints do not include or require it by default.
28
+
29
+ The few-shot release bundles all submodules into one end-to-end model (~6.8 GB).
30
+
31
+ ## Usage
32
+
33
+ ```python
34
+ from transformers import pipeline
35
+ import torch
36
+
37
+ MODEL = "/path/to/SkySensePlusPlus-transformers/skysensepp-swinv2-msl-hr"
38
+
39
+ pipe = pipeline(
40
+ task="image-feature-extraction",
41
+ model=MODEL,
42
+ trust_remote_code=True,
43
+ device="cpu",
44
+ )
45
+
46
+ hr_img = torch.randn(1, 3, 512, 512)
47
+ annotation = torch.zeros(1, 512, 512, dtype=torch.long) # semantic class indices
48
+
49
+ features = pipe(hr_img, annotation=annotation)
50
+ print(features["last_hidden_state"].shape) # (1, 2816, 16, 16)
51
+ ```
52
+
53
+ Sentinel-2 / Sentinel-1 backbones use the same pipeline pattern:
54
+
55
+ ```python
56
+ s2_pipe = pipeline(
57
+ task="image-feature-extraction",
58
+ model="/path/to/skysensepp-vit-msl-s2",
59
+ trust_remote_code=True,
60
+ device="cpu",
61
+ )
62
+
63
+ s2_img = torch.randn(1, 10, 16, 16)
64
+ s2_anno = torch.zeros(1, 16, 16, dtype=torch.long)
65
+ features = s2_pipe(s2_img, annotation=s2_anno)
66
+ print(features["last_hidden_state"].shape)
67
+ ```
68
+
69
+ SkySense++ MSL models require both imagery and a semantic annotation map. Use class index `0` for background/unlabeled regions during zero-shot feature extraction.
70
+
71
+ ### Optional fusion neck
72
+
73
+ ```python
74
+ fusion_pipe = pipeline(
75
+ task="skysensepp-fusion",
76
+ model="/path/to/skysensepp-fusion-neck",
77
+ trust_remote_code=True,
78
+ device="cpu",
79
+ )
80
+
81
+ # Concatenated HR + S2 + S1 stage-3 tokens per spatial location
82
+ hidden_states = torch.randn(256, 3, 2816)
83
+ fused = fusion_pipe(hidden_states)
84
+
85
+ print(fused["pooler_output"].shape) # (256, 1024)
86
+ ```
87
+
88
+ ### Few-shot / 1-shot segmentation
89
+
90
+ The full release model expects vertically stacked prompt+query inputs (prompt on top, query on bottom):
91
+
92
+ ```python
93
+ from transformers import pipeline
94
+ import torch
95
+
96
+ MODEL = "/path/to/SkySensePlusPlus-transformers/skysensepp-fewshot-release"
97
+
98
+ pipe = pipeline(
99
+ task="skysensepp-fewshot",
100
+ model=MODEL,
101
+ trust_remote_code=True,
102
+ device=0, # GPU recommended (~24 GB); CPU OOMs at 1024×512 HR
103
+ )
104
+
105
+ # Stacked HR (3, 1024, 512), S2/S1 with seq=2, RGB targets (ImageNet-normalized)
106
+ hr = torch.randn(1, 3, 1024, 512)
107
+ s2 = torch.randn(1, 10, 2, 32, 32)
108
+ s1 = torch.randn(1, 2, 2, 32, 32)
109
+ targets = torch.randn(1, 3, 1024, 512) # use real RGB annotation maps in practice
110
+ anno_mask = torch.zeros(1, 8, 4, dtype=torch.long)
111
+ anno_mask[:, 4:, :] = 1 # mask query (bottom) half
112
+
113
+ result = pipe(hr, s2_img=s2, s1_img=s1, targets=targets, anno_mask=anno_mask)
114
+ print(result["logits"].shape) # (1, 65, 512, 512) — query region only
115
+ ```
116
+
117
+ ## Conversion
118
+
119
+ Source project: `/home/czy/local/projects/SkySensePlusPlus-transformers`
120
+
121
+ ```bash
122
+ conda activate rsgen
123
+ python scripts/convert_checkpoint_to_hf.py \
124
+ --input-path /path/to/skysensepp_release_hr.pth \
125
+ --modality hr \
126
+ --output-dir /path/to/skysensepp-swinv2-msl-hr \
127
+ --clean-output
128
+
129
+ # Full few-shot release (~6.8 GB)
130
+ python scripts/convert_checkpoint_to_hf.py \
131
+ --input-path /path/to/skysensepp_release.ckpt \
132
+ --modality fewshot \
133
+ --output-dir /path/to/skysensepp-fewshot-release \
134
+ --clean-output
135
+ ```
136
+
137
+ ## Notes
138
+
139
+ - HR conversion skips Swin relative-position buffers (`relative_position_index`, `relative_coords_table`). These are **deterministically recomputed** at init from window geometry — not randomly initialized. Learned CPB weights (`cpb_mlp`, `logit_scale`) are loaded.
140
+ - The few-shot model uses the same 62 skipped HR buffers; all 1522 learned tensors load with 0 unexpected keys.
skysensepp-fewshot-release/config.json ADDED
@@ -0,0 +1,357 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "return_dict": true,
3
+ "output_hidden_states": false,
4
+ "dtype": "float32",
5
+ "chunk_size_feed_forward": 0,
6
+ "is_encoder_decoder": false,
7
+ "architectures": [
8
+ "SkySensePlusPlusModel"
9
+ ],
10
+ "id2label": {
11
+ "0": "LABEL_0",
12
+ "1": "LABEL_1"
13
+ },
14
+ "label2id": {
15
+ "LABEL_0": 0,
16
+ "LABEL_1": 1
17
+ },
18
+ "problem_type": null,
19
+ "_name_or_path": "",
20
+ "transformers_version": "5.0.0",
21
+ "sources": [
22
+ "hr",
23
+ "s2",
24
+ "s1"
25
+ ],
26
+ "vocabulary_size": 64,
27
+ "use_modal_vae": true,
28
+ "upsample_results": true,
29
+ "backbone_hr": {
30
+ "return_dict": true,
31
+ "output_hidden_states": false,
32
+ "dtype": null,
33
+ "chunk_size_feed_forward": 0,
34
+ "is_encoder_decoder": false,
35
+ "architectures": null,
36
+ "id2label": {
37
+ "0": "LABEL_0",
38
+ "1": "LABEL_1"
39
+ },
40
+ "label2id": {
41
+ "LABEL_0": 0,
42
+ "LABEL_1": 1
43
+ },
44
+ "problem_type": null,
45
+ "_name_or_path": "",
46
+ "arch": "huge",
47
+ "embed_dims": 352,
48
+ "depths": [
49
+ 2,
50
+ 2,
51
+ 18,
52
+ 2
53
+ ],
54
+ "num_heads": [
55
+ 8,
56
+ 16,
57
+ 32,
58
+ 64
59
+ ],
60
+ "extra_norm_every_n_blocks": 6,
61
+ "img_size": 512,
62
+ "patch_size": 4,
63
+ "in_channels": 3,
64
+ "window_size": 8,
65
+ "drop_rate": 0.0,
66
+ "drop_path_rate": 0.2,
67
+ "out_indices": [
68
+ 0,
69
+ 1,
70
+ 2,
71
+ 3
72
+ ],
73
+ "use_abs_pos_embed": false,
74
+ "with_cp": false,
75
+ "pad_small_map": false,
76
+ "pretrained_window_sizes": [
77
+ 0,
78
+ 0,
79
+ 0,
80
+ 0
81
+ ],
82
+ "is_post_norm_downsample": true,
83
+ "vocabulary_size": 64,
84
+ "num_vocabulary_tokens": 65,
85
+ "merge_stage": 2,
86
+ "use_attn": true,
87
+ "model_type": "skysensepp_swinv2_msl",
88
+ "output_attentions": false
89
+ },
90
+ "backbone_s2": {
91
+ "return_dict": true,
92
+ "output_hidden_states": false,
93
+ "dtype": null,
94
+ "chunk_size_feed_forward": 0,
95
+ "is_encoder_decoder": false,
96
+ "architectures": null,
97
+ "id2label": {
98
+ "0": "LABEL_0",
99
+ "1": "LABEL_1"
100
+ },
101
+ "label2id": {
102
+ "LABEL_0": 0,
103
+ "LABEL_1": 1
104
+ },
105
+ "problem_type": null,
106
+ "_name_or_path": "",
107
+ "img_size": 16,
108
+ "patch_size": 4,
109
+ "in_channels": 10,
110
+ "embed_dims": 1024,
111
+ "num_layers": 24,
112
+ "num_heads": 16,
113
+ "mlp_ratio": 4,
114
+ "out_indices": [
115
+ 5,
116
+ 11,
117
+ 17,
118
+ 23
119
+ ],
120
+ "qkv_bias": true,
121
+ "drop_rate": 0.0,
122
+ "attn_drop_rate": 0.0,
123
+ "drop_path_rate": 0.3,
124
+ "with_cls_token": false,
125
+ "output_cls_token": false,
126
+ "patch_norm": false,
127
+ "final_norm": false,
128
+ "with_cp": false,
129
+ "vocabulary_size": 64,
130
+ "num_vocabulary_tokens": 65,
131
+ "merge_stage": 4,
132
+ "use_attn": false,
133
+ "modality": "s2",
134
+ "model_type": "skysensepp_vit_msl",
135
+ "output_attentions": false
136
+ },
137
+ "backbone_s1": {
138
+ "return_dict": true,
139
+ "output_hidden_states": false,
140
+ "dtype": null,
141
+ "chunk_size_feed_forward": 0,
142
+ "is_encoder_decoder": false,
143
+ "architectures": null,
144
+ "id2label": {
145
+ "0": "LABEL_0",
146
+ "1": "LABEL_1"
147
+ },
148
+ "label2id": {
149
+ "LABEL_0": 0,
150
+ "LABEL_1": 1
151
+ },
152
+ "problem_type": null,
153
+ "_name_or_path": "",
154
+ "img_size": 16,
155
+ "patch_size": 4,
156
+ "in_channels": 2,
157
+ "embed_dims": 1024,
158
+ "num_layers": 24,
159
+ "num_heads": 16,
160
+ "mlp_ratio": 4,
161
+ "out_indices": [
162
+ 5,
163
+ 11,
164
+ 17,
165
+ 23
166
+ ],
167
+ "qkv_bias": true,
168
+ "drop_rate": 0.0,
169
+ "attn_drop_rate": 0.0,
170
+ "drop_path_rate": 0.3,
171
+ "with_cls_token": false,
172
+ "output_cls_token": false,
173
+ "patch_norm": false,
174
+ "final_norm": false,
175
+ "with_cp": false,
176
+ "vocabulary_size": 64,
177
+ "num_vocabulary_tokens": 65,
178
+ "merge_stage": 4,
179
+ "use_attn": false,
180
+ "modality": "s1",
181
+ "model_type": "skysensepp_vit_msl",
182
+ "output_attentions": false
183
+ },
184
+ "head_s2": {
185
+ "return_dict": true,
186
+ "output_hidden_states": false,
187
+ "dtype": null,
188
+ "chunk_size_feed_forward": 0,
189
+ "is_encoder_decoder": false,
190
+ "architectures": null,
191
+ "id2label": {
192
+ "0": "LABEL_0",
193
+ "1": "LABEL_1"
194
+ },
195
+ "label2id": {
196
+ "LABEL_0": 0,
197
+ "LABEL_1": 1
198
+ },
199
+ "problem_type": null,
200
+ "_name_or_path": "",
201
+ "in_dim": 1024,
202
+ "out_dim": 2816,
203
+ "up_scale": 4,
204
+ "model_type": "skysensepp_up_head",
205
+ "output_attentions": false
206
+ },
207
+ "head_s1": {
208
+ "return_dict": true,
209
+ "output_hidden_states": false,
210
+ "dtype": null,
211
+ "chunk_size_feed_forward": 0,
212
+ "is_encoder_decoder": false,
213
+ "architectures": null,
214
+ "id2label": {
215
+ "0": "LABEL_0",
216
+ "1": "LABEL_1"
217
+ },
218
+ "label2id": {
219
+ "LABEL_0": 0,
220
+ "LABEL_1": 1
221
+ },
222
+ "problem_type": null,
223
+ "_name_or_path": "",
224
+ "in_dim": 1024,
225
+ "out_dim": 2816,
226
+ "up_scale": 4,
227
+ "model_type": "skysensepp_up_head",
228
+ "output_attentions": false
229
+ },
230
+ "fusion": {
231
+ "return_dict": true,
232
+ "output_hidden_states": false,
233
+ "dtype": null,
234
+ "chunk_size_feed_forward": 0,
235
+ "is_encoder_decoder": false,
236
+ "architectures": null,
237
+ "id2label": {
238
+ "0": "LABEL_0",
239
+ "1": "LABEL_1"
240
+ },
241
+ "label2id": {
242
+ "LABEL_0": 0,
243
+ "LABEL_1": 1
244
+ },
245
+ "problem_type": null,
246
+ "_name_or_path": "",
247
+ "input_dims": 2816,
248
+ "embed_dims": 1024,
249
+ "num_layers": 24,
250
+ "num_heads": 16,
251
+ "mlp_ratio": 4,
252
+ "qkv_bias": true,
253
+ "drop_rate": 0.0,
254
+ "attn_drop_rate": 0.0,
255
+ "drop_path_rate": 0.3,
256
+ "with_cls_token": true,
257
+ "output_cls_token": true,
258
+ "with_cp": false,
259
+ "model_type": "skysensepp_fusion_neck",
260
+ "output_attentions": false
261
+ },
262
+ "modality_vae": {
263
+ "return_dict": true,
264
+ "output_hidden_states": false,
265
+ "dtype": null,
266
+ "chunk_size_feed_forward": 0,
267
+ "is_encoder_decoder": false,
268
+ "architectures": null,
269
+ "id2label": {
270
+ "0": "LABEL_0",
271
+ "1": "LABEL_1"
272
+ },
273
+ "label2id": {
274
+ "LABEL_0": 0,
275
+ "LABEL_1": 1
276
+ },
277
+ "problem_type": null,
278
+ "_name_or_path": "",
279
+ "input_shape_hr": [
280
+ 2816,
281
+ 32,
282
+ 16
283
+ ],
284
+ "input_shape_s2": [
285
+ 2816,
286
+ 32,
287
+ 16
288
+ ],
289
+ "input_shape_s1": [
290
+ 2816,
291
+ 32,
292
+ 16
293
+ ],
294
+ "conv_dim": 256,
295
+ "z_dim": 256,
296
+ "n_codebook": 8192,
297
+ "model_type": "skysensepp_modality_vae",
298
+ "output_attentions": false
299
+ },
300
+ "head_rec_hr": {
301
+ "return_dict": true,
302
+ "output_hidden_states": false,
303
+ "dtype": null,
304
+ "chunk_size_feed_forward": 0,
305
+ "is_encoder_decoder": false,
306
+ "architectures": null,
307
+ "id2label": {
308
+ "0": "LABEL_0",
309
+ "1": "LABEL_1"
310
+ },
311
+ "label2id": {
312
+ "LABEL_0": 0,
313
+ "LABEL_1": 1
314
+ },
315
+ "problem_type": null,
316
+ "_name_or_path": "",
317
+ "in_channels": [
318
+ 704,
319
+ 704,
320
+ 1408,
321
+ 2816,
322
+ 1024
323
+ ],
324
+ "channels": 512,
325
+ "num_classes": 65,
326
+ "pool_scales": [
327
+ 1,
328
+ 2,
329
+ 3,
330
+ 6
331
+ ],
332
+ "dropout_ratio": 0.1,
333
+ "align_corners": false,
334
+ "model_type": "skysensepp_uper_head",
335
+ "output_attentions": false
336
+ },
337
+ "model_type": "skysensepp",
338
+ "output_attentions": false,
339
+ "auto_map": {
340
+ "AutoConfig": "configuration_skysensepp.SkySensePlusPlusConfig",
341
+ "AutoModel": "modeling_skysensepp.SkySensePlusPlusModel"
342
+ },
343
+ "custom_pipelines": {
344
+ "skysensepp-fusion": {
345
+ "impl": "pipeline_skysensepp_fewshot.SkySensePlusPlusFewShotPipeline",
346
+ "pt": [
347
+ "AutoModel"
348
+ ]
349
+ },
350
+ "skysensepp-fewshot": {
351
+ "impl": "pipeline_skysensepp_fewshot.SkySensePlusPlusFewShotPipeline",
352
+ "pt": [
353
+ "AutoModel"
354
+ ]
355
+ }
356
+ }
357
+ }
skysensepp-fewshot-release/configuration_skysensepp.py ADDED
@@ -0,0 +1,281 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration classes for SkySense++ MSL backbones."""
2
+
3
+ from transformers import PretrainedConfig
4
+
5
+
6
+ class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig):
7
+ """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical)."""
8
+
9
+ model_type = "skysensepp_swinv2_msl"
10
+
11
+ arch_zoo = {
12
+ "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
13
+ "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
14
+ "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0},
15
+ "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0},
16
+ "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6},
17
+ "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6},
18
+ }
19
+
20
+ def __init__(
21
+ self,
22
+ arch="huge",
23
+ img_size=512,
24
+ patch_size=4,
25
+ in_channels=3,
26
+ window_size=8,
27
+ drop_rate=0.0,
28
+ drop_path_rate=0.2,
29
+ out_indices=(0, 1, 2, 3),
30
+ use_abs_pos_embed=False,
31
+ with_cp=False,
32
+ pad_small_map=False,
33
+ pretrained_window_sizes=(0, 0, 0, 0),
34
+ is_post_norm_downsample=True,
35
+ vocabulary_size=64,
36
+ merge_stage=2,
37
+ use_attn=True,
38
+ **kwargs,
39
+ ):
40
+ super().__init__(**kwargs)
41
+
42
+ arch = arch.lower()
43
+ if arch not in self.arch_zoo:
44
+ raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
45
+ arch_settings = self.arch_zoo[arch]
46
+
47
+ self.arch = arch
48
+ self.embed_dims = arch_settings["embed_dims"]
49
+ self.depths = arch_settings["depths"]
50
+ self.num_heads = arch_settings["num_heads"]
51
+ self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"]
52
+
53
+ self.img_size = img_size
54
+ self.patch_size = patch_size
55
+ self.in_channels = in_channels
56
+ self.window_size = window_size
57
+ self.drop_rate = drop_rate
58
+ self.drop_path_rate = drop_path_rate
59
+ self.out_indices = list(out_indices)
60
+ self.use_abs_pos_embed = use_abs_pos_embed
61
+ self.with_cp = with_cp
62
+ self.pad_small_map = pad_small_map
63
+ self.pretrained_window_sizes = list(pretrained_window_sizes)
64
+ self.is_post_norm_downsample = is_post_norm_downsample
65
+
66
+ self.vocabulary_size = vocabulary_size
67
+ self.num_vocabulary_tokens = vocabulary_size + 1
68
+ self.merge_stage = merge_stage
69
+ self.use_attn = use_attn
70
+
71
+
72
+ class SkySensePlusPlusViTMSLConfig(PretrainedConfig):
73
+ """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1)."""
74
+
75
+ model_type = "skysensepp_vit_msl"
76
+
77
+ def __init__(
78
+ self,
79
+ img_size=16,
80
+ patch_size=4,
81
+ in_channels=10,
82
+ embed_dims=1024,
83
+ num_layers=24,
84
+ num_heads=16,
85
+ mlp_ratio=4,
86
+ out_indices=(5, 11, 17, 23),
87
+ qkv_bias=True,
88
+ drop_rate=0.0,
89
+ attn_drop_rate=0.0,
90
+ drop_path_rate=0.3,
91
+ with_cls_token=False,
92
+ output_cls_token=False,
93
+ patch_norm=False,
94
+ final_norm=False,
95
+ with_cp=False,
96
+ vocabulary_size=64,
97
+ merge_stage=4,
98
+ use_attn=False,
99
+ modality="s2",
100
+ **kwargs,
101
+ ):
102
+ super().__init__(**kwargs)
103
+ self.img_size = img_size
104
+ self.patch_size = patch_size
105
+ self.in_channels = in_channels
106
+ self.embed_dims = embed_dims
107
+ self.num_layers = num_layers
108
+ self.num_heads = num_heads
109
+ self.mlp_ratio = mlp_ratio
110
+ self.out_indices = list(out_indices)
111
+ self.qkv_bias = qkv_bias
112
+ self.drop_rate = drop_rate
113
+ self.attn_drop_rate = attn_drop_rate
114
+ self.drop_path_rate = drop_path_rate
115
+ self.with_cls_token = with_cls_token
116
+ self.output_cls_token = output_cls_token
117
+ self.patch_norm = patch_norm
118
+ self.final_norm = final_norm
119
+ self.with_cp = with_cp
120
+ self.vocabulary_size = vocabulary_size
121
+ self.num_vocabulary_tokens = vocabulary_size + 1
122
+ self.merge_stage = merge_stage
123
+ self.use_attn = use_attn
124
+ self.modality = modality
125
+
126
+
127
+ class SkySensePlusPlusFusionNeckConfig(PretrainedConfig):
128
+ """Configuration for SkySense++ multi-modal fusion neck (TransformerEncoder).
129
+
130
+ Optional component — not used by default backbone checkpoints.
131
+ Fuses concatenated HR/S2/S1 stage-3 features (2816-dim) via a ViT encoder
132
+ with cls token output (1024-dim).
133
+ """
134
+
135
+ model_type = "skysensepp_fusion_neck"
136
+
137
+ def __init__(
138
+ self,
139
+ input_dims=2816,
140
+ embed_dims=1024,
141
+ num_layers=24,
142
+ num_heads=16,
143
+ mlp_ratio=4,
144
+ qkv_bias=True,
145
+ drop_rate=0.0,
146
+ attn_drop_rate=0.0,
147
+ drop_path_rate=0.3,
148
+ with_cls_token=True,
149
+ output_cls_token=True,
150
+ with_cp=False,
151
+ **kwargs,
152
+ ):
153
+ super().__init__(**kwargs)
154
+ self.input_dims = input_dims
155
+ self.embed_dims = embed_dims
156
+ self.num_layers = num_layers
157
+ self.num_heads = num_heads
158
+ self.mlp_ratio = mlp_ratio
159
+ self.qkv_bias = qkv_bias
160
+ self.drop_rate = drop_rate
161
+ self.attn_drop_rate = attn_drop_rate
162
+ self.drop_path_rate = drop_path_rate
163
+ self.with_cls_token = with_cls_token
164
+ self.output_cls_token = output_cls_token
165
+ self.with_cp = with_cp
166
+
167
+
168
+ class UPHeadConfig(PretrainedConfig):
169
+ model_type = "skysensepp_up_head"
170
+
171
+ def __init__(self, in_dim=1024, out_dim=2816, up_scale=4, **kwargs):
172
+ super().__init__(**kwargs)
173
+ self.in_dim = in_dim
174
+ self.out_dim = out_dim
175
+ self.up_scale = up_scale
176
+
177
+
178
+ class UPerHeadConfig(PretrainedConfig):
179
+ model_type = "skysensepp_uper_head"
180
+
181
+ def __init__(
182
+ self,
183
+ in_channels=(704, 704, 1408, 2816, 1024),
184
+ channels=512,
185
+ num_classes=65,
186
+ pool_scales=(1, 2, 3, 6),
187
+ dropout_ratio=0.1,
188
+ align_corners=False,
189
+ **kwargs,
190
+ ):
191
+ super().__init__(**kwargs)
192
+ self.in_channels = list(in_channels)
193
+ self.channels = channels
194
+ self.num_classes = num_classes
195
+ self.pool_scales = list(pool_scales)
196
+ self.dropout_ratio = dropout_ratio
197
+ self.align_corners = align_corners
198
+
199
+
200
+ class ModalityVAEConfig(PretrainedConfig):
201
+ model_type = "skysensepp_modality_vae"
202
+
203
+ def __init__(
204
+ self,
205
+ input_shape_hr=(2816, 32, 16),
206
+ input_shape_s2=(2816, 32, 16),
207
+ input_shape_s1=(2816, 32, 16),
208
+ conv_dim=256,
209
+ z_dim=256,
210
+ n_codebook=8192,
211
+ **kwargs,
212
+ ):
213
+ super().__init__(**kwargs)
214
+ self.input_shape_hr = list(input_shape_hr)
215
+ self.input_shape_s2 = list(input_shape_s2)
216
+ self.input_shape_s1 = list(input_shape_s1)
217
+ self.conv_dim = conv_dim
218
+ self.z_dim = z_dim
219
+ self.n_codebook = n_codebook
220
+
221
+
222
+ class SkySensePlusPlusConfig(PretrainedConfig):
223
+ """Full SkySense++ config for few-shot / 1-shot release checkpoint."""
224
+
225
+ model_type = "skysensepp"
226
+
227
+ def __init__(
228
+ self,
229
+ sources=("hr", "s2", "s1"),
230
+ vocabulary_size=64,
231
+ use_modal_vae=True,
232
+ upsample_results=True,
233
+ backbone_hr=None,
234
+ backbone_s2=None,
235
+ backbone_s1=None,
236
+ head_s2=None,
237
+ head_s1=None,
238
+ fusion=None,
239
+ modality_vae=None,
240
+ head_rec_hr=None,
241
+ **kwargs,
242
+ ):
243
+ super().__init__(**kwargs)
244
+ self.sources = list(sources)
245
+ self.vocabulary_size = vocabulary_size
246
+ self.use_modal_vae = use_modal_vae
247
+ self.upsample_results = upsample_results
248
+ self.backbone_hr = (
249
+ backbone_hr
250
+ if isinstance(backbone_hr, SkySensePlusPlusSwinV2MSLConfig)
251
+ else SkySensePlusPlusSwinV2MSLConfig(**(backbone_hr or {}))
252
+ )
253
+ self.backbone_s2 = (
254
+ backbone_s2
255
+ if isinstance(backbone_s2, SkySensePlusPlusViTMSLConfig)
256
+ else SkySensePlusPlusViTMSLConfig(**(backbone_s2 or {"modality": "s2"}))
257
+ )
258
+ self.backbone_s1 = (
259
+ backbone_s1
260
+ if isinstance(backbone_s1, SkySensePlusPlusViTMSLConfig)
261
+ else SkySensePlusPlusViTMSLConfig(
262
+ **(backbone_s1 or {"modality": "s1", "in_channels": 2})
263
+ )
264
+ )
265
+ self.head_s2 = head_s2 if isinstance(head_s2, UPHeadConfig) else UPHeadConfig(**(head_s2 or {}))
266
+ self.head_s1 = head_s1 if isinstance(head_s1, UPHeadConfig) else UPHeadConfig(**(head_s1 or {}))
267
+ self.fusion = (
268
+ fusion
269
+ if isinstance(fusion, SkySensePlusPlusFusionNeckConfig)
270
+ else SkySensePlusPlusFusionNeckConfig(**(fusion or {}))
271
+ )
272
+ self.modality_vae = (
273
+ modality_vae
274
+ if isinstance(modality_vae, ModalityVAEConfig)
275
+ else ModalityVAEConfig(**(modality_vae or {}))
276
+ )
277
+ self.head_rec_hr = (
278
+ head_rec_hr
279
+ if isinstance(head_rec_hr, UPerHeadConfig)
280
+ else UPerHeadConfig(**(head_rec_hr or {}))
281
+ )
skysensepp-fewshot-release/conversion_manifest.json ADDED
@@ -0,0 +1,1595 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release.ckpt",
3
+ "modality": "fewshot",
4
+ "model_class": "SkySensePlusPlusModel",
5
+ "num_tensors": 1522,
6
+ "missing_keys": [
7
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.relative_coords_table",
8
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.relative_position_index",
9
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.relative_coords_table",
10
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.relative_position_index",
11
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.relative_coords_table",
12
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.relative_position_index",
13
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.relative_coords_table",
14
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.relative_position_index",
15
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.relative_coords_table",
16
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.relative_position_index",
17
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.relative_coords_table",
18
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.relative_position_index",
19
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.relative_coords_table",
20
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.relative_position_index",
21
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.relative_coords_table",
22
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.relative_position_index",
23
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.relative_coords_table",
24
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.relative_position_index",
25
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.relative_coords_table",
26
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.relative_position_index",
27
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.relative_coords_table",
28
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.relative_position_index",
29
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.relative_coords_table",
30
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.relative_position_index",
31
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.relative_coords_table",
32
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.relative_position_index",
33
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.relative_coords_table",
34
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.relative_position_index",
35
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.relative_coords_table",
36
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.relative_position_index",
37
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.relative_coords_table",
38
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.relative_position_index",
39
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.relative_coords_table",
40
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.relative_position_index",
41
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.relative_coords_table",
42
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.relative_position_index",
43
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.relative_coords_table",
44
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.relative_position_index",
45
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.relative_coords_table",
46
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.relative_position_index",
47
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.relative_coords_table",
48
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.relative_position_index",
49
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.relative_coords_table",
50
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.relative_position_index",
51
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.relative_coords_table",
52
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.relative_position_index",
53
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.relative_coords_table",
54
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.relative_position_index",
55
+ "head_rec_hr.psp_modules.0.1.conv.bias",
56
+ "head_rec_hr.psp_modules.1.1.conv.bias",
57
+ "head_rec_hr.psp_modules.2.1.conv.bias",
58
+ "head_rec_hr.psp_modules.3.1.conv.bias",
59
+ "head_rec_hr.bottleneck.conv.bias",
60
+ "head_rec_hr.lateral_convs.0.conv.bias",
61
+ "head_rec_hr.lateral_convs.1.conv.bias",
62
+ "head_rec_hr.lateral_convs.2.conv.bias",
63
+ "head_rec_hr.lateral_convs.3.conv.bias",
64
+ "head_rec_hr.fpn_convs.0.conv.bias",
65
+ "head_rec_hr.fpn_convs.1.conv.bias",
66
+ "head_rec_hr.fpn_convs.2.conv.bias",
67
+ "head_rec_hr.fpn_convs.3.conv.bias",
68
+ "head_rec_hr.fpn_bottleneck.conv.bias"
69
+ ],
70
+ "unexpected_keys": [],
71
+ "tensor_names": [
72
+ "backbone_hr.attn1.attn.in_proj_bias",
73
+ "backbone_hr.attn1.attn.in_proj_weight",
74
+ "backbone_hr.attn1.attn.out_proj.bias",
75
+ "backbone_hr.attn1.attn.out_proj.weight",
76
+ "backbone_hr.attn1.proj_in.bias",
77
+ "backbone_hr.attn1.proj_in.weight",
78
+ "backbone_hr.attn1.proj_out.bias",
79
+ "backbone_hr.attn1.proj_out.weight",
80
+ "backbone_hr.attn2.attn.in_proj_bias",
81
+ "backbone_hr.attn2.attn.in_proj_weight",
82
+ "backbone_hr.attn2.attn.out_proj.bias",
83
+ "backbone_hr.attn2.attn.out_proj.weight",
84
+ "backbone_hr.attn2.proj_in.bias",
85
+ "backbone_hr.attn2.proj_in.weight",
86
+ "backbone_hr.attn2.proj_out.bias",
87
+ "backbone_hr.attn2.proj_out.weight",
88
+ "backbone_hr.attn3.attn.in_proj_bias",
89
+ "backbone_hr.attn3.attn.in_proj_weight",
90
+ "backbone_hr.attn3.attn.out_proj.bias",
91
+ "backbone_hr.attn3.attn.out_proj.weight",
92
+ "backbone_hr.attn3.proj_in.bias",
93
+ "backbone_hr.attn3.proj_in.weight",
94
+ "backbone_hr.attn3.proj_out.bias",
95
+ "backbone_hr.attn3.proj_out.weight",
96
+ "backbone_hr.mask_token",
97
+ "backbone_hr.norm0.bias",
98
+ "backbone_hr.norm0.weight",
99
+ "backbone_hr.norm1.bias",
100
+ "backbone_hr.norm1.weight",
101
+ "backbone_hr.norm2.bias",
102
+ "backbone_hr.norm2.weight",
103
+ "backbone_hr.norm3.bias",
104
+ "backbone_hr.norm3.weight",
105
+ "backbone_hr.norm_attn.bias",
106
+ "backbone_hr.norm_attn.weight",
107
+ "backbone_hr.patch_embed.norm.bias",
108
+ "backbone_hr.patch_embed.norm.weight",
109
+ "backbone_hr.patch_embed.projection.bias",
110
+ "backbone_hr.patch_embed.projection.weight",
111
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.cpb_mlp.0.bias",
112
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.cpb_mlp.0.weight",
113
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.cpb_mlp.2.weight",
114
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.logit_scale",
115
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.proj.bias",
116
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.proj.weight",
117
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.q_bias",
118
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.qkv.weight",
119
+ "backbone_hr.stages.0.blocks.0.attn.w_msa.v_bias",
120
+ "backbone_hr.stages.0.blocks.0.ffn.layers.0.bias",
121
+ "backbone_hr.stages.0.blocks.0.ffn.layers.0.weight",
122
+ "backbone_hr.stages.0.blocks.0.ffn.layers.3.bias",
123
+ "backbone_hr.stages.0.blocks.0.ffn.layers.3.weight",
124
+ "backbone_hr.stages.0.blocks.0.norm1.bias",
125
+ "backbone_hr.stages.0.blocks.0.norm1.weight",
126
+ "backbone_hr.stages.0.blocks.0.norm2.bias",
127
+ "backbone_hr.stages.0.blocks.0.norm2.weight",
128
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.cpb_mlp.0.bias",
129
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.cpb_mlp.0.weight",
130
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.cpb_mlp.2.weight",
131
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.logit_scale",
132
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.proj.bias",
133
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.proj.weight",
134
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.q_bias",
135
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.qkv.weight",
136
+ "backbone_hr.stages.0.blocks.1.attn.w_msa.v_bias",
137
+ "backbone_hr.stages.0.blocks.1.ffn.layers.0.bias",
138
+ "backbone_hr.stages.0.blocks.1.ffn.layers.0.weight",
139
+ "backbone_hr.stages.0.blocks.1.ffn.layers.3.bias",
140
+ "backbone_hr.stages.0.blocks.1.ffn.layers.3.weight",
141
+ "backbone_hr.stages.0.blocks.1.norm1.bias",
142
+ "backbone_hr.stages.0.blocks.1.norm1.weight",
143
+ "backbone_hr.stages.0.blocks.1.norm2.bias",
144
+ "backbone_hr.stages.0.blocks.1.norm2.weight",
145
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.cpb_mlp.0.bias",
146
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.cpb_mlp.0.weight",
147
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.cpb_mlp.2.weight",
148
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.logit_scale",
149
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.proj.bias",
150
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.proj.weight",
151
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.q_bias",
152
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.qkv.weight",
153
+ "backbone_hr.stages.1.blocks.0.attn.w_msa.v_bias",
154
+ "backbone_hr.stages.1.blocks.0.ffn.layers.0.bias",
155
+ "backbone_hr.stages.1.blocks.0.ffn.layers.0.weight",
156
+ "backbone_hr.stages.1.blocks.0.ffn.layers.3.bias",
157
+ "backbone_hr.stages.1.blocks.0.ffn.layers.3.weight",
158
+ "backbone_hr.stages.1.blocks.0.norm1.bias",
159
+ "backbone_hr.stages.1.blocks.0.norm1.weight",
160
+ "backbone_hr.stages.1.blocks.0.norm2.bias",
161
+ "backbone_hr.stages.1.blocks.0.norm2.weight",
162
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.cpb_mlp.0.bias",
163
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.cpb_mlp.0.weight",
164
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.cpb_mlp.2.weight",
165
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.logit_scale",
166
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.proj.bias",
167
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.proj.weight",
168
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.q_bias",
169
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.qkv.weight",
170
+ "backbone_hr.stages.1.blocks.1.attn.w_msa.v_bias",
171
+ "backbone_hr.stages.1.blocks.1.ffn.layers.0.bias",
172
+ "backbone_hr.stages.1.blocks.1.ffn.layers.0.weight",
173
+ "backbone_hr.stages.1.blocks.1.ffn.layers.3.bias",
174
+ "backbone_hr.stages.1.blocks.1.ffn.layers.3.weight",
175
+ "backbone_hr.stages.1.blocks.1.norm1.bias",
176
+ "backbone_hr.stages.1.blocks.1.norm1.weight",
177
+ "backbone_hr.stages.1.blocks.1.norm2.bias",
178
+ "backbone_hr.stages.1.blocks.1.norm2.weight",
179
+ "backbone_hr.stages.1.downsample.norm.bias",
180
+ "backbone_hr.stages.1.downsample.norm.weight",
181
+ "backbone_hr.stages.1.downsample.reduction.weight",
182
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.cpb_mlp.0.bias",
183
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.cpb_mlp.0.weight",
184
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.cpb_mlp.2.weight",
185
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.logit_scale",
186
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.proj.bias",
187
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.proj.weight",
188
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.q_bias",
189
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.qkv.weight",
190
+ "backbone_hr.stages.2.blocks.0.attn.w_msa.v_bias",
191
+ "backbone_hr.stages.2.blocks.0.ffn.layers.0.bias",
192
+ "backbone_hr.stages.2.blocks.0.ffn.layers.0.weight",
193
+ "backbone_hr.stages.2.blocks.0.ffn.layers.3.bias",
194
+ "backbone_hr.stages.2.blocks.0.ffn.layers.3.weight",
195
+ "backbone_hr.stages.2.blocks.0.norm1.bias",
196
+ "backbone_hr.stages.2.blocks.0.norm1.weight",
197
+ "backbone_hr.stages.2.blocks.0.norm2.bias",
198
+ "backbone_hr.stages.2.blocks.0.norm2.weight",
199
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.cpb_mlp.0.bias",
200
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.cpb_mlp.0.weight",
201
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.cpb_mlp.2.weight",
202
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.logit_scale",
203
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.proj.bias",
204
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.proj.weight",
205
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.q_bias",
206
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.qkv.weight",
207
+ "backbone_hr.stages.2.blocks.1.attn.w_msa.v_bias",
208
+ "backbone_hr.stages.2.blocks.1.ffn.layers.0.bias",
209
+ "backbone_hr.stages.2.blocks.1.ffn.layers.0.weight",
210
+ "backbone_hr.stages.2.blocks.1.ffn.layers.3.bias",
211
+ "backbone_hr.stages.2.blocks.1.ffn.layers.3.weight",
212
+ "backbone_hr.stages.2.blocks.1.norm1.bias",
213
+ "backbone_hr.stages.2.blocks.1.norm1.weight",
214
+ "backbone_hr.stages.2.blocks.1.norm2.bias",
215
+ "backbone_hr.stages.2.blocks.1.norm2.weight",
216
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.cpb_mlp.0.bias",
217
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.cpb_mlp.0.weight",
218
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.cpb_mlp.2.weight",
219
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.logit_scale",
220
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.proj.bias",
221
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.proj.weight",
222
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.q_bias",
223
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.qkv.weight",
224
+ "backbone_hr.stages.2.blocks.10.attn.w_msa.v_bias",
225
+ "backbone_hr.stages.2.blocks.10.ffn.layers.0.bias",
226
+ "backbone_hr.stages.2.blocks.10.ffn.layers.0.weight",
227
+ "backbone_hr.stages.2.blocks.10.ffn.layers.3.bias",
228
+ "backbone_hr.stages.2.blocks.10.ffn.layers.3.weight",
229
+ "backbone_hr.stages.2.blocks.10.norm1.bias",
230
+ "backbone_hr.stages.2.blocks.10.norm1.weight",
231
+ "backbone_hr.stages.2.blocks.10.norm2.bias",
232
+ "backbone_hr.stages.2.blocks.10.norm2.weight",
233
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.cpb_mlp.0.bias",
234
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.cpb_mlp.0.weight",
235
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.cpb_mlp.2.weight",
236
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.logit_scale",
237
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.proj.bias",
238
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.proj.weight",
239
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.q_bias",
240
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.qkv.weight",
241
+ "backbone_hr.stages.2.blocks.11.attn.w_msa.v_bias",
242
+ "backbone_hr.stages.2.blocks.11.ffn.layers.0.bias",
243
+ "backbone_hr.stages.2.blocks.11.ffn.layers.0.weight",
244
+ "backbone_hr.stages.2.blocks.11.ffn.layers.3.bias",
245
+ "backbone_hr.stages.2.blocks.11.ffn.layers.3.weight",
246
+ "backbone_hr.stages.2.blocks.11.norm1.bias",
247
+ "backbone_hr.stages.2.blocks.11.norm1.weight",
248
+ "backbone_hr.stages.2.blocks.11.norm2.bias",
249
+ "backbone_hr.stages.2.blocks.11.norm2.weight",
250
+ "backbone_hr.stages.2.blocks.11.norm3.bias",
251
+ "backbone_hr.stages.2.blocks.11.norm3.weight",
252
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.cpb_mlp.0.bias",
253
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.cpb_mlp.0.weight",
254
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.cpb_mlp.2.weight",
255
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.logit_scale",
256
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.proj.bias",
257
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.proj.weight",
258
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.q_bias",
259
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.qkv.weight",
260
+ "backbone_hr.stages.2.blocks.12.attn.w_msa.v_bias",
261
+ "backbone_hr.stages.2.blocks.12.ffn.layers.0.bias",
262
+ "backbone_hr.stages.2.blocks.12.ffn.layers.0.weight",
263
+ "backbone_hr.stages.2.blocks.12.ffn.layers.3.bias",
264
+ "backbone_hr.stages.2.blocks.12.ffn.layers.3.weight",
265
+ "backbone_hr.stages.2.blocks.12.norm1.bias",
266
+ "backbone_hr.stages.2.blocks.12.norm1.weight",
267
+ "backbone_hr.stages.2.blocks.12.norm2.bias",
268
+ "backbone_hr.stages.2.blocks.12.norm2.weight",
269
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.cpb_mlp.0.bias",
270
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.cpb_mlp.0.weight",
271
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.cpb_mlp.2.weight",
272
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.logit_scale",
273
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.proj.bias",
274
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.proj.weight",
275
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.q_bias",
276
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.qkv.weight",
277
+ "backbone_hr.stages.2.blocks.13.attn.w_msa.v_bias",
278
+ "backbone_hr.stages.2.blocks.13.ffn.layers.0.bias",
279
+ "backbone_hr.stages.2.blocks.13.ffn.layers.0.weight",
280
+ "backbone_hr.stages.2.blocks.13.ffn.layers.3.bias",
281
+ "backbone_hr.stages.2.blocks.13.ffn.layers.3.weight",
282
+ "backbone_hr.stages.2.blocks.13.norm1.bias",
283
+ "backbone_hr.stages.2.blocks.13.norm1.weight",
284
+ "backbone_hr.stages.2.blocks.13.norm2.bias",
285
+ "backbone_hr.stages.2.blocks.13.norm2.weight",
286
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.cpb_mlp.0.bias",
287
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.cpb_mlp.0.weight",
288
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.cpb_mlp.2.weight",
289
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.logit_scale",
290
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.proj.bias",
291
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.proj.weight",
292
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.q_bias",
293
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.qkv.weight",
294
+ "backbone_hr.stages.2.blocks.14.attn.w_msa.v_bias",
295
+ "backbone_hr.stages.2.blocks.14.ffn.layers.0.bias",
296
+ "backbone_hr.stages.2.blocks.14.ffn.layers.0.weight",
297
+ "backbone_hr.stages.2.blocks.14.ffn.layers.3.bias",
298
+ "backbone_hr.stages.2.blocks.14.ffn.layers.3.weight",
299
+ "backbone_hr.stages.2.blocks.14.norm1.bias",
300
+ "backbone_hr.stages.2.blocks.14.norm1.weight",
301
+ "backbone_hr.stages.2.blocks.14.norm2.bias",
302
+ "backbone_hr.stages.2.blocks.14.norm2.weight",
303
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.cpb_mlp.0.bias",
304
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.cpb_mlp.0.weight",
305
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.cpb_mlp.2.weight",
306
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.logit_scale",
307
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.proj.bias",
308
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.proj.weight",
309
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.q_bias",
310
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.qkv.weight",
311
+ "backbone_hr.stages.2.blocks.15.attn.w_msa.v_bias",
312
+ "backbone_hr.stages.2.blocks.15.ffn.layers.0.bias",
313
+ "backbone_hr.stages.2.blocks.15.ffn.layers.0.weight",
314
+ "backbone_hr.stages.2.blocks.15.ffn.layers.3.bias",
315
+ "backbone_hr.stages.2.blocks.15.ffn.layers.3.weight",
316
+ "backbone_hr.stages.2.blocks.15.norm1.bias",
317
+ "backbone_hr.stages.2.blocks.15.norm1.weight",
318
+ "backbone_hr.stages.2.blocks.15.norm2.bias",
319
+ "backbone_hr.stages.2.blocks.15.norm2.weight",
320
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.cpb_mlp.0.bias",
321
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.cpb_mlp.0.weight",
322
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.cpb_mlp.2.weight",
323
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.logit_scale",
324
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.proj.bias",
325
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.proj.weight",
326
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.q_bias",
327
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.qkv.weight",
328
+ "backbone_hr.stages.2.blocks.16.attn.w_msa.v_bias",
329
+ "backbone_hr.stages.2.blocks.16.ffn.layers.0.bias",
330
+ "backbone_hr.stages.2.blocks.16.ffn.layers.0.weight",
331
+ "backbone_hr.stages.2.blocks.16.ffn.layers.3.bias",
332
+ "backbone_hr.stages.2.blocks.16.ffn.layers.3.weight",
333
+ "backbone_hr.stages.2.blocks.16.norm1.bias",
334
+ "backbone_hr.stages.2.blocks.16.norm1.weight",
335
+ "backbone_hr.stages.2.blocks.16.norm2.bias",
336
+ "backbone_hr.stages.2.blocks.16.norm2.weight",
337
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.cpb_mlp.0.bias",
338
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.cpb_mlp.0.weight",
339
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.cpb_mlp.2.weight",
340
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.logit_scale",
341
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.proj.bias",
342
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.proj.weight",
343
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.q_bias",
344
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.qkv.weight",
345
+ "backbone_hr.stages.2.blocks.17.attn.w_msa.v_bias",
346
+ "backbone_hr.stages.2.blocks.17.ffn.layers.0.bias",
347
+ "backbone_hr.stages.2.blocks.17.ffn.layers.0.weight",
348
+ "backbone_hr.stages.2.blocks.17.ffn.layers.3.bias",
349
+ "backbone_hr.stages.2.blocks.17.ffn.layers.3.weight",
350
+ "backbone_hr.stages.2.blocks.17.norm1.bias",
351
+ "backbone_hr.stages.2.blocks.17.norm1.weight",
352
+ "backbone_hr.stages.2.blocks.17.norm2.bias",
353
+ "backbone_hr.stages.2.blocks.17.norm2.weight",
354
+ "backbone_hr.stages.2.blocks.17.norm3.bias",
355
+ "backbone_hr.stages.2.blocks.17.norm3.weight",
356
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.cpb_mlp.0.bias",
357
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.cpb_mlp.0.weight",
358
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.cpb_mlp.2.weight",
359
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.logit_scale",
360
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.proj.bias",
361
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.proj.weight",
362
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.q_bias",
363
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.qkv.weight",
364
+ "backbone_hr.stages.2.blocks.2.attn.w_msa.v_bias",
365
+ "backbone_hr.stages.2.blocks.2.ffn.layers.0.bias",
366
+ "backbone_hr.stages.2.blocks.2.ffn.layers.0.weight",
367
+ "backbone_hr.stages.2.blocks.2.ffn.layers.3.bias",
368
+ "backbone_hr.stages.2.blocks.2.ffn.layers.3.weight",
369
+ "backbone_hr.stages.2.blocks.2.norm1.bias",
370
+ "backbone_hr.stages.2.blocks.2.norm1.weight",
371
+ "backbone_hr.stages.2.blocks.2.norm2.bias",
372
+ "backbone_hr.stages.2.blocks.2.norm2.weight",
373
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.cpb_mlp.0.bias",
374
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.cpb_mlp.0.weight",
375
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.cpb_mlp.2.weight",
376
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.logit_scale",
377
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.proj.bias",
378
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.proj.weight",
379
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.q_bias",
380
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.qkv.weight",
381
+ "backbone_hr.stages.2.blocks.3.attn.w_msa.v_bias",
382
+ "backbone_hr.stages.2.blocks.3.ffn.layers.0.bias",
383
+ "backbone_hr.stages.2.blocks.3.ffn.layers.0.weight",
384
+ "backbone_hr.stages.2.blocks.3.ffn.layers.3.bias",
385
+ "backbone_hr.stages.2.blocks.3.ffn.layers.3.weight",
386
+ "backbone_hr.stages.2.blocks.3.norm1.bias",
387
+ "backbone_hr.stages.2.blocks.3.norm1.weight",
388
+ "backbone_hr.stages.2.blocks.3.norm2.bias",
389
+ "backbone_hr.stages.2.blocks.3.norm2.weight",
390
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.cpb_mlp.0.bias",
391
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.cpb_mlp.0.weight",
392
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.cpb_mlp.2.weight",
393
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.logit_scale",
394
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.proj.bias",
395
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.proj.weight",
396
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.q_bias",
397
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.qkv.weight",
398
+ "backbone_hr.stages.2.blocks.4.attn.w_msa.v_bias",
399
+ "backbone_hr.stages.2.blocks.4.ffn.layers.0.bias",
400
+ "backbone_hr.stages.2.blocks.4.ffn.layers.0.weight",
401
+ "backbone_hr.stages.2.blocks.4.ffn.layers.3.bias",
402
+ "backbone_hr.stages.2.blocks.4.ffn.layers.3.weight",
403
+ "backbone_hr.stages.2.blocks.4.norm1.bias",
404
+ "backbone_hr.stages.2.blocks.4.norm1.weight",
405
+ "backbone_hr.stages.2.blocks.4.norm2.bias",
406
+ "backbone_hr.stages.2.blocks.4.norm2.weight",
407
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.cpb_mlp.0.bias",
408
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.cpb_mlp.0.weight",
409
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.cpb_mlp.2.weight",
410
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.logit_scale",
411
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.proj.bias",
412
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.proj.weight",
413
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.q_bias",
414
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.qkv.weight",
415
+ "backbone_hr.stages.2.blocks.5.attn.w_msa.v_bias",
416
+ "backbone_hr.stages.2.blocks.5.ffn.layers.0.bias",
417
+ "backbone_hr.stages.2.blocks.5.ffn.layers.0.weight",
418
+ "backbone_hr.stages.2.blocks.5.ffn.layers.3.bias",
419
+ "backbone_hr.stages.2.blocks.5.ffn.layers.3.weight",
420
+ "backbone_hr.stages.2.blocks.5.norm1.bias",
421
+ "backbone_hr.stages.2.blocks.5.norm1.weight",
422
+ "backbone_hr.stages.2.blocks.5.norm2.bias",
423
+ "backbone_hr.stages.2.blocks.5.norm2.weight",
424
+ "backbone_hr.stages.2.blocks.5.norm3.bias",
425
+ "backbone_hr.stages.2.blocks.5.norm3.weight",
426
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.cpb_mlp.0.bias",
427
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.cpb_mlp.0.weight",
428
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.cpb_mlp.2.weight",
429
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.logit_scale",
430
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.proj.bias",
431
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.proj.weight",
432
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.q_bias",
433
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.qkv.weight",
434
+ "backbone_hr.stages.2.blocks.6.attn.w_msa.v_bias",
435
+ "backbone_hr.stages.2.blocks.6.ffn.layers.0.bias",
436
+ "backbone_hr.stages.2.blocks.6.ffn.layers.0.weight",
437
+ "backbone_hr.stages.2.blocks.6.ffn.layers.3.bias",
438
+ "backbone_hr.stages.2.blocks.6.ffn.layers.3.weight",
439
+ "backbone_hr.stages.2.blocks.6.norm1.bias",
440
+ "backbone_hr.stages.2.blocks.6.norm1.weight",
441
+ "backbone_hr.stages.2.blocks.6.norm2.bias",
442
+ "backbone_hr.stages.2.blocks.6.norm2.weight",
443
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.cpb_mlp.0.bias",
444
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.cpb_mlp.0.weight",
445
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.cpb_mlp.2.weight",
446
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.logit_scale",
447
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.proj.bias",
448
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.proj.weight",
449
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.q_bias",
450
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.qkv.weight",
451
+ "backbone_hr.stages.2.blocks.7.attn.w_msa.v_bias",
452
+ "backbone_hr.stages.2.blocks.7.ffn.layers.0.bias",
453
+ "backbone_hr.stages.2.blocks.7.ffn.layers.0.weight",
454
+ "backbone_hr.stages.2.blocks.7.ffn.layers.3.bias",
455
+ "backbone_hr.stages.2.blocks.7.ffn.layers.3.weight",
456
+ "backbone_hr.stages.2.blocks.7.norm1.bias",
457
+ "backbone_hr.stages.2.blocks.7.norm1.weight",
458
+ "backbone_hr.stages.2.blocks.7.norm2.bias",
459
+ "backbone_hr.stages.2.blocks.7.norm2.weight",
460
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.cpb_mlp.0.bias",
461
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.cpb_mlp.0.weight",
462
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.cpb_mlp.2.weight",
463
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.logit_scale",
464
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.proj.bias",
465
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.proj.weight",
466
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.q_bias",
467
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.qkv.weight",
468
+ "backbone_hr.stages.2.blocks.8.attn.w_msa.v_bias",
469
+ "backbone_hr.stages.2.blocks.8.ffn.layers.0.bias",
470
+ "backbone_hr.stages.2.blocks.8.ffn.layers.0.weight",
471
+ "backbone_hr.stages.2.blocks.8.ffn.layers.3.bias",
472
+ "backbone_hr.stages.2.blocks.8.ffn.layers.3.weight",
473
+ "backbone_hr.stages.2.blocks.8.norm1.bias",
474
+ "backbone_hr.stages.2.blocks.8.norm1.weight",
475
+ "backbone_hr.stages.2.blocks.8.norm2.bias",
476
+ "backbone_hr.stages.2.blocks.8.norm2.weight",
477
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.cpb_mlp.0.bias",
478
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.cpb_mlp.0.weight",
479
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.cpb_mlp.2.weight",
480
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.logit_scale",
481
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.proj.bias",
482
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.proj.weight",
483
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.q_bias",
484
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.qkv.weight",
485
+ "backbone_hr.stages.2.blocks.9.attn.w_msa.v_bias",
486
+ "backbone_hr.stages.2.blocks.9.ffn.layers.0.bias",
487
+ "backbone_hr.stages.2.blocks.9.ffn.layers.0.weight",
488
+ "backbone_hr.stages.2.blocks.9.ffn.layers.3.bias",
489
+ "backbone_hr.stages.2.blocks.9.ffn.layers.3.weight",
490
+ "backbone_hr.stages.2.blocks.9.norm1.bias",
491
+ "backbone_hr.stages.2.blocks.9.norm1.weight",
492
+ "backbone_hr.stages.2.blocks.9.norm2.bias",
493
+ "backbone_hr.stages.2.blocks.9.norm2.weight",
494
+ "backbone_hr.stages.2.downsample.norm.bias",
495
+ "backbone_hr.stages.2.downsample.norm.weight",
496
+ "backbone_hr.stages.2.downsample.reduction.weight",
497
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.cpb_mlp.0.bias",
498
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.cpb_mlp.0.weight",
499
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.cpb_mlp.2.weight",
500
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.logit_scale",
501
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.proj.bias",
502
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.proj.weight",
503
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.q_bias",
504
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.qkv.weight",
505
+ "backbone_hr.stages.3.blocks.0.attn.w_msa.v_bias",
506
+ "backbone_hr.stages.3.blocks.0.ffn.layers.0.bias",
507
+ "backbone_hr.stages.3.blocks.0.ffn.layers.0.weight",
508
+ "backbone_hr.stages.3.blocks.0.ffn.layers.3.bias",
509
+ "backbone_hr.stages.3.blocks.0.ffn.layers.3.weight",
510
+ "backbone_hr.stages.3.blocks.0.norm1.bias",
511
+ "backbone_hr.stages.3.blocks.0.norm1.weight",
512
+ "backbone_hr.stages.3.blocks.0.norm2.bias",
513
+ "backbone_hr.stages.3.blocks.0.norm2.weight",
514
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.cpb_mlp.0.bias",
515
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.cpb_mlp.0.weight",
516
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.cpb_mlp.2.weight",
517
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.logit_scale",
518
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.proj.bias",
519
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.proj.weight",
520
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.q_bias",
521
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.qkv.weight",
522
+ "backbone_hr.stages.3.blocks.1.attn.w_msa.v_bias",
523
+ "backbone_hr.stages.3.blocks.1.ffn.layers.0.bias",
524
+ "backbone_hr.stages.3.blocks.1.ffn.layers.0.weight",
525
+ "backbone_hr.stages.3.blocks.1.ffn.layers.3.bias",
526
+ "backbone_hr.stages.3.blocks.1.ffn.layers.3.weight",
527
+ "backbone_hr.stages.3.blocks.1.norm1.bias",
528
+ "backbone_hr.stages.3.blocks.1.norm1.weight",
529
+ "backbone_hr.stages.3.blocks.1.norm2.bias",
530
+ "backbone_hr.stages.3.blocks.1.norm2.weight",
531
+ "backbone_hr.stages.3.downsample.norm.bias",
532
+ "backbone_hr.stages.3.downsample.norm.weight",
533
+ "backbone_hr.stages.3.downsample.reduction.weight",
534
+ "backbone_hr.vocabulary_token",
535
+ "backbone_hr.vocabulary_weight",
536
+ "backbone_s1.cls_token",
537
+ "backbone_s1.layers.0.attn.in_proj_bias",
538
+ "backbone_s1.layers.0.attn.in_proj_weight",
539
+ "backbone_s1.layers.0.attn.out_proj.bias",
540
+ "backbone_s1.layers.0.attn.out_proj.weight",
541
+ "backbone_s1.layers.0.ffn.layers.0.bias",
542
+ "backbone_s1.layers.0.ffn.layers.0.weight",
543
+ "backbone_s1.layers.0.ffn.layers.3.bias",
544
+ "backbone_s1.layers.0.ffn.layers.3.weight",
545
+ "backbone_s1.layers.0.norm1.bias",
546
+ "backbone_s1.layers.0.norm1.weight",
547
+ "backbone_s1.layers.0.norm2.bias",
548
+ "backbone_s1.layers.0.norm2.weight",
549
+ "backbone_s1.layers.1.attn.in_proj_bias",
550
+ "backbone_s1.layers.1.attn.in_proj_weight",
551
+ "backbone_s1.layers.1.attn.out_proj.bias",
552
+ "backbone_s1.layers.1.attn.out_proj.weight",
553
+ "backbone_s1.layers.1.ffn.layers.0.bias",
554
+ "backbone_s1.layers.1.ffn.layers.0.weight",
555
+ "backbone_s1.layers.1.ffn.layers.3.bias",
556
+ "backbone_s1.layers.1.ffn.layers.3.weight",
557
+ "backbone_s1.layers.1.norm1.bias",
558
+ "backbone_s1.layers.1.norm1.weight",
559
+ "backbone_s1.layers.1.norm2.bias",
560
+ "backbone_s1.layers.1.norm2.weight",
561
+ "backbone_s1.layers.10.attn.in_proj_bias",
562
+ "backbone_s1.layers.10.attn.in_proj_weight",
563
+ "backbone_s1.layers.10.attn.out_proj.bias",
564
+ "backbone_s1.layers.10.attn.out_proj.weight",
565
+ "backbone_s1.layers.10.ffn.layers.0.bias",
566
+ "backbone_s1.layers.10.ffn.layers.0.weight",
567
+ "backbone_s1.layers.10.ffn.layers.3.bias",
568
+ "backbone_s1.layers.10.ffn.layers.3.weight",
569
+ "backbone_s1.layers.10.norm1.bias",
570
+ "backbone_s1.layers.10.norm1.weight",
571
+ "backbone_s1.layers.10.norm2.bias",
572
+ "backbone_s1.layers.10.norm2.weight",
573
+ "backbone_s1.layers.11.attn.in_proj_bias",
574
+ "backbone_s1.layers.11.attn.in_proj_weight",
575
+ "backbone_s1.layers.11.attn.out_proj.bias",
576
+ "backbone_s1.layers.11.attn.out_proj.weight",
577
+ "backbone_s1.layers.11.ffn.layers.0.bias",
578
+ "backbone_s1.layers.11.ffn.layers.0.weight",
579
+ "backbone_s1.layers.11.ffn.layers.3.bias",
580
+ "backbone_s1.layers.11.ffn.layers.3.weight",
581
+ "backbone_s1.layers.11.norm1.bias",
582
+ "backbone_s1.layers.11.norm1.weight",
583
+ "backbone_s1.layers.11.norm2.bias",
584
+ "backbone_s1.layers.11.norm2.weight",
585
+ "backbone_s1.layers.12.attn.in_proj_bias",
586
+ "backbone_s1.layers.12.attn.in_proj_weight",
587
+ "backbone_s1.layers.12.attn.out_proj.bias",
588
+ "backbone_s1.layers.12.attn.out_proj.weight",
589
+ "backbone_s1.layers.12.ffn.layers.0.bias",
590
+ "backbone_s1.layers.12.ffn.layers.0.weight",
591
+ "backbone_s1.layers.12.ffn.layers.3.bias",
592
+ "backbone_s1.layers.12.ffn.layers.3.weight",
593
+ "backbone_s1.layers.12.norm1.bias",
594
+ "backbone_s1.layers.12.norm1.weight",
595
+ "backbone_s1.layers.12.norm2.bias",
596
+ "backbone_s1.layers.12.norm2.weight",
597
+ "backbone_s1.layers.13.attn.in_proj_bias",
598
+ "backbone_s1.layers.13.attn.in_proj_weight",
599
+ "backbone_s1.layers.13.attn.out_proj.bias",
600
+ "backbone_s1.layers.13.attn.out_proj.weight",
601
+ "backbone_s1.layers.13.ffn.layers.0.bias",
602
+ "backbone_s1.layers.13.ffn.layers.0.weight",
603
+ "backbone_s1.layers.13.ffn.layers.3.bias",
604
+ "backbone_s1.layers.13.ffn.layers.3.weight",
605
+ "backbone_s1.layers.13.norm1.bias",
606
+ "backbone_s1.layers.13.norm1.weight",
607
+ "backbone_s1.layers.13.norm2.bias",
608
+ "backbone_s1.layers.13.norm2.weight",
609
+ "backbone_s1.layers.14.attn.in_proj_bias",
610
+ "backbone_s1.layers.14.attn.in_proj_weight",
611
+ "backbone_s1.layers.14.attn.out_proj.bias",
612
+ "backbone_s1.layers.14.attn.out_proj.weight",
613
+ "backbone_s1.layers.14.ffn.layers.0.bias",
614
+ "backbone_s1.layers.14.ffn.layers.0.weight",
615
+ "backbone_s1.layers.14.ffn.layers.3.bias",
616
+ "backbone_s1.layers.14.ffn.layers.3.weight",
617
+ "backbone_s1.layers.14.norm1.bias",
618
+ "backbone_s1.layers.14.norm1.weight",
619
+ "backbone_s1.layers.14.norm2.bias",
620
+ "backbone_s1.layers.14.norm2.weight",
621
+ "backbone_s1.layers.15.attn.in_proj_bias",
622
+ "backbone_s1.layers.15.attn.in_proj_weight",
623
+ "backbone_s1.layers.15.attn.out_proj.bias",
624
+ "backbone_s1.layers.15.attn.out_proj.weight",
625
+ "backbone_s1.layers.15.ffn.layers.0.bias",
626
+ "backbone_s1.layers.15.ffn.layers.0.weight",
627
+ "backbone_s1.layers.15.ffn.layers.3.bias",
628
+ "backbone_s1.layers.15.ffn.layers.3.weight",
629
+ "backbone_s1.layers.15.norm1.bias",
630
+ "backbone_s1.layers.15.norm1.weight",
631
+ "backbone_s1.layers.15.norm2.bias",
632
+ "backbone_s1.layers.15.norm2.weight",
633
+ "backbone_s1.layers.16.attn.in_proj_bias",
634
+ "backbone_s1.layers.16.attn.in_proj_weight",
635
+ "backbone_s1.layers.16.attn.out_proj.bias",
636
+ "backbone_s1.layers.16.attn.out_proj.weight",
637
+ "backbone_s1.layers.16.ffn.layers.0.bias",
638
+ "backbone_s1.layers.16.ffn.layers.0.weight",
639
+ "backbone_s1.layers.16.ffn.layers.3.bias",
640
+ "backbone_s1.layers.16.ffn.layers.3.weight",
641
+ "backbone_s1.layers.16.norm1.bias",
642
+ "backbone_s1.layers.16.norm1.weight",
643
+ "backbone_s1.layers.16.norm2.bias",
644
+ "backbone_s1.layers.16.norm2.weight",
645
+ "backbone_s1.layers.17.attn.in_proj_bias",
646
+ "backbone_s1.layers.17.attn.in_proj_weight",
647
+ "backbone_s1.layers.17.attn.out_proj.bias",
648
+ "backbone_s1.layers.17.attn.out_proj.weight",
649
+ "backbone_s1.layers.17.ffn.layers.0.bias",
650
+ "backbone_s1.layers.17.ffn.layers.0.weight",
651
+ "backbone_s1.layers.17.ffn.layers.3.bias",
652
+ "backbone_s1.layers.17.ffn.layers.3.weight",
653
+ "backbone_s1.layers.17.norm1.bias",
654
+ "backbone_s1.layers.17.norm1.weight",
655
+ "backbone_s1.layers.17.norm2.bias",
656
+ "backbone_s1.layers.17.norm2.weight",
657
+ "backbone_s1.layers.18.attn.in_proj_bias",
658
+ "backbone_s1.layers.18.attn.in_proj_weight",
659
+ "backbone_s1.layers.18.attn.out_proj.bias",
660
+ "backbone_s1.layers.18.attn.out_proj.weight",
661
+ "backbone_s1.layers.18.ffn.layers.0.bias",
662
+ "backbone_s1.layers.18.ffn.layers.0.weight",
663
+ "backbone_s1.layers.18.ffn.layers.3.bias",
664
+ "backbone_s1.layers.18.ffn.layers.3.weight",
665
+ "backbone_s1.layers.18.norm1.bias",
666
+ "backbone_s1.layers.18.norm1.weight",
667
+ "backbone_s1.layers.18.norm2.bias",
668
+ "backbone_s1.layers.18.norm2.weight",
669
+ "backbone_s1.layers.19.attn.in_proj_bias",
670
+ "backbone_s1.layers.19.attn.in_proj_weight",
671
+ "backbone_s1.layers.19.attn.out_proj.bias",
672
+ "backbone_s1.layers.19.attn.out_proj.weight",
673
+ "backbone_s1.layers.19.ffn.layers.0.bias",
674
+ "backbone_s1.layers.19.ffn.layers.0.weight",
675
+ "backbone_s1.layers.19.ffn.layers.3.bias",
676
+ "backbone_s1.layers.19.ffn.layers.3.weight",
677
+ "backbone_s1.layers.19.norm1.bias",
678
+ "backbone_s1.layers.19.norm1.weight",
679
+ "backbone_s1.layers.19.norm2.bias",
680
+ "backbone_s1.layers.19.norm2.weight",
681
+ "backbone_s1.layers.2.attn.in_proj_bias",
682
+ "backbone_s1.layers.2.attn.in_proj_weight",
683
+ "backbone_s1.layers.2.attn.out_proj.bias",
684
+ "backbone_s1.layers.2.attn.out_proj.weight",
685
+ "backbone_s1.layers.2.ffn.layers.0.bias",
686
+ "backbone_s1.layers.2.ffn.layers.0.weight",
687
+ "backbone_s1.layers.2.ffn.layers.3.bias",
688
+ "backbone_s1.layers.2.ffn.layers.3.weight",
689
+ "backbone_s1.layers.2.norm1.bias",
690
+ "backbone_s1.layers.2.norm1.weight",
691
+ "backbone_s1.layers.2.norm2.bias",
692
+ "backbone_s1.layers.2.norm2.weight",
693
+ "backbone_s1.layers.20.attn.in_proj_bias",
694
+ "backbone_s1.layers.20.attn.in_proj_weight",
695
+ "backbone_s1.layers.20.attn.out_proj.bias",
696
+ "backbone_s1.layers.20.attn.out_proj.weight",
697
+ "backbone_s1.layers.20.ffn.layers.0.bias",
698
+ "backbone_s1.layers.20.ffn.layers.0.weight",
699
+ "backbone_s1.layers.20.ffn.layers.3.bias",
700
+ "backbone_s1.layers.20.ffn.layers.3.weight",
701
+ "backbone_s1.layers.20.norm1.bias",
702
+ "backbone_s1.layers.20.norm1.weight",
703
+ "backbone_s1.layers.20.norm2.bias",
704
+ "backbone_s1.layers.20.norm2.weight",
705
+ "backbone_s1.layers.21.attn.in_proj_bias",
706
+ "backbone_s1.layers.21.attn.in_proj_weight",
707
+ "backbone_s1.layers.21.attn.out_proj.bias",
708
+ "backbone_s1.layers.21.attn.out_proj.weight",
709
+ "backbone_s1.layers.21.ffn.layers.0.bias",
710
+ "backbone_s1.layers.21.ffn.layers.0.weight",
711
+ "backbone_s1.layers.21.ffn.layers.3.bias",
712
+ "backbone_s1.layers.21.ffn.layers.3.weight",
713
+ "backbone_s1.layers.21.norm1.bias",
714
+ "backbone_s1.layers.21.norm1.weight",
715
+ "backbone_s1.layers.21.norm2.bias",
716
+ "backbone_s1.layers.21.norm2.weight",
717
+ "backbone_s1.layers.22.attn.in_proj_bias",
718
+ "backbone_s1.layers.22.attn.in_proj_weight",
719
+ "backbone_s1.layers.22.attn.out_proj.bias",
720
+ "backbone_s1.layers.22.attn.out_proj.weight",
721
+ "backbone_s1.layers.22.ffn.layers.0.bias",
722
+ "backbone_s1.layers.22.ffn.layers.0.weight",
723
+ "backbone_s1.layers.22.ffn.layers.3.bias",
724
+ "backbone_s1.layers.22.ffn.layers.3.weight",
725
+ "backbone_s1.layers.22.norm1.bias",
726
+ "backbone_s1.layers.22.norm1.weight",
727
+ "backbone_s1.layers.22.norm2.bias",
728
+ "backbone_s1.layers.22.norm2.weight",
729
+ "backbone_s1.layers.23.attn.in_proj_bias",
730
+ "backbone_s1.layers.23.attn.in_proj_weight",
731
+ "backbone_s1.layers.23.attn.out_proj.bias",
732
+ "backbone_s1.layers.23.attn.out_proj.weight",
733
+ "backbone_s1.layers.23.ffn.layers.0.bias",
734
+ "backbone_s1.layers.23.ffn.layers.0.weight",
735
+ "backbone_s1.layers.23.ffn.layers.3.bias",
736
+ "backbone_s1.layers.23.ffn.layers.3.weight",
737
+ "backbone_s1.layers.23.norm1.bias",
738
+ "backbone_s1.layers.23.norm1.weight",
739
+ "backbone_s1.layers.23.norm2.bias",
740
+ "backbone_s1.layers.23.norm2.weight",
741
+ "backbone_s1.layers.3.attn.in_proj_bias",
742
+ "backbone_s1.layers.3.attn.in_proj_weight",
743
+ "backbone_s1.layers.3.attn.out_proj.bias",
744
+ "backbone_s1.layers.3.attn.out_proj.weight",
745
+ "backbone_s1.layers.3.ffn.layers.0.bias",
746
+ "backbone_s1.layers.3.ffn.layers.0.weight",
747
+ "backbone_s1.layers.3.ffn.layers.3.bias",
748
+ "backbone_s1.layers.3.ffn.layers.3.weight",
749
+ "backbone_s1.layers.3.norm1.bias",
750
+ "backbone_s1.layers.3.norm1.weight",
751
+ "backbone_s1.layers.3.norm2.bias",
752
+ "backbone_s1.layers.3.norm2.weight",
753
+ "backbone_s1.layers.4.attn.in_proj_bias",
754
+ "backbone_s1.layers.4.attn.in_proj_weight",
755
+ "backbone_s1.layers.4.attn.out_proj.bias",
756
+ "backbone_s1.layers.4.attn.out_proj.weight",
757
+ "backbone_s1.layers.4.ffn.layers.0.bias",
758
+ "backbone_s1.layers.4.ffn.layers.0.weight",
759
+ "backbone_s1.layers.4.ffn.layers.3.bias",
760
+ "backbone_s1.layers.4.ffn.layers.3.weight",
761
+ "backbone_s1.layers.4.norm1.bias",
762
+ "backbone_s1.layers.4.norm1.weight",
763
+ "backbone_s1.layers.4.norm2.bias",
764
+ "backbone_s1.layers.4.norm2.weight",
765
+ "backbone_s1.layers.5.attn.in_proj_bias",
766
+ "backbone_s1.layers.5.attn.in_proj_weight",
767
+ "backbone_s1.layers.5.attn.out_proj.bias",
768
+ "backbone_s1.layers.5.attn.out_proj.weight",
769
+ "backbone_s1.layers.5.ffn.layers.0.bias",
770
+ "backbone_s1.layers.5.ffn.layers.0.weight",
771
+ "backbone_s1.layers.5.ffn.layers.3.bias",
772
+ "backbone_s1.layers.5.ffn.layers.3.weight",
773
+ "backbone_s1.layers.5.norm1.bias",
774
+ "backbone_s1.layers.5.norm1.weight",
775
+ "backbone_s1.layers.5.norm2.bias",
776
+ "backbone_s1.layers.5.norm2.weight",
777
+ "backbone_s1.layers.6.attn.in_proj_bias",
778
+ "backbone_s1.layers.6.attn.in_proj_weight",
779
+ "backbone_s1.layers.6.attn.out_proj.bias",
780
+ "backbone_s1.layers.6.attn.out_proj.weight",
781
+ "backbone_s1.layers.6.ffn.layers.0.bias",
782
+ "backbone_s1.layers.6.ffn.layers.0.weight",
783
+ "backbone_s1.layers.6.ffn.layers.3.bias",
784
+ "backbone_s1.layers.6.ffn.layers.3.weight",
785
+ "backbone_s1.layers.6.norm1.bias",
786
+ "backbone_s1.layers.6.norm1.weight",
787
+ "backbone_s1.layers.6.norm2.bias",
788
+ "backbone_s1.layers.6.norm2.weight",
789
+ "backbone_s1.layers.7.attn.in_proj_bias",
790
+ "backbone_s1.layers.7.attn.in_proj_weight",
791
+ "backbone_s1.layers.7.attn.out_proj.bias",
792
+ "backbone_s1.layers.7.attn.out_proj.weight",
793
+ "backbone_s1.layers.7.ffn.layers.0.bias",
794
+ "backbone_s1.layers.7.ffn.layers.0.weight",
795
+ "backbone_s1.layers.7.ffn.layers.3.bias",
796
+ "backbone_s1.layers.7.ffn.layers.3.weight",
797
+ "backbone_s1.layers.7.norm1.bias",
798
+ "backbone_s1.layers.7.norm1.weight",
799
+ "backbone_s1.layers.7.norm2.bias",
800
+ "backbone_s1.layers.7.norm2.weight",
801
+ "backbone_s1.layers.8.attn.in_proj_bias",
802
+ "backbone_s1.layers.8.attn.in_proj_weight",
803
+ "backbone_s1.layers.8.attn.out_proj.bias",
804
+ "backbone_s1.layers.8.attn.out_proj.weight",
805
+ "backbone_s1.layers.8.ffn.layers.0.bias",
806
+ "backbone_s1.layers.8.ffn.layers.0.weight",
807
+ "backbone_s1.layers.8.ffn.layers.3.bias",
808
+ "backbone_s1.layers.8.ffn.layers.3.weight",
809
+ "backbone_s1.layers.8.norm1.bias",
810
+ "backbone_s1.layers.8.norm1.weight",
811
+ "backbone_s1.layers.8.norm2.bias",
812
+ "backbone_s1.layers.8.norm2.weight",
813
+ "backbone_s1.layers.9.attn.in_proj_bias",
814
+ "backbone_s1.layers.9.attn.in_proj_weight",
815
+ "backbone_s1.layers.9.attn.out_proj.bias",
816
+ "backbone_s1.layers.9.attn.out_proj.weight",
817
+ "backbone_s1.layers.9.ffn.layers.0.bias",
818
+ "backbone_s1.layers.9.ffn.layers.0.weight",
819
+ "backbone_s1.layers.9.ffn.layers.3.bias",
820
+ "backbone_s1.layers.9.ffn.layers.3.weight",
821
+ "backbone_s1.layers.9.norm1.bias",
822
+ "backbone_s1.layers.9.norm1.weight",
823
+ "backbone_s1.layers.9.norm2.bias",
824
+ "backbone_s1.layers.9.norm2.weight",
825
+ "backbone_s1.mask_token",
826
+ "backbone_s1.patch_embed.projection.bias",
827
+ "backbone_s1.patch_embed.projection.weight",
828
+ "backbone_s1.pos_embed",
829
+ "backbone_s1.vocabulary_token",
830
+ "backbone_s1.vocabulary_weight",
831
+ "backbone_s2.cls_token",
832
+ "backbone_s2.layers.0.attn.in_proj_bias",
833
+ "backbone_s2.layers.0.attn.in_proj_weight",
834
+ "backbone_s2.layers.0.attn.out_proj.bias",
835
+ "backbone_s2.layers.0.attn.out_proj.weight",
836
+ "backbone_s2.layers.0.ffn.layers.0.bias",
837
+ "backbone_s2.layers.0.ffn.layers.0.weight",
838
+ "backbone_s2.layers.0.ffn.layers.3.bias",
839
+ "backbone_s2.layers.0.ffn.layers.3.weight",
840
+ "backbone_s2.layers.0.norm1.bias",
841
+ "backbone_s2.layers.0.norm1.weight",
842
+ "backbone_s2.layers.0.norm2.bias",
843
+ "backbone_s2.layers.0.norm2.weight",
844
+ "backbone_s2.layers.1.attn.in_proj_bias",
845
+ "backbone_s2.layers.1.attn.in_proj_weight",
846
+ "backbone_s2.layers.1.attn.out_proj.bias",
847
+ "backbone_s2.layers.1.attn.out_proj.weight",
848
+ "backbone_s2.layers.1.ffn.layers.0.bias",
849
+ "backbone_s2.layers.1.ffn.layers.0.weight",
850
+ "backbone_s2.layers.1.ffn.layers.3.bias",
851
+ "backbone_s2.layers.1.ffn.layers.3.weight",
852
+ "backbone_s2.layers.1.norm1.bias",
853
+ "backbone_s2.layers.1.norm1.weight",
854
+ "backbone_s2.layers.1.norm2.bias",
855
+ "backbone_s2.layers.1.norm2.weight",
856
+ "backbone_s2.layers.10.attn.in_proj_bias",
857
+ "backbone_s2.layers.10.attn.in_proj_weight",
858
+ "backbone_s2.layers.10.attn.out_proj.bias",
859
+ "backbone_s2.layers.10.attn.out_proj.weight",
860
+ "backbone_s2.layers.10.ffn.layers.0.bias",
861
+ "backbone_s2.layers.10.ffn.layers.0.weight",
862
+ "backbone_s2.layers.10.ffn.layers.3.bias",
863
+ "backbone_s2.layers.10.ffn.layers.3.weight",
864
+ "backbone_s2.layers.10.norm1.bias",
865
+ "backbone_s2.layers.10.norm1.weight",
866
+ "backbone_s2.layers.10.norm2.bias",
867
+ "backbone_s2.layers.10.norm2.weight",
868
+ "backbone_s2.layers.11.attn.in_proj_bias",
869
+ "backbone_s2.layers.11.attn.in_proj_weight",
870
+ "backbone_s2.layers.11.attn.out_proj.bias",
871
+ "backbone_s2.layers.11.attn.out_proj.weight",
872
+ "backbone_s2.layers.11.ffn.layers.0.bias",
873
+ "backbone_s2.layers.11.ffn.layers.0.weight",
874
+ "backbone_s2.layers.11.ffn.layers.3.bias",
875
+ "backbone_s2.layers.11.ffn.layers.3.weight",
876
+ "backbone_s2.layers.11.norm1.bias",
877
+ "backbone_s2.layers.11.norm1.weight",
878
+ "backbone_s2.layers.11.norm2.bias",
879
+ "backbone_s2.layers.11.norm2.weight",
880
+ "backbone_s2.layers.12.attn.in_proj_bias",
881
+ "backbone_s2.layers.12.attn.in_proj_weight",
882
+ "backbone_s2.layers.12.attn.out_proj.bias",
883
+ "backbone_s2.layers.12.attn.out_proj.weight",
884
+ "backbone_s2.layers.12.ffn.layers.0.bias",
885
+ "backbone_s2.layers.12.ffn.layers.0.weight",
886
+ "backbone_s2.layers.12.ffn.layers.3.bias",
887
+ "backbone_s2.layers.12.ffn.layers.3.weight",
888
+ "backbone_s2.layers.12.norm1.bias",
889
+ "backbone_s2.layers.12.norm1.weight",
890
+ "backbone_s2.layers.12.norm2.bias",
891
+ "backbone_s2.layers.12.norm2.weight",
892
+ "backbone_s2.layers.13.attn.in_proj_bias",
893
+ "backbone_s2.layers.13.attn.in_proj_weight",
894
+ "backbone_s2.layers.13.attn.out_proj.bias",
895
+ "backbone_s2.layers.13.attn.out_proj.weight",
896
+ "backbone_s2.layers.13.ffn.layers.0.bias",
897
+ "backbone_s2.layers.13.ffn.layers.0.weight",
898
+ "backbone_s2.layers.13.ffn.layers.3.bias",
899
+ "backbone_s2.layers.13.ffn.layers.3.weight",
900
+ "backbone_s2.layers.13.norm1.bias",
901
+ "backbone_s2.layers.13.norm1.weight",
902
+ "backbone_s2.layers.13.norm2.bias",
903
+ "backbone_s2.layers.13.norm2.weight",
904
+ "backbone_s2.layers.14.attn.in_proj_bias",
905
+ "backbone_s2.layers.14.attn.in_proj_weight",
906
+ "backbone_s2.layers.14.attn.out_proj.bias",
907
+ "backbone_s2.layers.14.attn.out_proj.weight",
908
+ "backbone_s2.layers.14.ffn.layers.0.bias",
909
+ "backbone_s2.layers.14.ffn.layers.0.weight",
910
+ "backbone_s2.layers.14.ffn.layers.3.bias",
911
+ "backbone_s2.layers.14.ffn.layers.3.weight",
912
+ "backbone_s2.layers.14.norm1.bias",
913
+ "backbone_s2.layers.14.norm1.weight",
914
+ "backbone_s2.layers.14.norm2.bias",
915
+ "backbone_s2.layers.14.norm2.weight",
916
+ "backbone_s2.layers.15.attn.in_proj_bias",
917
+ "backbone_s2.layers.15.attn.in_proj_weight",
918
+ "backbone_s2.layers.15.attn.out_proj.bias",
919
+ "backbone_s2.layers.15.attn.out_proj.weight",
920
+ "backbone_s2.layers.15.ffn.layers.0.bias",
921
+ "backbone_s2.layers.15.ffn.layers.0.weight",
922
+ "backbone_s2.layers.15.ffn.layers.3.bias",
923
+ "backbone_s2.layers.15.ffn.layers.3.weight",
924
+ "backbone_s2.layers.15.norm1.bias",
925
+ "backbone_s2.layers.15.norm1.weight",
926
+ "backbone_s2.layers.15.norm2.bias",
927
+ "backbone_s2.layers.15.norm2.weight",
928
+ "backbone_s2.layers.16.attn.in_proj_bias",
929
+ "backbone_s2.layers.16.attn.in_proj_weight",
930
+ "backbone_s2.layers.16.attn.out_proj.bias",
931
+ "backbone_s2.layers.16.attn.out_proj.weight",
932
+ "backbone_s2.layers.16.ffn.layers.0.bias",
933
+ "backbone_s2.layers.16.ffn.layers.0.weight",
934
+ "backbone_s2.layers.16.ffn.layers.3.bias",
935
+ "backbone_s2.layers.16.ffn.layers.3.weight",
936
+ "backbone_s2.layers.16.norm1.bias",
937
+ "backbone_s2.layers.16.norm1.weight",
938
+ "backbone_s2.layers.16.norm2.bias",
939
+ "backbone_s2.layers.16.norm2.weight",
940
+ "backbone_s2.layers.17.attn.in_proj_bias",
941
+ "backbone_s2.layers.17.attn.in_proj_weight",
942
+ "backbone_s2.layers.17.attn.out_proj.bias",
943
+ "backbone_s2.layers.17.attn.out_proj.weight",
944
+ "backbone_s2.layers.17.ffn.layers.0.bias",
945
+ "backbone_s2.layers.17.ffn.layers.0.weight",
946
+ "backbone_s2.layers.17.ffn.layers.3.bias",
947
+ "backbone_s2.layers.17.ffn.layers.3.weight",
948
+ "backbone_s2.layers.17.norm1.bias",
949
+ "backbone_s2.layers.17.norm1.weight",
950
+ "backbone_s2.layers.17.norm2.bias",
951
+ "backbone_s2.layers.17.norm2.weight",
952
+ "backbone_s2.layers.18.attn.in_proj_bias",
953
+ "backbone_s2.layers.18.attn.in_proj_weight",
954
+ "backbone_s2.layers.18.attn.out_proj.bias",
955
+ "backbone_s2.layers.18.attn.out_proj.weight",
956
+ "backbone_s2.layers.18.ffn.layers.0.bias",
957
+ "backbone_s2.layers.18.ffn.layers.0.weight",
958
+ "backbone_s2.layers.18.ffn.layers.3.bias",
959
+ "backbone_s2.layers.18.ffn.layers.3.weight",
960
+ "backbone_s2.layers.18.norm1.bias",
961
+ "backbone_s2.layers.18.norm1.weight",
962
+ "backbone_s2.layers.18.norm2.bias",
963
+ "backbone_s2.layers.18.norm2.weight",
964
+ "backbone_s2.layers.19.attn.in_proj_bias",
965
+ "backbone_s2.layers.19.attn.in_proj_weight",
966
+ "backbone_s2.layers.19.attn.out_proj.bias",
967
+ "backbone_s2.layers.19.attn.out_proj.weight",
968
+ "backbone_s2.layers.19.ffn.layers.0.bias",
969
+ "backbone_s2.layers.19.ffn.layers.0.weight",
970
+ "backbone_s2.layers.19.ffn.layers.3.bias",
971
+ "backbone_s2.layers.19.ffn.layers.3.weight",
972
+ "backbone_s2.layers.19.norm1.bias",
973
+ "backbone_s2.layers.19.norm1.weight",
974
+ "backbone_s2.layers.19.norm2.bias",
975
+ "backbone_s2.layers.19.norm2.weight",
976
+ "backbone_s2.layers.2.attn.in_proj_bias",
977
+ "backbone_s2.layers.2.attn.in_proj_weight",
978
+ "backbone_s2.layers.2.attn.out_proj.bias",
979
+ "backbone_s2.layers.2.attn.out_proj.weight",
980
+ "backbone_s2.layers.2.ffn.layers.0.bias",
981
+ "backbone_s2.layers.2.ffn.layers.0.weight",
982
+ "backbone_s2.layers.2.ffn.layers.3.bias",
983
+ "backbone_s2.layers.2.ffn.layers.3.weight",
984
+ "backbone_s2.layers.2.norm1.bias",
985
+ "backbone_s2.layers.2.norm1.weight",
986
+ "backbone_s2.layers.2.norm2.bias",
987
+ "backbone_s2.layers.2.norm2.weight",
988
+ "backbone_s2.layers.20.attn.in_proj_bias",
989
+ "backbone_s2.layers.20.attn.in_proj_weight",
990
+ "backbone_s2.layers.20.attn.out_proj.bias",
991
+ "backbone_s2.layers.20.attn.out_proj.weight",
992
+ "backbone_s2.layers.20.ffn.layers.0.bias",
993
+ "backbone_s2.layers.20.ffn.layers.0.weight",
994
+ "backbone_s2.layers.20.ffn.layers.3.bias",
995
+ "backbone_s2.layers.20.ffn.layers.3.weight",
996
+ "backbone_s2.layers.20.norm1.bias",
997
+ "backbone_s2.layers.20.norm1.weight",
998
+ "backbone_s2.layers.20.norm2.bias",
999
+ "backbone_s2.layers.20.norm2.weight",
1000
+ "backbone_s2.layers.21.attn.in_proj_bias",
1001
+ "backbone_s2.layers.21.attn.in_proj_weight",
1002
+ "backbone_s2.layers.21.attn.out_proj.bias",
1003
+ "backbone_s2.layers.21.attn.out_proj.weight",
1004
+ "backbone_s2.layers.21.ffn.layers.0.bias",
1005
+ "backbone_s2.layers.21.ffn.layers.0.weight",
1006
+ "backbone_s2.layers.21.ffn.layers.3.bias",
1007
+ "backbone_s2.layers.21.ffn.layers.3.weight",
1008
+ "backbone_s2.layers.21.norm1.bias",
1009
+ "backbone_s2.layers.21.norm1.weight",
1010
+ "backbone_s2.layers.21.norm2.bias",
1011
+ "backbone_s2.layers.21.norm2.weight",
1012
+ "backbone_s2.layers.22.attn.in_proj_bias",
1013
+ "backbone_s2.layers.22.attn.in_proj_weight",
1014
+ "backbone_s2.layers.22.attn.out_proj.bias",
1015
+ "backbone_s2.layers.22.attn.out_proj.weight",
1016
+ "backbone_s2.layers.22.ffn.layers.0.bias",
1017
+ "backbone_s2.layers.22.ffn.layers.0.weight",
1018
+ "backbone_s2.layers.22.ffn.layers.3.bias",
1019
+ "backbone_s2.layers.22.ffn.layers.3.weight",
1020
+ "backbone_s2.layers.22.norm1.bias",
1021
+ "backbone_s2.layers.22.norm1.weight",
1022
+ "backbone_s2.layers.22.norm2.bias",
1023
+ "backbone_s2.layers.22.norm2.weight",
1024
+ "backbone_s2.layers.23.attn.in_proj_bias",
1025
+ "backbone_s2.layers.23.attn.in_proj_weight",
1026
+ "backbone_s2.layers.23.attn.out_proj.bias",
1027
+ "backbone_s2.layers.23.attn.out_proj.weight",
1028
+ "backbone_s2.layers.23.ffn.layers.0.bias",
1029
+ "backbone_s2.layers.23.ffn.layers.0.weight",
1030
+ "backbone_s2.layers.23.ffn.layers.3.bias",
1031
+ "backbone_s2.layers.23.ffn.layers.3.weight",
1032
+ "backbone_s2.layers.23.norm1.bias",
1033
+ "backbone_s2.layers.23.norm1.weight",
1034
+ "backbone_s2.layers.23.norm2.bias",
1035
+ "backbone_s2.layers.23.norm2.weight",
1036
+ "backbone_s2.layers.3.attn.in_proj_bias",
1037
+ "backbone_s2.layers.3.attn.in_proj_weight",
1038
+ "backbone_s2.layers.3.attn.out_proj.bias",
1039
+ "backbone_s2.layers.3.attn.out_proj.weight",
1040
+ "backbone_s2.layers.3.ffn.layers.0.bias",
1041
+ "backbone_s2.layers.3.ffn.layers.0.weight",
1042
+ "backbone_s2.layers.3.ffn.layers.3.bias",
1043
+ "backbone_s2.layers.3.ffn.layers.3.weight",
1044
+ "backbone_s2.layers.3.norm1.bias",
1045
+ "backbone_s2.layers.3.norm1.weight",
1046
+ "backbone_s2.layers.3.norm2.bias",
1047
+ "backbone_s2.layers.3.norm2.weight",
1048
+ "backbone_s2.layers.4.attn.in_proj_bias",
1049
+ "backbone_s2.layers.4.attn.in_proj_weight",
1050
+ "backbone_s2.layers.4.attn.out_proj.bias",
1051
+ "backbone_s2.layers.4.attn.out_proj.weight",
1052
+ "backbone_s2.layers.4.ffn.layers.0.bias",
1053
+ "backbone_s2.layers.4.ffn.layers.0.weight",
1054
+ "backbone_s2.layers.4.ffn.layers.3.bias",
1055
+ "backbone_s2.layers.4.ffn.layers.3.weight",
1056
+ "backbone_s2.layers.4.norm1.bias",
1057
+ "backbone_s2.layers.4.norm1.weight",
1058
+ "backbone_s2.layers.4.norm2.bias",
1059
+ "backbone_s2.layers.4.norm2.weight",
1060
+ "backbone_s2.layers.5.attn.in_proj_bias",
1061
+ "backbone_s2.layers.5.attn.in_proj_weight",
1062
+ "backbone_s2.layers.5.attn.out_proj.bias",
1063
+ "backbone_s2.layers.5.attn.out_proj.weight",
1064
+ "backbone_s2.layers.5.ffn.layers.0.bias",
1065
+ "backbone_s2.layers.5.ffn.layers.0.weight",
1066
+ "backbone_s2.layers.5.ffn.layers.3.bias",
1067
+ "backbone_s2.layers.5.ffn.layers.3.weight",
1068
+ "backbone_s2.layers.5.norm1.bias",
1069
+ "backbone_s2.layers.5.norm1.weight",
1070
+ "backbone_s2.layers.5.norm2.bias",
1071
+ "backbone_s2.layers.5.norm2.weight",
1072
+ "backbone_s2.layers.6.attn.in_proj_bias",
1073
+ "backbone_s2.layers.6.attn.in_proj_weight",
1074
+ "backbone_s2.layers.6.attn.out_proj.bias",
1075
+ "backbone_s2.layers.6.attn.out_proj.weight",
1076
+ "backbone_s2.layers.6.ffn.layers.0.bias",
1077
+ "backbone_s2.layers.6.ffn.layers.0.weight",
1078
+ "backbone_s2.layers.6.ffn.layers.3.bias",
1079
+ "backbone_s2.layers.6.ffn.layers.3.weight",
1080
+ "backbone_s2.layers.6.norm1.bias",
1081
+ "backbone_s2.layers.6.norm1.weight",
1082
+ "backbone_s2.layers.6.norm2.bias",
1083
+ "backbone_s2.layers.6.norm2.weight",
1084
+ "backbone_s2.layers.7.attn.in_proj_bias",
1085
+ "backbone_s2.layers.7.attn.in_proj_weight",
1086
+ "backbone_s2.layers.7.attn.out_proj.bias",
1087
+ "backbone_s2.layers.7.attn.out_proj.weight",
1088
+ "backbone_s2.layers.7.ffn.layers.0.bias",
1089
+ "backbone_s2.layers.7.ffn.layers.0.weight",
1090
+ "backbone_s2.layers.7.ffn.layers.3.bias",
1091
+ "backbone_s2.layers.7.ffn.layers.3.weight",
1092
+ "backbone_s2.layers.7.norm1.bias",
1093
+ "backbone_s2.layers.7.norm1.weight",
1094
+ "backbone_s2.layers.7.norm2.bias",
1095
+ "backbone_s2.layers.7.norm2.weight",
1096
+ "backbone_s2.layers.8.attn.in_proj_bias",
1097
+ "backbone_s2.layers.8.attn.in_proj_weight",
1098
+ "backbone_s2.layers.8.attn.out_proj.bias",
1099
+ "backbone_s2.layers.8.attn.out_proj.weight",
1100
+ "backbone_s2.layers.8.ffn.layers.0.bias",
1101
+ "backbone_s2.layers.8.ffn.layers.0.weight",
1102
+ "backbone_s2.layers.8.ffn.layers.3.bias",
1103
+ "backbone_s2.layers.8.ffn.layers.3.weight",
1104
+ "backbone_s2.layers.8.norm1.bias",
1105
+ "backbone_s2.layers.8.norm1.weight",
1106
+ "backbone_s2.layers.8.norm2.bias",
1107
+ "backbone_s2.layers.8.norm2.weight",
1108
+ "backbone_s2.layers.9.attn.in_proj_bias",
1109
+ "backbone_s2.layers.9.attn.in_proj_weight",
1110
+ "backbone_s2.layers.9.attn.out_proj.bias",
1111
+ "backbone_s2.layers.9.attn.out_proj.weight",
1112
+ "backbone_s2.layers.9.ffn.layers.0.bias",
1113
+ "backbone_s2.layers.9.ffn.layers.0.weight",
1114
+ "backbone_s2.layers.9.ffn.layers.3.bias",
1115
+ "backbone_s2.layers.9.ffn.layers.3.weight",
1116
+ "backbone_s2.layers.9.norm1.bias",
1117
+ "backbone_s2.layers.9.norm1.weight",
1118
+ "backbone_s2.layers.9.norm2.bias",
1119
+ "backbone_s2.layers.9.norm2.weight",
1120
+ "backbone_s2.mask_token",
1121
+ "backbone_s2.patch_embed.projection.bias",
1122
+ "backbone_s2.patch_embed.projection.weight",
1123
+ "backbone_s2.pos_embed",
1124
+ "backbone_s2.vocabulary_token",
1125
+ "backbone_s2.vocabulary_weight",
1126
+ "fusion.cls_token",
1127
+ "fusion.layers.0.attn.in_proj_bias",
1128
+ "fusion.layers.0.attn.in_proj_weight",
1129
+ "fusion.layers.0.attn.out_proj.bias",
1130
+ "fusion.layers.0.attn.out_proj.weight",
1131
+ "fusion.layers.0.ffn.layers.0.bias",
1132
+ "fusion.layers.0.ffn.layers.0.weight",
1133
+ "fusion.layers.0.ffn.layers.3.bias",
1134
+ "fusion.layers.0.ffn.layers.3.weight",
1135
+ "fusion.layers.0.norm1.bias",
1136
+ "fusion.layers.0.norm1.weight",
1137
+ "fusion.layers.0.norm2.bias",
1138
+ "fusion.layers.0.norm2.weight",
1139
+ "fusion.layers.1.attn.in_proj_bias",
1140
+ "fusion.layers.1.attn.in_proj_weight",
1141
+ "fusion.layers.1.attn.out_proj.bias",
1142
+ "fusion.layers.1.attn.out_proj.weight",
1143
+ "fusion.layers.1.ffn.layers.0.bias",
1144
+ "fusion.layers.1.ffn.layers.0.weight",
1145
+ "fusion.layers.1.ffn.layers.3.bias",
1146
+ "fusion.layers.1.ffn.layers.3.weight",
1147
+ "fusion.layers.1.norm1.bias",
1148
+ "fusion.layers.1.norm1.weight",
1149
+ "fusion.layers.1.norm2.bias",
1150
+ "fusion.layers.1.norm2.weight",
1151
+ "fusion.layers.10.attn.in_proj_bias",
1152
+ "fusion.layers.10.attn.in_proj_weight",
1153
+ "fusion.layers.10.attn.out_proj.bias",
1154
+ "fusion.layers.10.attn.out_proj.weight",
1155
+ "fusion.layers.10.ffn.layers.0.bias",
1156
+ "fusion.layers.10.ffn.layers.0.weight",
1157
+ "fusion.layers.10.ffn.layers.3.bias",
1158
+ "fusion.layers.10.ffn.layers.3.weight",
1159
+ "fusion.layers.10.norm1.bias",
1160
+ "fusion.layers.10.norm1.weight",
1161
+ "fusion.layers.10.norm2.bias",
1162
+ "fusion.layers.10.norm2.weight",
1163
+ "fusion.layers.11.attn.in_proj_bias",
1164
+ "fusion.layers.11.attn.in_proj_weight",
1165
+ "fusion.layers.11.attn.out_proj.bias",
1166
+ "fusion.layers.11.attn.out_proj.weight",
1167
+ "fusion.layers.11.ffn.layers.0.bias",
1168
+ "fusion.layers.11.ffn.layers.0.weight",
1169
+ "fusion.layers.11.ffn.layers.3.bias",
1170
+ "fusion.layers.11.ffn.layers.3.weight",
1171
+ "fusion.layers.11.norm1.bias",
1172
+ "fusion.layers.11.norm1.weight",
1173
+ "fusion.layers.11.norm2.bias",
1174
+ "fusion.layers.11.norm2.weight",
1175
+ "fusion.layers.12.attn.in_proj_bias",
1176
+ "fusion.layers.12.attn.in_proj_weight",
1177
+ "fusion.layers.12.attn.out_proj.bias",
1178
+ "fusion.layers.12.attn.out_proj.weight",
1179
+ "fusion.layers.12.ffn.layers.0.bias",
1180
+ "fusion.layers.12.ffn.layers.0.weight",
1181
+ "fusion.layers.12.ffn.layers.3.bias",
1182
+ "fusion.layers.12.ffn.layers.3.weight",
1183
+ "fusion.layers.12.norm1.bias",
1184
+ "fusion.layers.12.norm1.weight",
1185
+ "fusion.layers.12.norm2.bias",
1186
+ "fusion.layers.12.norm2.weight",
1187
+ "fusion.layers.13.attn.in_proj_bias",
1188
+ "fusion.layers.13.attn.in_proj_weight",
1189
+ "fusion.layers.13.attn.out_proj.bias",
1190
+ "fusion.layers.13.attn.out_proj.weight",
1191
+ "fusion.layers.13.ffn.layers.0.bias",
1192
+ "fusion.layers.13.ffn.layers.0.weight",
1193
+ "fusion.layers.13.ffn.layers.3.bias",
1194
+ "fusion.layers.13.ffn.layers.3.weight",
1195
+ "fusion.layers.13.norm1.bias",
1196
+ "fusion.layers.13.norm1.weight",
1197
+ "fusion.layers.13.norm2.bias",
1198
+ "fusion.layers.13.norm2.weight",
1199
+ "fusion.layers.14.attn.in_proj_bias",
1200
+ "fusion.layers.14.attn.in_proj_weight",
1201
+ "fusion.layers.14.attn.out_proj.bias",
1202
+ "fusion.layers.14.attn.out_proj.weight",
1203
+ "fusion.layers.14.ffn.layers.0.bias",
1204
+ "fusion.layers.14.ffn.layers.0.weight",
1205
+ "fusion.layers.14.ffn.layers.3.bias",
1206
+ "fusion.layers.14.ffn.layers.3.weight",
1207
+ "fusion.layers.14.norm1.bias",
1208
+ "fusion.layers.14.norm1.weight",
1209
+ "fusion.layers.14.norm2.bias",
1210
+ "fusion.layers.14.norm2.weight",
1211
+ "fusion.layers.15.attn.in_proj_bias",
1212
+ "fusion.layers.15.attn.in_proj_weight",
1213
+ "fusion.layers.15.attn.out_proj.bias",
1214
+ "fusion.layers.15.attn.out_proj.weight",
1215
+ "fusion.layers.15.ffn.layers.0.bias",
1216
+ "fusion.layers.15.ffn.layers.0.weight",
1217
+ "fusion.layers.15.ffn.layers.3.bias",
1218
+ "fusion.layers.15.ffn.layers.3.weight",
1219
+ "fusion.layers.15.norm1.bias",
1220
+ "fusion.layers.15.norm1.weight",
1221
+ "fusion.layers.15.norm2.bias",
1222
+ "fusion.layers.15.norm2.weight",
1223
+ "fusion.layers.16.attn.in_proj_bias",
1224
+ "fusion.layers.16.attn.in_proj_weight",
1225
+ "fusion.layers.16.attn.out_proj.bias",
1226
+ "fusion.layers.16.attn.out_proj.weight",
1227
+ "fusion.layers.16.ffn.layers.0.bias",
1228
+ "fusion.layers.16.ffn.layers.0.weight",
1229
+ "fusion.layers.16.ffn.layers.3.bias",
1230
+ "fusion.layers.16.ffn.layers.3.weight",
1231
+ "fusion.layers.16.norm1.bias",
1232
+ "fusion.layers.16.norm1.weight",
1233
+ "fusion.layers.16.norm2.bias",
1234
+ "fusion.layers.16.norm2.weight",
1235
+ "fusion.layers.17.attn.in_proj_bias",
1236
+ "fusion.layers.17.attn.in_proj_weight",
1237
+ "fusion.layers.17.attn.out_proj.bias",
1238
+ "fusion.layers.17.attn.out_proj.weight",
1239
+ "fusion.layers.17.ffn.layers.0.bias",
1240
+ "fusion.layers.17.ffn.layers.0.weight",
1241
+ "fusion.layers.17.ffn.layers.3.bias",
1242
+ "fusion.layers.17.ffn.layers.3.weight",
1243
+ "fusion.layers.17.norm1.bias",
1244
+ "fusion.layers.17.norm1.weight",
1245
+ "fusion.layers.17.norm2.bias",
1246
+ "fusion.layers.17.norm2.weight",
1247
+ "fusion.layers.18.attn.in_proj_bias",
1248
+ "fusion.layers.18.attn.in_proj_weight",
1249
+ "fusion.layers.18.attn.out_proj.bias",
1250
+ "fusion.layers.18.attn.out_proj.weight",
1251
+ "fusion.layers.18.ffn.layers.0.bias",
1252
+ "fusion.layers.18.ffn.layers.0.weight",
1253
+ "fusion.layers.18.ffn.layers.3.bias",
1254
+ "fusion.layers.18.ffn.layers.3.weight",
1255
+ "fusion.layers.18.norm1.bias",
1256
+ "fusion.layers.18.norm1.weight",
1257
+ "fusion.layers.18.norm2.bias",
1258
+ "fusion.layers.18.norm2.weight",
1259
+ "fusion.layers.19.attn.in_proj_bias",
1260
+ "fusion.layers.19.attn.in_proj_weight",
1261
+ "fusion.layers.19.attn.out_proj.bias",
1262
+ "fusion.layers.19.attn.out_proj.weight",
1263
+ "fusion.layers.19.ffn.layers.0.bias",
1264
+ "fusion.layers.19.ffn.layers.0.weight",
1265
+ "fusion.layers.19.ffn.layers.3.bias",
1266
+ "fusion.layers.19.ffn.layers.3.weight",
1267
+ "fusion.layers.19.norm1.bias",
1268
+ "fusion.layers.19.norm1.weight",
1269
+ "fusion.layers.19.norm2.bias",
1270
+ "fusion.layers.19.norm2.weight",
1271
+ "fusion.layers.2.attn.in_proj_bias",
1272
+ "fusion.layers.2.attn.in_proj_weight",
1273
+ "fusion.layers.2.attn.out_proj.bias",
1274
+ "fusion.layers.2.attn.out_proj.weight",
1275
+ "fusion.layers.2.ffn.layers.0.bias",
1276
+ "fusion.layers.2.ffn.layers.0.weight",
1277
+ "fusion.layers.2.ffn.layers.3.bias",
1278
+ "fusion.layers.2.ffn.layers.3.weight",
1279
+ "fusion.layers.2.norm1.bias",
1280
+ "fusion.layers.2.norm1.weight",
1281
+ "fusion.layers.2.norm2.bias",
1282
+ "fusion.layers.2.norm2.weight",
1283
+ "fusion.layers.20.attn.in_proj_bias",
1284
+ "fusion.layers.20.attn.in_proj_weight",
1285
+ "fusion.layers.20.attn.out_proj.bias",
1286
+ "fusion.layers.20.attn.out_proj.weight",
1287
+ "fusion.layers.20.ffn.layers.0.bias",
1288
+ "fusion.layers.20.ffn.layers.0.weight",
1289
+ "fusion.layers.20.ffn.layers.3.bias",
1290
+ "fusion.layers.20.ffn.layers.3.weight",
1291
+ "fusion.layers.20.norm1.bias",
1292
+ "fusion.layers.20.norm1.weight",
1293
+ "fusion.layers.20.norm2.bias",
1294
+ "fusion.layers.20.norm2.weight",
1295
+ "fusion.layers.21.attn.in_proj_bias",
1296
+ "fusion.layers.21.attn.in_proj_weight",
1297
+ "fusion.layers.21.attn.out_proj.bias",
1298
+ "fusion.layers.21.attn.out_proj.weight",
1299
+ "fusion.layers.21.ffn.layers.0.bias",
1300
+ "fusion.layers.21.ffn.layers.0.weight",
1301
+ "fusion.layers.21.ffn.layers.3.bias",
1302
+ "fusion.layers.21.ffn.layers.3.weight",
1303
+ "fusion.layers.21.norm1.bias",
1304
+ "fusion.layers.21.norm1.weight",
1305
+ "fusion.layers.21.norm2.bias",
1306
+ "fusion.layers.21.norm2.weight",
1307
+ "fusion.layers.22.attn.in_proj_bias",
1308
+ "fusion.layers.22.attn.in_proj_weight",
1309
+ "fusion.layers.22.attn.out_proj.bias",
1310
+ "fusion.layers.22.attn.out_proj.weight",
1311
+ "fusion.layers.22.ffn.layers.0.bias",
1312
+ "fusion.layers.22.ffn.layers.0.weight",
1313
+ "fusion.layers.22.ffn.layers.3.bias",
1314
+ "fusion.layers.22.ffn.layers.3.weight",
1315
+ "fusion.layers.22.norm1.bias",
1316
+ "fusion.layers.22.norm1.weight",
1317
+ "fusion.layers.22.norm2.bias",
1318
+ "fusion.layers.22.norm2.weight",
1319
+ "fusion.layers.23.attn.in_proj_bias",
1320
+ "fusion.layers.23.attn.in_proj_weight",
1321
+ "fusion.layers.23.attn.out_proj.bias",
1322
+ "fusion.layers.23.attn.out_proj.weight",
1323
+ "fusion.layers.23.ffn.layers.0.bias",
1324
+ "fusion.layers.23.ffn.layers.0.weight",
1325
+ "fusion.layers.23.ffn.layers.3.bias",
1326
+ "fusion.layers.23.ffn.layers.3.weight",
1327
+ "fusion.layers.23.norm1.bias",
1328
+ "fusion.layers.23.norm1.weight",
1329
+ "fusion.layers.23.norm2.bias",
1330
+ "fusion.layers.23.norm2.weight",
1331
+ "fusion.layers.3.attn.in_proj_bias",
1332
+ "fusion.layers.3.attn.in_proj_weight",
1333
+ "fusion.layers.3.attn.out_proj.bias",
1334
+ "fusion.layers.3.attn.out_proj.weight",
1335
+ "fusion.layers.3.ffn.layers.0.bias",
1336
+ "fusion.layers.3.ffn.layers.0.weight",
1337
+ "fusion.layers.3.ffn.layers.3.bias",
1338
+ "fusion.layers.3.ffn.layers.3.weight",
1339
+ "fusion.layers.3.norm1.bias",
1340
+ "fusion.layers.3.norm1.weight",
1341
+ "fusion.layers.3.norm2.bias",
1342
+ "fusion.layers.3.norm2.weight",
1343
+ "fusion.layers.4.attn.in_proj_bias",
1344
+ "fusion.layers.4.attn.in_proj_weight",
1345
+ "fusion.layers.4.attn.out_proj.bias",
1346
+ "fusion.layers.4.attn.out_proj.weight",
1347
+ "fusion.layers.4.ffn.layers.0.bias",
1348
+ "fusion.layers.4.ffn.layers.0.weight",
1349
+ "fusion.layers.4.ffn.layers.3.bias",
1350
+ "fusion.layers.4.ffn.layers.3.weight",
1351
+ "fusion.layers.4.norm1.bias",
1352
+ "fusion.layers.4.norm1.weight",
1353
+ "fusion.layers.4.norm2.bias",
1354
+ "fusion.layers.4.norm2.weight",
1355
+ "fusion.layers.5.attn.in_proj_bias",
1356
+ "fusion.layers.5.attn.in_proj_weight",
1357
+ "fusion.layers.5.attn.out_proj.bias",
1358
+ "fusion.layers.5.attn.out_proj.weight",
1359
+ "fusion.layers.5.ffn.layers.0.bias",
1360
+ "fusion.layers.5.ffn.layers.0.weight",
1361
+ "fusion.layers.5.ffn.layers.3.bias",
1362
+ "fusion.layers.5.ffn.layers.3.weight",
1363
+ "fusion.layers.5.norm1.bias",
1364
+ "fusion.layers.5.norm1.weight",
1365
+ "fusion.layers.5.norm2.bias",
1366
+ "fusion.layers.5.norm2.weight",
1367
+ "fusion.layers.6.attn.in_proj_bias",
1368
+ "fusion.layers.6.attn.in_proj_weight",
1369
+ "fusion.layers.6.attn.out_proj.bias",
1370
+ "fusion.layers.6.attn.out_proj.weight",
1371
+ "fusion.layers.6.ffn.layers.0.bias",
1372
+ "fusion.layers.6.ffn.layers.0.weight",
1373
+ "fusion.layers.6.ffn.layers.3.bias",
1374
+ "fusion.layers.6.ffn.layers.3.weight",
1375
+ "fusion.layers.6.norm1.bias",
1376
+ "fusion.layers.6.norm1.weight",
1377
+ "fusion.layers.6.norm2.bias",
1378
+ "fusion.layers.6.norm2.weight",
1379
+ "fusion.layers.7.attn.in_proj_bias",
1380
+ "fusion.layers.7.attn.in_proj_weight",
1381
+ "fusion.layers.7.attn.out_proj.bias",
1382
+ "fusion.layers.7.attn.out_proj.weight",
1383
+ "fusion.layers.7.ffn.layers.0.bias",
1384
+ "fusion.layers.7.ffn.layers.0.weight",
1385
+ "fusion.layers.7.ffn.layers.3.bias",
1386
+ "fusion.layers.7.ffn.layers.3.weight",
1387
+ "fusion.layers.7.norm1.bias",
1388
+ "fusion.layers.7.norm1.weight",
1389
+ "fusion.layers.7.norm2.bias",
1390
+ "fusion.layers.7.norm2.weight",
1391
+ "fusion.layers.8.attn.in_proj_bias",
1392
+ "fusion.layers.8.attn.in_proj_weight",
1393
+ "fusion.layers.8.attn.out_proj.bias",
1394
+ "fusion.layers.8.attn.out_proj.weight",
1395
+ "fusion.layers.8.ffn.layers.0.bias",
1396
+ "fusion.layers.8.ffn.layers.0.weight",
1397
+ "fusion.layers.8.ffn.layers.3.bias",
1398
+ "fusion.layers.8.ffn.layers.3.weight",
1399
+ "fusion.layers.8.norm1.bias",
1400
+ "fusion.layers.8.norm1.weight",
1401
+ "fusion.layers.8.norm2.bias",
1402
+ "fusion.layers.8.norm2.weight",
1403
+ "fusion.layers.9.attn.in_proj_bias",
1404
+ "fusion.layers.9.attn.in_proj_weight",
1405
+ "fusion.layers.9.attn.out_proj.bias",
1406
+ "fusion.layers.9.attn.out_proj.weight",
1407
+ "fusion.layers.9.ffn.layers.0.bias",
1408
+ "fusion.layers.9.ffn.layers.0.weight",
1409
+ "fusion.layers.9.ffn.layers.3.bias",
1410
+ "fusion.layers.9.ffn.layers.3.weight",
1411
+ "fusion.layers.9.norm1.bias",
1412
+ "fusion.layers.9.norm1.weight",
1413
+ "fusion.layers.9.norm2.bias",
1414
+ "fusion.layers.9.norm2.weight",
1415
+ "fusion.porj_linear.bias",
1416
+ "fusion.porj_linear.weight",
1417
+ "head_rec_hr.bottleneck.bn.bias",
1418
+ "head_rec_hr.bottleneck.bn.num_batches_tracked",
1419
+ "head_rec_hr.bottleneck.bn.running_mean",
1420
+ "head_rec_hr.bottleneck.bn.running_var",
1421
+ "head_rec_hr.bottleneck.bn.weight",
1422
+ "head_rec_hr.bottleneck.conv.weight",
1423
+ "head_rec_hr.conv_seg.bias",
1424
+ "head_rec_hr.conv_seg.weight",
1425
+ "head_rec_hr.fpn_bottleneck.bn.bias",
1426
+ "head_rec_hr.fpn_bottleneck.bn.num_batches_tracked",
1427
+ "head_rec_hr.fpn_bottleneck.bn.running_mean",
1428
+ "head_rec_hr.fpn_bottleneck.bn.running_var",
1429
+ "head_rec_hr.fpn_bottleneck.bn.weight",
1430
+ "head_rec_hr.fpn_bottleneck.conv.weight",
1431
+ "head_rec_hr.fpn_convs.0.bn.bias",
1432
+ "head_rec_hr.fpn_convs.0.bn.num_batches_tracked",
1433
+ "head_rec_hr.fpn_convs.0.bn.running_mean",
1434
+ "head_rec_hr.fpn_convs.0.bn.running_var",
1435
+ "head_rec_hr.fpn_convs.0.bn.weight",
1436
+ "head_rec_hr.fpn_convs.0.conv.weight",
1437
+ "head_rec_hr.fpn_convs.1.bn.bias",
1438
+ "head_rec_hr.fpn_convs.1.bn.num_batches_tracked",
1439
+ "head_rec_hr.fpn_convs.1.bn.running_mean",
1440
+ "head_rec_hr.fpn_convs.1.bn.running_var",
1441
+ "head_rec_hr.fpn_convs.1.bn.weight",
1442
+ "head_rec_hr.fpn_convs.1.conv.weight",
1443
+ "head_rec_hr.fpn_convs.2.bn.bias",
1444
+ "head_rec_hr.fpn_convs.2.bn.num_batches_tracked",
1445
+ "head_rec_hr.fpn_convs.2.bn.running_mean",
1446
+ "head_rec_hr.fpn_convs.2.bn.running_var",
1447
+ "head_rec_hr.fpn_convs.2.bn.weight",
1448
+ "head_rec_hr.fpn_convs.2.conv.weight",
1449
+ "head_rec_hr.fpn_convs.3.bn.bias",
1450
+ "head_rec_hr.fpn_convs.3.bn.num_batches_tracked",
1451
+ "head_rec_hr.fpn_convs.3.bn.running_mean",
1452
+ "head_rec_hr.fpn_convs.3.bn.running_var",
1453
+ "head_rec_hr.fpn_convs.3.bn.weight",
1454
+ "head_rec_hr.fpn_convs.3.conv.weight",
1455
+ "head_rec_hr.lateral_convs.0.bn.bias",
1456
+ "head_rec_hr.lateral_convs.0.bn.num_batches_tracked",
1457
+ "head_rec_hr.lateral_convs.0.bn.running_mean",
1458
+ "head_rec_hr.lateral_convs.0.bn.running_var",
1459
+ "head_rec_hr.lateral_convs.0.bn.weight",
1460
+ "head_rec_hr.lateral_convs.0.conv.weight",
1461
+ "head_rec_hr.lateral_convs.1.bn.bias",
1462
+ "head_rec_hr.lateral_convs.1.bn.num_batches_tracked",
1463
+ "head_rec_hr.lateral_convs.1.bn.running_mean",
1464
+ "head_rec_hr.lateral_convs.1.bn.running_var",
1465
+ "head_rec_hr.lateral_convs.1.bn.weight",
1466
+ "head_rec_hr.lateral_convs.1.conv.weight",
1467
+ "head_rec_hr.lateral_convs.2.bn.bias",
1468
+ "head_rec_hr.lateral_convs.2.bn.num_batches_tracked",
1469
+ "head_rec_hr.lateral_convs.2.bn.running_mean",
1470
+ "head_rec_hr.lateral_convs.2.bn.running_var",
1471
+ "head_rec_hr.lateral_convs.2.bn.weight",
1472
+ "head_rec_hr.lateral_convs.2.conv.weight",
1473
+ "head_rec_hr.lateral_convs.3.bn.bias",
1474
+ "head_rec_hr.lateral_convs.3.bn.num_batches_tracked",
1475
+ "head_rec_hr.lateral_convs.3.bn.running_mean",
1476
+ "head_rec_hr.lateral_convs.3.bn.running_var",
1477
+ "head_rec_hr.lateral_convs.3.bn.weight",
1478
+ "head_rec_hr.lateral_convs.3.conv.weight",
1479
+ "head_rec_hr.psp_modules.0.1.bn.bias",
1480
+ "head_rec_hr.psp_modules.0.1.bn.num_batches_tracked",
1481
+ "head_rec_hr.psp_modules.0.1.bn.running_mean",
1482
+ "head_rec_hr.psp_modules.0.1.bn.running_var",
1483
+ "head_rec_hr.psp_modules.0.1.bn.weight",
1484
+ "head_rec_hr.psp_modules.0.1.conv.weight",
1485
+ "head_rec_hr.psp_modules.1.1.bn.bias",
1486
+ "head_rec_hr.psp_modules.1.1.bn.num_batches_tracked",
1487
+ "head_rec_hr.psp_modules.1.1.bn.running_mean",
1488
+ "head_rec_hr.psp_modules.1.1.bn.running_var",
1489
+ "head_rec_hr.psp_modules.1.1.bn.weight",
1490
+ "head_rec_hr.psp_modules.1.1.conv.weight",
1491
+ "head_rec_hr.psp_modules.2.1.bn.bias",
1492
+ "head_rec_hr.psp_modules.2.1.bn.num_batches_tracked",
1493
+ "head_rec_hr.psp_modules.2.1.bn.running_mean",
1494
+ "head_rec_hr.psp_modules.2.1.bn.running_var",
1495
+ "head_rec_hr.psp_modules.2.1.bn.weight",
1496
+ "head_rec_hr.psp_modules.2.1.conv.weight",
1497
+ "head_rec_hr.psp_modules.3.1.bn.bias",
1498
+ "head_rec_hr.psp_modules.3.1.bn.num_batches_tracked",
1499
+ "head_rec_hr.psp_modules.3.1.bn.running_mean",
1500
+ "head_rec_hr.psp_modules.3.1.bn.running_var",
1501
+ "head_rec_hr.psp_modules.3.1.bn.weight",
1502
+ "head_rec_hr.psp_modules.3.1.conv.weight",
1503
+ "head_s1.decoder.0.bias",
1504
+ "head_s1.decoder.0.weight",
1505
+ "head_s2.decoder.0.bias",
1506
+ "head_s2.decoder.0.weight",
1507
+ "modality_vae.vae_hr.codebook.weight",
1508
+ "modality_vae.vae_hr.conv1.bias",
1509
+ "modality_vae.vae_hr.conv1.weight",
1510
+ "modality_vae.vae_hr.conv2.bias",
1511
+ "modality_vae.vae_hr.conv2.weight",
1512
+ "modality_vae.vae_hr.dec_block1.0.bias",
1513
+ "modality_vae.vae_hr.dec_block1.0.weight",
1514
+ "modality_vae.vae_hr.dec_block1.2.bias",
1515
+ "modality_vae.vae_hr.dec_block1.2.weight",
1516
+ "modality_vae.vae_hr.dec_block2.0.bias",
1517
+ "modality_vae.vae_hr.dec_block2.0.weight",
1518
+ "modality_vae.vae_hr.dec_block2.2.bias",
1519
+ "modality_vae.vae_hr.dec_block2.2.weight",
1520
+ "modality_vae.vae_hr.enc_block1.0.bias",
1521
+ "modality_vae.vae_hr.enc_block1.0.weight",
1522
+ "modality_vae.vae_hr.enc_block1.2.bias",
1523
+ "modality_vae.vae_hr.enc_block1.2.weight",
1524
+ "modality_vae.vae_hr.enc_block2.0.bias",
1525
+ "modality_vae.vae_hr.enc_block2.0.weight",
1526
+ "modality_vae.vae_hr.enc_block2.2.bias",
1527
+ "modality_vae.vae_hr.enc_block2.2.weight",
1528
+ "modality_vae.vae_hr.gamma_1",
1529
+ "modality_vae.vae_hr.gamma_2",
1530
+ "modality_vae.vae_hr.gamma_3",
1531
+ "modality_vae.vae_hr.gamma_4",
1532
+ "modality_vae.vae_hr.logit_conv.bias",
1533
+ "modality_vae.vae_hr.logit_conv.weight",
1534
+ "modality_vae.vae_hr.rec_conv.bias",
1535
+ "modality_vae.vae_hr.rec_conv.weight",
1536
+ "modality_vae.vae_s1.codebook.weight",
1537
+ "modality_vae.vae_s1.conv1.bias",
1538
+ "modality_vae.vae_s1.conv1.weight",
1539
+ "modality_vae.vae_s1.conv2.bias",
1540
+ "modality_vae.vae_s1.conv2.weight",
1541
+ "modality_vae.vae_s1.dec_block1.0.bias",
1542
+ "modality_vae.vae_s1.dec_block1.0.weight",
1543
+ "modality_vae.vae_s1.dec_block1.2.bias",
1544
+ "modality_vae.vae_s1.dec_block1.2.weight",
1545
+ "modality_vae.vae_s1.dec_block2.0.bias",
1546
+ "modality_vae.vae_s1.dec_block2.0.weight",
1547
+ "modality_vae.vae_s1.dec_block2.2.bias",
1548
+ "modality_vae.vae_s1.dec_block2.2.weight",
1549
+ "modality_vae.vae_s1.enc_block1.0.bias",
1550
+ "modality_vae.vae_s1.enc_block1.0.weight",
1551
+ "modality_vae.vae_s1.enc_block1.2.bias",
1552
+ "modality_vae.vae_s1.enc_block1.2.weight",
1553
+ "modality_vae.vae_s1.enc_block2.0.bias",
1554
+ "modality_vae.vae_s1.enc_block2.0.weight",
1555
+ "modality_vae.vae_s1.enc_block2.2.bias",
1556
+ "modality_vae.vae_s1.enc_block2.2.weight",
1557
+ "modality_vae.vae_s1.gamma_1",
1558
+ "modality_vae.vae_s1.gamma_2",
1559
+ "modality_vae.vae_s1.gamma_3",
1560
+ "modality_vae.vae_s1.gamma_4",
1561
+ "modality_vae.vae_s1.logit_conv.bias",
1562
+ "modality_vae.vae_s1.logit_conv.weight",
1563
+ "modality_vae.vae_s1.rec_conv.bias",
1564
+ "modality_vae.vae_s1.rec_conv.weight",
1565
+ "modality_vae.vae_s2.codebook.weight",
1566
+ "modality_vae.vae_s2.conv1.bias",
1567
+ "modality_vae.vae_s2.conv1.weight",
1568
+ "modality_vae.vae_s2.conv2.bias",
1569
+ "modality_vae.vae_s2.conv2.weight",
1570
+ "modality_vae.vae_s2.dec_block1.0.bias",
1571
+ "modality_vae.vae_s2.dec_block1.0.weight",
1572
+ "modality_vae.vae_s2.dec_block1.2.bias",
1573
+ "modality_vae.vae_s2.dec_block1.2.weight",
1574
+ "modality_vae.vae_s2.dec_block2.0.bias",
1575
+ "modality_vae.vae_s2.dec_block2.0.weight",
1576
+ "modality_vae.vae_s2.dec_block2.2.bias",
1577
+ "modality_vae.vae_s2.dec_block2.2.weight",
1578
+ "modality_vae.vae_s2.enc_block1.0.bias",
1579
+ "modality_vae.vae_s2.enc_block1.0.weight",
1580
+ "modality_vae.vae_s2.enc_block1.2.bias",
1581
+ "modality_vae.vae_s2.enc_block1.2.weight",
1582
+ "modality_vae.vae_s2.enc_block2.0.bias",
1583
+ "modality_vae.vae_s2.enc_block2.0.weight",
1584
+ "modality_vae.vae_s2.enc_block2.2.bias",
1585
+ "modality_vae.vae_s2.enc_block2.2.weight",
1586
+ "modality_vae.vae_s2.gamma_1",
1587
+ "modality_vae.vae_s2.gamma_2",
1588
+ "modality_vae.vae_s2.gamma_3",
1589
+ "modality_vae.vae_s2.gamma_4",
1590
+ "modality_vae.vae_s2.logit_conv.bias",
1591
+ "modality_vae.vae_s2.logit_conv.weight",
1592
+ "modality_vae.vae_s2.rec_conv.bias",
1593
+ "modality_vae.vae_s2.rec_conv.weight"
1594
+ ]
1595
+ }
skysensepp-fewshot-release/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e6d35f7768cce0bb019d3f7089265ff6099168d651ac3ae5b4abc0139dfb759c
3
+ size 7238917044
skysensepp-fewshot-release/modeling_skysensepp.py ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Full SkySense++ model for few-shot / 1-shot segmentation."""
2
+
3
+ from dataclasses import dataclass
4
+ from typing import Dict, List, Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ from transformers import PreTrainedModel
10
+ from transformers.modeling_outputs import ModelOutput
11
+
12
+ from .configuration_skysensepp import SkySensePlusPlusConfig
13
+ from .modeling_utils import DropPath as _DropPath # noqa: F401 — bundled for remote code
14
+ from .modeling_skysensepp_components import ModalityCompletion, UPerHead, UPHead
15
+ from .modeling_skysensepp_fusion_neck import SkySensePlusPlusFusionNeckModel
16
+ from .modeling_skysensepp_swinv2_msl import SkySensePlusPlusSwinV2MSLModel
17
+ from .modeling_skysensepp_vit_msl import SkySensePlusPlusViTMSLModel
18
+
19
+ IMAGENET_MEAN = (0.485, 0.456, 0.406)
20
+ IMAGENET_STD = (0.229, 0.224, 0.225)
21
+
22
+
23
+ @dataclass
24
+ class SkySensePlusPlusOutput(ModelOutput):
25
+ logits: Optional[torch.FloatTensor] = None
26
+ mapped_targets: Optional[torch.LongTensor] = None
27
+ idx_2_color: Optional[dict] = None
28
+ mask_hr: Optional[torch.Tensor] = None
29
+ vae_out: Optional[dict] = None
30
+
31
+
32
+ class SkySensePlusPlusPreTrainedModel(PreTrainedModel):
33
+ config_class = SkySensePlusPlusConfig
34
+ base_model_prefix = "skysensepp"
35
+ supports_gradient_checkpointing = False
36
+
37
+
38
+ class SkySensePlusPlusModel(SkySensePlusPlusPreTrainedModel):
39
+ """End-to-end SkySense++ pipeline matching the released few-shot checkpoint."""
40
+
41
+ def __init__(self, config: SkySensePlusPlusConfig):
42
+ super().__init__(config)
43
+ self.sources = list(config.sources)
44
+ self.use_modal_vae = config.use_modal_vae
45
+ self.vocabulary_size = config.vocabulary_size
46
+ self.vocabulary = list(range(1, config.vocabulary_size + 1))
47
+
48
+ if "hr" in self.sources:
49
+ self.backbone_hr = SkySensePlusPlusSwinV2MSLModel(config.backbone_hr)
50
+ if "s2" in self.sources:
51
+ self.backbone_s2 = SkySensePlusPlusViTMSLModel(config.backbone_s2)
52
+ self.head_s2 = UPHead(config.head_s2.in_dim, config.head_s2.out_dim, config.head_s2.up_scale)
53
+ if "s1" in self.sources:
54
+ self.backbone_s1 = SkySensePlusPlusViTMSLModel(config.backbone_s1)
55
+ self.head_s1 = UPHead(config.head_s1.in_dim, config.head_s1.out_dim, config.head_s1.up_scale)
56
+
57
+ self.fusion = SkySensePlusPlusFusionNeckModel(config.fusion)
58
+ if self.use_modal_vae:
59
+ self.modality_vae = ModalityCompletion(
60
+ input_shape_hr=tuple(config.modality_vae.input_shape_hr),
61
+ input_shape_s2=tuple(config.modality_vae.input_shape_s2),
62
+ input_shape_s1=tuple(config.modality_vae.input_shape_s1),
63
+ conv_dim=config.modality_vae.conv_dim,
64
+ z_dim=config.modality_vae.z_dim,
65
+ n_codebook=config.modality_vae.n_codebook,
66
+ )
67
+ self.head_rec_hr = UPerHead(
68
+ in_channels=config.head_rec_hr.in_channels,
69
+ channels=config.head_rec_hr.channels,
70
+ num_classes=config.head_rec_hr.num_classes,
71
+ pool_scales=tuple(config.head_rec_hr.pool_scales),
72
+ dropout_ratio=config.head_rec_hr.dropout_ratio,
73
+ align_corners=config.head_rec_hr.align_corners,
74
+ )
75
+ self.post_init()
76
+
77
+ def convert_target(self, target: torch.Tensor):
78
+ mean = target.new_tensor(IMAGENET_MEAN).reshape(1, 3, 1, 1)
79
+ std = target.new_tensor(IMAGENET_STD).reshape(1, 3, 1, 1)
80
+ target = ((target * std + mean) * 255).to(torch.long)
81
+ target = target[:, 0] * 256 * 256 + target[:, 1] * 256 + target[:, 2]
82
+ target = target.type(torch.long)
83
+ unique_target = target.unique()
84
+ target_index = torch.searchsorted(unique_target, target)
85
+ no_bg = unique_target[0].item() > 0
86
+ if no_bg:
87
+ target_index = target_index + 1
88
+ target_index_unique = target_index.unique().tolist()
89
+ vocab = target.new_tensor([0] + self.vocabulary)
90
+ mapped_target = target_index.clone()
91
+ idx_2_color = {}
92
+ for value in target_index_unique:
93
+ mapped_target[target_index == value] = vocab[value]
94
+ idx_2_color[vocab[value].item()] = unique_target[value - 1 if no_bg else value].item()
95
+ return mapped_target, idx_2_color
96
+
97
+ def forward(
98
+ self,
99
+ hr_img: Optional[torch.Tensor] = None,
100
+ s2_img: Optional[torch.Tensor] = None,
101
+ s1_img: Optional[torch.Tensor] = None,
102
+ targets: Optional[torch.Tensor] = None,
103
+ anno_mask: Optional[torch.Tensor] = None,
104
+ modality_flags: Optional[torch.Tensor] = None,
105
+ return_dict: Optional[bool] = None,
106
+ ) -> Union[Dict, SkySensePlusPlusOutput]:
107
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
108
+ output: Dict = {}
109
+
110
+ if targets is None:
111
+ raise ValueError("SkySense++ few-shot forward requires `targets` for annotation conditioning.")
112
+ anno_img, idx_2_color = self.convert_target(targets)
113
+ output["mapped_targets"] = anno_img
114
+ output["idx_2_color"] = idx_2_color
115
+
116
+ anno_s2 = anno_img[:, 15::32, 15::32]
117
+ anno_s1 = anno_s2
118
+
119
+ if anno_mask is not None:
120
+ batch_size, mask_h, mask_w = anno_mask.shape
121
+ block_size = 32
122
+ anno_mask_hr = (
123
+ anno_mask.unsqueeze(-1)
124
+ .unsqueeze(-1)
125
+ .repeat(1, 1, 1, block_size, block_size)
126
+ .permute(0, 1, 3, 2, 4)
127
+ .reshape(batch_size, mask_h * block_size, mask_w * block_size)
128
+ .contiguous()
129
+ )
130
+ else:
131
+ anno_mask_hr = None
132
+
133
+ if "hr" in self.sources:
134
+ hr_features = self.backbone_hr(hr_img, anno_img, anno_mask_hr, return_dict=False)
135
+ output["mask_hr"] = anno_mask_hr
136
+
137
+ batch_size = hr_img.shape[0]
138
+ seq_len_s2 = s2_img.shape[2] if s2_img is not None else 1
139
+ seq_len_s1 = s1_img.shape[2] if s1_img is not None else 1
140
+
141
+ if "s2" in self.sources:
142
+ b, c, seq, h, w = s2_img.shape
143
+ s2_flat = s2_img.permute(0, 2, 1, 3, 4).reshape(b * seq, c, h, w).contiguous()
144
+ s2_features = self.backbone_s2(s2_flat, anno_s2, anno_mask, return_dict=False)
145
+ s2_features = self.head_s2(s2_features[-1])
146
+ s2_features = [s2_features]
147
+
148
+ if "s1" in self.sources:
149
+ b, c, seq, h, w = s1_img.shape
150
+ s1_flat = s1_img.permute(0, 2, 1, 3, 4).reshape(b * seq, c, h, w).contiguous()
151
+ s1_features = self.backbone_s1(s1_flat, anno_s1, anno_mask, return_dict=False)
152
+ s1_features = self.head_s1(s1_features[-1])
153
+ s1_features = [s1_features]
154
+
155
+ hr_features_stage3 = hr_features[-1]
156
+ s2_features_stage3 = s2_features[-1]
157
+ s1_features_stage3 = s1_features[-1]
158
+
159
+ if modality_flags is None:
160
+ modality_flags = torch.tensor([[0, 0, 1]] * batch_size, device=hr_img.device, dtype=torch.float32)
161
+
162
+ if self.use_modal_vae:
163
+ vae_out = self.modality_vae(hr_features_stage3, s2_features_stage3, s1_features_stage3, modality_flags)
164
+ hr_features_stage3 = vae_out["hr_out"]
165
+ s2_features_stage3 = vae_out["s2_out"]
166
+ s1_features_stage3 = vae_out["s1_out"]
167
+ output["vae_out"] = vae_out
168
+
169
+ _, c3, h3, w3 = hr_features_stage3.shape
170
+ hr_tokens = hr_features_stage3.permute(0, 2, 3, 1).reshape(batch_size * h3 * w3, c3).unsqueeze(1)
171
+
172
+ _, c3_s2, h3_s2, w3_s2 = s2_features_stage3.shape
173
+ s2_tokens = (
174
+ s2_features_stage3.reshape(batch_size, seq_len_s2, c3_s2, h3_s2, w3_s2)
175
+ .permute(0, 3, 4, 1, 2)
176
+ .reshape(batch_size, h3_s2 * w3_s2, seq_len_s2, c3_s2)
177
+ .reshape(batch_size * h3_s2 * w3_s2, seq_len_s2, c3_s2)
178
+ .contiguous()
179
+ )
180
+ features_stage3 = torch.cat((hr_tokens, s2_tokens), dim=1)
181
+
182
+ _, c3_s1, h3_s1, w3_s1 = s1_features_stage3.shape
183
+ s1_tokens = (
184
+ s1_features_stage3.reshape(batch_size, seq_len_s1, c3_s1, h3_s1, w3_s1)
185
+ .permute(0, 3, 4, 1, 2)
186
+ .reshape(batch_size, h3_s1 * w3_s1, seq_len_s1, c3_s1)
187
+ .reshape(batch_size * h3_s1 * w3_s1, seq_len_s1, c3_s1)
188
+ .contiguous()
189
+ )
190
+ features_stage3 = torch.cat((features_stage3, s1_tokens), dim=1)
191
+
192
+ fusion_out = self.fusion(features_stage3, return_dict=True)
193
+ cls_token = fusion_out.pooler_output.reshape(batch_size, h3, w3, -1).permute(0, 3, 1, 2).contiguous()
194
+
195
+ hr_rec_inputs = list(hr_features)
196
+ feat_stage1 = hr_rec_inputs[0]
197
+ if feat_stage1.shape[-1] == feat_stage1.shape[-2]:
198
+ left, right = torch.split(feat_stage1, feat_stage1.shape[-1] // 2, dim=-1)
199
+ hr_rec_inputs[0] = torch.cat((left, right), dim=1)
200
+
201
+ logits_hr = self.head_rec_hr([*hr_rec_inputs, cls_token])
202
+ if self.config.upsample_results:
203
+ logits_hr = F.interpolate(logits_hr.float(), scale_factor=4, mode="bilinear", align_corners=True)
204
+ output["logits_hr"] = logits_hr
205
+
206
+ if not return_dict:
207
+ return output
208
+ return SkySensePlusPlusOutput(
209
+ logits=logits_hr,
210
+ mapped_targets=output.get("mapped_targets"),
211
+ idx_2_color=output.get("idx_2_color"),
212
+ mask_hr=output.get("mask_hr"),
213
+ vae_out=output.get("vae_out"),
214
+ )
skysensepp-fewshot-release/modeling_skysensepp_components.py ADDED
@@ -0,0 +1,238 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Shared heads and necks for the full SkySense++ model."""
2
+
3
+ from typing import List, Sequence, Tuple
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.nn.functional as F
8
+
9
+
10
+ def resize_tensor(x: torch.Tensor, size: Tuple[int, int], align_corners: bool = False) -> torch.Tensor:
11
+ return F.interpolate(x, size=size, mode="bilinear", align_corners=align_corners)
12
+
13
+
14
+ class ConvModule(nn.Module):
15
+ def __init__(self, in_channels: int, out_channels: int, kernel_size: int, padding: int = 0):
16
+ super().__init__()
17
+ self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
18
+ self.bn = nn.BatchNorm2d(out_channels)
19
+ self.relu = nn.ReLU(inplace=False)
20
+
21
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
22
+ return self.relu(self.bn(self.conv(x)))
23
+
24
+
25
+ class UPHead(nn.Module):
26
+ def __init__(self, in_dim: int, out_dim: int, up_scale: int):
27
+ super().__init__()
28
+ self.decoder = nn.Sequential(
29
+ nn.Conv2d(in_dim, up_scale**2 * out_dim, kernel_size=1),
30
+ nn.PixelShuffle(up_scale),
31
+ )
32
+
33
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
34
+ return self.decoder(x)
35
+
36
+
37
+ class PPM(nn.ModuleList):
38
+ def __init__(
39
+ self,
40
+ pool_scales: Sequence[int],
41
+ in_channels: int,
42
+ channels: int,
43
+ align_corners: bool = False,
44
+ ):
45
+ super().__init__()
46
+ self.align_corners = align_corners
47
+ for pool_scale in pool_scales:
48
+ self.append(
49
+ nn.Sequential(
50
+ nn.AdaptiveAvgPool2d(pool_scale),
51
+ ConvModule(in_channels, channels, kernel_size=1),
52
+ )
53
+ )
54
+
55
+ def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
56
+ outputs = []
57
+ for module in self:
58
+ out = module(x)
59
+ out = resize_tensor(out, x.shape[2:], align_corners=self.align_corners)
60
+ outputs.append(out)
61
+ return outputs
62
+
63
+
64
+ class UPerHead(nn.Module):
65
+ def __init__(
66
+ self,
67
+ in_channels: Sequence[int] = (704, 704, 1408, 2816, 1024),
68
+ channels: int = 512,
69
+ num_classes: int = 65,
70
+ pool_scales: Sequence[int] = (1, 2, 3, 6),
71
+ dropout_ratio: float = 0.1,
72
+ align_corners: bool = False,
73
+ ):
74
+ super().__init__()
75
+ self.in_channels = list(in_channels)
76
+ self.channels = channels
77
+ self.align_corners = align_corners
78
+ self.psp_modules = PPM(pool_scales, self.in_channels[-1], channels, align_corners)
79
+ self.bottleneck = ConvModule(
80
+ self.in_channels[-1] + len(pool_scales) * channels,
81
+ channels,
82
+ kernel_size=3,
83
+ padding=1,
84
+ )
85
+ self.lateral_convs = nn.ModuleList()
86
+ self.fpn_convs = nn.ModuleList()
87
+ for in_ch in self.in_channels[:-1]:
88
+ self.lateral_convs.append(ConvModule(in_ch, channels, kernel_size=1))
89
+ self.fpn_convs.append(ConvModule(channels, channels, kernel_size=3, padding=1))
90
+ self.fpn_bottleneck = ConvModule(len(self.in_channels) * channels, channels, kernel_size=3, padding=1)
91
+ self.dropout = nn.Dropout2d(dropout_ratio) if dropout_ratio > 0 else nn.Identity()
92
+ self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1)
93
+
94
+ def psp_forward(self, inputs: List[torch.Tensor]) -> torch.Tensor:
95
+ x = inputs[-1]
96
+ psp_outs = [x, *self.psp_modules(x)]
97
+ return self.bottleneck(torch.cat(psp_outs, dim=1))
98
+
99
+ def forward(self, inputs: List[torch.Tensor]) -> torch.Tensor:
100
+ laterals = [conv(inputs[i]) for i, conv in enumerate(self.lateral_convs)]
101
+ laterals.append(self.psp_forward(inputs))
102
+
103
+ for i in range(len(laterals) - 1, 0, -1):
104
+ laterals[i - 1] = laterals[i - 1] + resize_tensor(
105
+ laterals[i], laterals[i - 1].shape[2:], align_corners=self.align_corners
106
+ )
107
+
108
+ fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(len(laterals) - 1)]
109
+ fpn_outs.append(laterals[-1])
110
+ for i in range(len(fpn_outs) - 1, 0, -1):
111
+ fpn_outs[i] = resize_tensor(fpn_outs[i], fpn_outs[0].shape[2:], align_corners=self.align_corners)
112
+ output = self.fpn_bottleneck(torch.cat(fpn_outs, dim=1))
113
+ output = self.dropout(output)
114
+ return self.conv_seg(output)
115
+
116
+
117
+ class BFloat16UpsampleNearest2d(nn.Module):
118
+ def __init__(self, scale_factor: int, mode: str = "bilinear"):
119
+ super().__init__()
120
+ self.scale_factor = scale_factor
121
+ self.mode = mode
122
+
123
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
124
+ upsampled = F.interpolate(x.float(), scale_factor=self.scale_factor, mode=self.mode)
125
+ return upsampled.to(x.dtype)
126
+
127
+
128
+ class ConvVQVAEv2(nn.Module):
129
+ def __init__(self, input_shape: Tuple[int, int, int], conv_dim: int, z_dim: int, num_tokens: int = 8192, temp: float = 0.9):
130
+ super().__init__()
131
+ self.temp = temp
132
+ self.codebook = nn.Embedding(num_tokens, z_dim)
133
+ self.relu = nn.LeakyReLU()
134
+ self.pool = nn.AvgPool2d(2)
135
+ self.conv1 = nn.Conv2d(input_shape[0], conv_dim, 5, stride=1, padding=2)
136
+ self.enc_block1 = nn.Sequential(
137
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
138
+ nn.LeakyReLU(),
139
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
140
+ nn.LeakyReLU(),
141
+ )
142
+ self.gamma_1 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1)))
143
+ self.enc_block2 = nn.Sequential(
144
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
145
+ nn.LeakyReLU(),
146
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
147
+ nn.LeakyReLU(),
148
+ )
149
+ self.gamma_2 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1)))
150
+ self.logit_conv = nn.Conv2d(conv_dim, num_tokens, 1)
151
+ self.unpool = BFloat16UpsampleNearest2d(scale_factor=2)
152
+ self.conv2 = nn.Conv2d(z_dim, conv_dim, 3, stride=1, padding=1)
153
+ self.dec_block1 = nn.Sequential(
154
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
155
+ nn.LeakyReLU(),
156
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
157
+ nn.LeakyReLU(),
158
+ )
159
+ self.gamma_3 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1)))
160
+ self.dec_block2 = nn.Sequential(
161
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
162
+ nn.LeakyReLU(),
163
+ nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1),
164
+ nn.LeakyReLU(),
165
+ )
166
+ self.gamma_4 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1)))
167
+ self.rec_conv = nn.Conv2d(conv_dim, input_shape[0], 3, stride=1, padding=1)
168
+
169
+ def forward_encoder(self, x: torch.Tensor) -> torch.Tensor:
170
+ x = self.relu(self.conv1(x))
171
+ x = x + self.gamma_1 * self.enc_block1(x)
172
+ x = self.pool(x)
173
+ x = x + self.gamma_2 * self.enc_block2(x)
174
+ x = self.pool(x)
175
+ return self.logit_conv(x)
176
+
177
+ def forward_decoder(self, logits: torch.Tensor):
178
+ soft_one_hot = F.softmax(logits * (self.temp * 10), dim=1)
179
+ sampled = torch.einsum("bnhw,nd->bdhw", soft_one_hot, self.codebook.weight)
180
+ x = self.relu(self.conv2(sampled))
181
+ x = self.unpool(x)
182
+ x = x + self.gamma_3 * self.dec_block1(x)
183
+ x = self.unpool(x)
184
+ x = x + self.gamma_4 * self.dec_block2(x)
185
+ return self.rec_conv(x), soft_one_hot
186
+
187
+
188
+ class ModalityCompletion(nn.Module):
189
+ def __init__(
190
+ self,
191
+ input_shape_hr: Tuple[int, int, int] = (2816, 32, 16),
192
+ input_shape_s2: Tuple[int, int, int] = (2816, 32, 16),
193
+ input_shape_s1: Tuple[int, int, int] = (2816, 32, 16),
194
+ conv_dim: int = 256,
195
+ z_dim: int = 256,
196
+ n_codebook: int = 8192,
197
+ ):
198
+ super().__init__()
199
+ self.vae_hr = ConvVQVAEv2(input_shape_hr, conv_dim, z_dim, num_tokens=n_codebook)
200
+ self.vae_s2 = ConvVQVAEv2(input_shape_s2, conv_dim, z_dim, num_tokens=n_codebook)
201
+ self.vae_s1 = ConvVQVAEv2(input_shape_s1, conv_dim, z_dim, num_tokens=n_codebook)
202
+
203
+ def forward(
204
+ self,
205
+ feat_hr: torch.Tensor,
206
+ feat_s2: torch.Tensor,
207
+ feat_s1: torch.Tensor,
208
+ modality_info: torch.Tensor,
209
+ ) -> dict[str, torch.Tensor]:
210
+ logits_hr = self.vae_hr.forward_encoder(feat_hr)
211
+ logits_s2 = self.vae_s2.forward_encoder(feat_s2)
212
+ logits_s1 = self.vae_s1.forward_encoder(feat_s1)
213
+
214
+ flag_hr = modality_info[:, 0][:, None, None, None]
215
+ flag_s2 = modality_info[:, 1][:, None, None, None]
216
+ flag_s1 = modality_info[:, 2][:, None, None, None]
217
+
218
+ mean_logits_hr_s2 = logits_hr * flag_hr + logits_s2 * flag_s2
219
+ mean_logits_hr_s1 = logits_hr * flag_hr + logits_s1 * flag_s1
220
+ mean_logits_s1_s2 = logits_s1 * flag_s1 + logits_s2 * flag_s2
221
+
222
+ logits_hr_rec = logits_hr * flag_hr + mean_logits_s1_s2 * (1.0 - flag_hr)
223
+ logits_s2_rec = logits_s2 * flag_s2 + mean_logits_hr_s1 * (1.0 - flag_s2)
224
+ logits_s1_rec = logits_s1 * flag_s1 + mean_logits_hr_s2 * (1.0 - flag_s1)
225
+
226
+ g_hr, _ = self.vae_hr.forward_decoder(logits_hr_rec)
227
+ g_s2, _ = self.vae_s2.forward_decoder(logits_s2_rec)
228
+ g_s1, _ = self.vae_s1.forward_decoder(logits_s1_rec)
229
+
230
+ inv_hr = 1.0 - flag_hr
231
+ inv_s2 = 1.0 - flag_s2
232
+ inv_s1 = 1.0 - flag_s1
233
+
234
+ return {
235
+ "hr_out": feat_hr * flag_hr + g_hr * inv_hr,
236
+ "s2_out": feat_s2 * flag_s2 + g_s2 * inv_s2,
237
+ "s1_out": feat_s1 * flag_s1 + g_s1 * inv_s1,
238
+ }
skysensepp-fewshot-release/modeling_skysensepp_fusion_neck.py ADDED
@@ -0,0 +1,164 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ fusion neck (TransformerEncoder) — optional multi-modal fusion module."""
2
+
3
+ from typing import Optional, Tuple, Union
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.utils.checkpoint as cp
8
+ from transformers import PreTrainedModel
9
+ from transformers.modeling_outputs import BaseModelOutputWithPooling
10
+
11
+ from .configuration_skysensepp import SkySensePlusPlusFusionNeckConfig
12
+ from .modeling_utils import DropPath, FFN
13
+
14
+
15
+ class FusionEncoderLayer(nn.Module):
16
+ def __init__(
17
+ self,
18
+ embed_dims: int,
19
+ num_heads: int,
20
+ feedforward_channels: int,
21
+ drop_rate: float = 0.0,
22
+ attn_drop_rate: float = 0.0,
23
+ drop_path_rate: float = 0.0,
24
+ qkv_bias: bool = True,
25
+ with_cp: bool = False,
26
+ ):
27
+ super().__init__()
28
+ self.with_cp = with_cp
29
+ self.norm1 = nn.LayerNorm(embed_dims)
30
+ self.attn = nn.MultiheadAttention(
31
+ embed_dim=embed_dims,
32
+ num_heads=num_heads,
33
+ dropout=attn_drop_rate,
34
+ bias=qkv_bias,
35
+ batch_first=True,
36
+ )
37
+ self.proj_drop = nn.Dropout(drop_rate)
38
+ self.norm2 = nn.LayerNorm(embed_dims)
39
+ self.ffn = FFN(
40
+ embed_dims=embed_dims,
41
+ feedforward_channels=feedforward_channels,
42
+ num_fcs=2,
43
+ ffn_drop=drop_rate,
44
+ drop_path=drop_path_rate,
45
+ act_layer=nn.GELU,
46
+ add_identity=True,
47
+ )
48
+ self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
49
+
50
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
51
+ def _inner_forward(x):
52
+ residual = x
53
+ x_norm = self.norm1(x)
54
+ attn_out, _ = self.attn(x_norm, x_norm, x_norm)
55
+ attn_out = self.proj_drop(attn_out)
56
+ x = residual + self.drop_path(attn_out)
57
+ return self.ffn(self.norm2(x), identity=x)
58
+
59
+ if self.with_cp and x.requires_grad:
60
+ return cp.checkpoint(_inner_forward, x, use_reentrant=False)
61
+ return _inner_forward(x)
62
+
63
+
64
+ class SkySensePlusPlusFusionNeckPreTrainedModel(PreTrainedModel):
65
+ config_class = SkySensePlusPlusFusionNeckConfig
66
+ base_model_prefix = "skysensepp_fusion_neck"
67
+ supports_gradient_checkpointing = True
68
+
69
+ def _init_weights(self, module):
70
+ if isinstance(module, nn.Linear):
71
+ nn.init.trunc_normal_(module.weight, std=0.02)
72
+ if module.bias is not None:
73
+ nn.init.zeros_(module.bias)
74
+ elif isinstance(module, nn.LayerNorm):
75
+ nn.init.ones_(module.weight)
76
+ nn.init.zeros_(module.bias)
77
+
78
+
79
+ class SkySensePlusPlusFusionNeckModel(SkySensePlusPlusFusionNeckPreTrainedModel):
80
+ """Fuses per-location multi-modal tokens into a cls-token representation.
81
+
82
+ Input shape: ``(batch, num_modalities, input_dims)`` — e.g. concatenated
83
+ HR + S2 + S1 stage-3 features with ``input_dims=2816``.
84
+
85
+ Output: cls token embedding ``(batch, embed_dims)`` when
86
+ ``output_cls_token=True`` (default).
87
+ """
88
+
89
+ def __init__(self, config: SkySensePlusPlusFusionNeckConfig):
90
+ super().__init__(config)
91
+
92
+ # Original checkpoint uses the typo `porj_linear`.
93
+ self.porj_linear = nn.Linear(config.input_dims, config.embed_dims)
94
+ self.with_cls_token = config.with_cls_token
95
+ self.output_cls_token = config.output_cls_token
96
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
97
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
98
+
99
+ num_layers = config.num_layers
100
+ if num_layers > 1:
101
+ dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
102
+ else:
103
+ dpr = [0.0]
104
+
105
+ self.layers = nn.ModuleList()
106
+ for i in range(config.num_layers):
107
+ self.layers.append(
108
+ FusionEncoderLayer(
109
+ embed_dims=config.embed_dims,
110
+ num_heads=config.num_heads,
111
+ feedforward_channels=config.mlp_ratio * config.embed_dims,
112
+ attn_drop_rate=config.attn_drop_rate,
113
+ drop_rate=config.drop_rate,
114
+ drop_path_rate=dpr[i],
115
+ qkv_bias=config.qkv_bias,
116
+ with_cp=config.with_cp,
117
+ )
118
+ )
119
+
120
+ self.post_init()
121
+
122
+ def forward(
123
+ self,
124
+ hidden_states: torch.Tensor,
125
+ output_hidden_states: Optional[bool] = None,
126
+ return_dict: Optional[bool] = None,
127
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
128
+ """Forward pass.
129
+
130
+ Args:
131
+ hidden_states: ``(batch, seq_len, input_dims)`` fused modality tokens.
132
+ """
133
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
134
+
135
+ x = self.porj_linear(hidden_states)
136
+ cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
137
+ x = torch.cat((cls_tokens, x), dim=1)
138
+ if not self.with_cls_token:
139
+ x = x[:, 1:]
140
+
141
+ all_hidden_states = () if output_hidden_states else None
142
+ for layer in self.layers:
143
+ x = layer(x)
144
+ if output_hidden_states:
145
+ all_hidden_states = all_hidden_states + (x,)
146
+
147
+ if self.output_cls_token:
148
+ pooler = x[:, 0]
149
+ last_hidden = pooler.unsqueeze(1)
150
+ elif self.with_cls_token:
151
+ pooler = None
152
+ last_hidden = x[:, 1:]
153
+ else:
154
+ pooler = None
155
+ last_hidden = x
156
+
157
+ if not return_dict:
158
+ return (last_hidden, pooler) if pooler is not None else (last_hidden,)
159
+
160
+ return BaseModelOutputWithPooling(
161
+ last_hidden_state=last_hidden,
162
+ pooler_output=pooler,
163
+ hidden_states=all_hidden_states,
164
+ )
skysensepp-fewshot-release/modeling_skysensepp_swinv2_msl.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ from copy import deepcopy
4
+ from typing import Optional, Sequence, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig
14
+ from .modeling_utils import (
15
+ DropPath,
16
+ FFN,
17
+ PatchEmbed,
18
+ PatchMerging,
19
+ ShiftWindowMSA,
20
+ to_2tuple,
21
+ )
22
+
23
+
24
+ class SwinBlockV2(nn.Module):
25
+ def __init__(
26
+ self,
27
+ embed_dims: int,
28
+ num_heads: int,
29
+ window_size: int = 8,
30
+ shift: bool = False,
31
+ extra_norm: bool = False,
32
+ ffn_ratio: float = 4.0,
33
+ drop_path: float = 0.0,
34
+ pad_small_map: bool = False,
35
+ with_cp: bool = False,
36
+ pretrained_window_size: int = 0,
37
+ ):
38
+ super().__init__()
39
+ self.with_cp = with_cp
40
+ self.extra_norm = extra_norm
41
+ self.attn = ShiftWindowMSA(
42
+ embed_dims=embed_dims,
43
+ num_heads=num_heads,
44
+ window_size=window_size,
45
+ shift_size=window_size // 2 if shift else 0,
46
+ drop_path=drop_path,
47
+ pad_small_map=pad_small_map,
48
+ pretrained_window_size=pretrained_window_size,
49
+ )
50
+ self.norm1 = nn.LayerNorm(embed_dims)
51
+ self.ffn = FFN(
52
+ embed_dims=embed_dims,
53
+ feedforward_channels=int(embed_dims * ffn_ratio),
54
+ num_fcs=2,
55
+ drop_path=drop_path,
56
+ act_layer=nn.GELU,
57
+ add_identity=False,
58
+ )
59
+ self.norm2 = nn.LayerNorm(embed_dims)
60
+ if self.extra_norm:
61
+ self.norm3 = nn.LayerNorm(embed_dims)
62
+
63
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
64
+ def _inner_forward(x):
65
+ identity = x
66
+ x = self.attn(x, hw_shape)
67
+ x = self.norm1(x)
68
+ x = x + identity
69
+
70
+ identity = x
71
+ x = self.ffn(x)
72
+ x = self.norm2(x)
73
+ x = x + identity
74
+
75
+ if self.extra_norm:
76
+ x = self.norm3(x)
77
+ return x
78
+
79
+ if self.with_cp and x.requires_grad:
80
+ x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
81
+ else:
82
+ x = _inner_forward(x)
83
+ return x
84
+
85
+
86
+ class SwinBlockV2Sequence(nn.Module):
87
+ def __init__(
88
+ self,
89
+ embed_dims: int,
90
+ depth: int,
91
+ num_heads: int,
92
+ window_size: int = 8,
93
+ downsample: bool = False,
94
+ drop_paths: Union[Sequence[float], float] = 0.0,
95
+ with_cp: bool = False,
96
+ pad_small_map: bool = False,
97
+ extra_norm_every_n_blocks: int = 0,
98
+ pretrained_window_size: int = 0,
99
+ is_post_norm_downsample: bool = True,
100
+ ):
101
+ super().__init__()
102
+ if not isinstance(drop_paths, Sequence):
103
+ drop_paths = [drop_paths] * depth
104
+
105
+ if downsample:
106
+ self.out_channels = 2 * embed_dims
107
+ self.downsample = PatchMerging(
108
+ in_channels=embed_dims,
109
+ out_channels=self.out_channels,
110
+ is_post_norm=is_post_norm_downsample,
111
+ )
112
+ else:
113
+ self.out_channels = embed_dims
114
+ self.downsample = None
115
+
116
+ self.blocks = nn.ModuleList()
117
+ for i in range(depth):
118
+ extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0
119
+ self.blocks.append(
120
+ SwinBlockV2(
121
+ embed_dims=self.out_channels,
122
+ num_heads=num_heads,
123
+ window_size=window_size,
124
+ shift=(i % 2 == 1),
125
+ extra_norm=extra_norm,
126
+ drop_path=drop_paths[i],
127
+ with_cp=with_cp,
128
+ pad_small_map=pad_small_map,
129
+ pretrained_window_size=pretrained_window_size,
130
+ )
131
+ )
132
+
133
+ def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
134
+ if self.downsample is not None:
135
+ x, out_shape = self.downsample(x, in_shape)
136
+ else:
137
+ out_shape = in_shape
138
+
139
+ for block in self.blocks:
140
+ x = block(x, out_shape)
141
+ return x, out_shape
142
+
143
+
144
+ class ProjMHSA(nn.Module):
145
+ """Projected multi-head self-attention used in SkySense++ HR backbone."""
146
+
147
+ def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True):
148
+ super().__init__()
149
+ self.proj_in = nn.Linear(embed_dims, proj_dims)
150
+ self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias)
151
+ self.proj_out = nn.Linear(proj_dims, embed_dims)
152
+
153
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
154
+ x = self.proj_in(x)
155
+ x, _ = self.attn(x, x, x)
156
+ return self.proj_out(x)
157
+
158
+
159
+ class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel):
160
+ config_class = SkySensePlusPlusSwinV2MSLConfig
161
+ base_model_prefix = "skysensepp_swinv2_msl"
162
+ supports_gradient_checkpointing = True
163
+
164
+ def _init_weights(self, module):
165
+ if isinstance(module, nn.Linear):
166
+ nn.init.trunc_normal_(module.weight, std=0.02)
167
+ if module.bias is not None:
168
+ nn.init.zeros_(module.bias)
169
+ elif isinstance(module, nn.LayerNorm):
170
+ nn.init.ones_(module.weight)
171
+ nn.init.zeros_(module.bias)
172
+ elif isinstance(module, nn.Conv2d):
173
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
174
+ if module.bias is not None:
175
+ nn.init.zeros_(module.bias)
176
+
177
+
178
+ class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel):
179
+ """SkySense++ HR backbone with semantic vocabulary and annotation conditioning."""
180
+
181
+ def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig):
182
+ super().__init__(config)
183
+
184
+ self.num_layers = len(config.depths)
185
+ self.out_indices = config.out_indices
186
+ self.merge_stage = config.merge_stage
187
+ self.use_attn = config.use_attn
188
+ self.patch_size = config.patch_size
189
+
190
+ if isinstance(config.window_size, int):
191
+ window_sizes = [config.window_size] * self.num_layers
192
+ else:
193
+ window_sizes = list(config.window_size)
194
+
195
+ self.patch_embed = PatchEmbed(
196
+ in_channels=config.in_channels,
197
+ embed_dims=config.embed_dims,
198
+ kernel_size=config.patch_size,
199
+ stride=config.patch_size,
200
+ norm_layer=nn.LayerNorm,
201
+ input_size=config.img_size,
202
+ )
203
+
204
+ self.use_abs_pos_embed = config.use_abs_pos_embed
205
+ if self.use_abs_pos_embed:
206
+ patch_resolution = self.patch_embed.init_out_size
207
+ num_patches = patch_resolution[0] * patch_resolution[1]
208
+ self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
209
+
210
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
211
+
212
+ total_depth = sum(config.depths)
213
+ if total_depth > 1:
214
+ dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
215
+ else:
216
+ dpr = [0.0]
217
+
218
+ self.stages = nn.ModuleList()
219
+ embed_dims_list = [config.embed_dims]
220
+ for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)):
221
+ stage = SwinBlockV2Sequence(
222
+ embed_dims=embed_dims_list[-1],
223
+ depth=depth,
224
+ num_heads=num_heads,
225
+ window_size=window_sizes[i],
226
+ downsample=(i > 0),
227
+ drop_paths=dpr[:depth],
228
+ with_cp=config.with_cp,
229
+ pad_small_map=config.pad_small_map,
230
+ extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
231
+ pretrained_window_size=config.pretrained_window_sizes[i],
232
+ is_post_norm_downsample=config.is_post_norm_downsample,
233
+ )
234
+ self.stages.append(stage)
235
+ dpr = dpr[depth:]
236
+ embed_dims_list.append(stage.out_channels)
237
+
238
+ for i in self.out_indices:
239
+ self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1]))
240
+
241
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
242
+ self.vocabulary_token = nn.Parameter(
243
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
244
+ )
245
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
246
+
247
+ if self.use_attn:
248
+ self.attn1 = ProjMHSA(352, 256, num_heads=16)
249
+ self.attn2 = ProjMHSA(704, 512, num_heads=16)
250
+ self.attn3 = ProjMHSA(1408, 1024, num_heads=16)
251
+ self.norm_attn = nn.LayerNorm(1408)
252
+
253
+ self.post_init()
254
+
255
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
256
+ batch_size, height, width = anno_img.shape
257
+ ann_token = torch.index_select(
258
+ self.vocabulary_token, 0, anno_img.reshape(-1)
259
+ ).reshape(batch_size, height, width, -1)
260
+
261
+ num_patch_h = height // self.patch_size
262
+ num_patch_w = width // self.patch_size
263
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
264
+ weight = (
265
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
266
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
267
+ .reshape(1, height, width, 1)
268
+ )
269
+ ann_token = ann_token * weight
270
+ ann_token = F.avg_pool2d(
271
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
272
+ )
273
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
274
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
275
+ )
276
+
277
+ def forward(
278
+ self,
279
+ pixel_values: torch.Tensor,
280
+ annotation: torch.Tensor,
281
+ mask: Optional[torch.Tensor] = None,
282
+ output_hidden_states: Optional[bool] = None,
283
+ return_dict: Optional[bool] = None,
284
+ ) -> Union[Tuple, BaseModelOutput]:
285
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
286
+
287
+ x, hw_shape = self.patch_embed(pixel_values)
288
+ y = self.create_ann_token(annotation)
289
+ batch_size, num_tokens, channels = y.shape
290
+
291
+ if mask is not None:
292
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
293
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
294
+ y = y * (1.0 - weight) + mask_tokens * weight
295
+
296
+ if self.merge_stage == 0:
297
+ x = (x + y) * 0.5
298
+ else:
299
+ x = x.reshape(batch_size, *hw_shape, channels)
300
+ y = y.reshape(batch_size, *hw_shape, channels)
301
+ x = torch.cat((x, y), dim=2)
302
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
303
+ x = x.reshape(batch_size, -1, channels)
304
+
305
+ if self.use_abs_pos_embed:
306
+ x = x + self.absolute_pos_embed
307
+ x = self.drop_after_pos(x)
308
+
309
+ all_hidden_states = () if output_hidden_states else None
310
+ feature_maps = []
311
+ merge_idx = self.merge_stage - 1
312
+
313
+ for i, stage in enumerate(self.stages):
314
+ x, hw_shape = stage(x, hw_shape)
315
+ if i == merge_idx:
316
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
317
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
318
+ x = x.reshape(batch_size, -1, x.shape[-1])
319
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
320
+
321
+ if self.use_attn:
322
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
323
+ if i <= len(attention_blocks) - 1:
324
+ x = x + attention_blocks[i](x)
325
+ if i == len(attention_blocks) - 1:
326
+ x = self.norm_attn(x)
327
+
328
+ if output_hidden_states:
329
+ all_hidden_states = all_hidden_states + (x,)
330
+
331
+ if i in self.out_indices:
332
+ norm_layer = getattr(self, f"norm{i}")
333
+ out = norm_layer(x)
334
+ out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous()
335
+ feature_maps.append(out)
336
+
337
+ if not return_dict:
338
+ return tuple(feature_maps)
339
+
340
+ return BaseModelOutput(
341
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
342
+ hidden_states=all_hidden_states,
343
+ )
skysensepp-fewshot-release/modeling_skysensepp_vit_msl.py ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ import math
4
+ from typing import Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig
14
+ from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
15
+
16
+
17
+ class TransformerEncoderLayer(nn.Module):
18
+ def __init__(
19
+ self,
20
+ embed_dims: int,
21
+ num_heads: int,
22
+ feedforward_channels: int,
23
+ drop_rate: float = 0.0,
24
+ attn_drop_rate: float = 0.0,
25
+ drop_path_rate: float = 0.0,
26
+ num_fcs: int = 2,
27
+ qkv_bias: bool = True,
28
+ with_cp: bool = False,
29
+ ):
30
+ super().__init__()
31
+ self.with_cp = with_cp
32
+ self.norm1 = nn.LayerNorm(embed_dims)
33
+ self.attn = nn.MultiheadAttention(
34
+ embed_dim=embed_dims,
35
+ num_heads=num_heads,
36
+ dropout=attn_drop_rate,
37
+ bias=qkv_bias,
38
+ batch_first=True,
39
+ )
40
+ self.proj_drop = nn.Dropout(drop_rate)
41
+ self.norm2 = nn.LayerNorm(embed_dims)
42
+ self.ffn = FFN(
43
+ embed_dims=embed_dims,
44
+ feedforward_channels=feedforward_channels,
45
+ num_fcs=num_fcs,
46
+ ffn_drop=drop_rate,
47
+ drop_path=drop_path_rate,
48
+ act_layer=nn.GELU,
49
+ add_identity=True,
50
+ )
51
+ self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
52
+
53
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
54
+ def _inner_forward(x):
55
+ residual = x
56
+ x_norm = self.norm1(x)
57
+ attn_out, _ = self.attn(x_norm, x_norm, x_norm)
58
+ attn_out = self.proj_drop(attn_out)
59
+ x = residual + self.drop_path(attn_out)
60
+ return self.ffn(self.norm2(x), identity=x)
61
+
62
+ if self.with_cp and x.requires_grad:
63
+ return cp.checkpoint(_inner_forward, x, use_reentrant=False)
64
+ return _inner_forward(x)
65
+
66
+
67
+ class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel):
68
+ config_class = SkySensePlusPlusViTMSLConfig
69
+ base_model_prefix = "skysensepp_vit_msl"
70
+ supports_gradient_checkpointing = True
71
+
72
+ def _init_weights(self, module):
73
+ if isinstance(module, nn.Linear):
74
+ nn.init.trunc_normal_(module.weight, std=0.02)
75
+ if module.bias is not None:
76
+ nn.init.zeros_(module.bias)
77
+ elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
78
+ nn.init.ones_(module.weight)
79
+ nn.init.zeros_(module.bias)
80
+ elif isinstance(module, nn.Conv2d):
81
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
82
+ if module.bias is not None:
83
+ nn.init.zeros_(module.bias)
84
+
85
+
86
+ class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel):
87
+ """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning."""
88
+
89
+ def __init__(self, config: SkySensePlusPlusViTMSLConfig):
90
+ super().__init__(config)
91
+
92
+ img_size = to_2tuple(config.img_size)
93
+ self.img_size = img_size
94
+ self.patch_size = config.patch_size
95
+ self.with_cls_token = config.with_cls_token
96
+ self.output_cls_token = config.output_cls_token
97
+ self.merge_stage = config.merge_stage
98
+ self.use_attn = config.use_attn
99
+ self.interpolate_mode = "bicubic"
100
+
101
+ self.patch_embed = PatchEmbed(
102
+ in_channels=config.in_channels,
103
+ embed_dims=config.embed_dims,
104
+ kernel_size=config.patch_size,
105
+ stride=config.patch_size,
106
+ norm_layer=nn.LayerNorm if config.patch_norm else None,
107
+ )
108
+
109
+ num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size)
110
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
111
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
112
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
113
+
114
+ out_indices = list(config.out_indices)
115
+ self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices]
116
+
117
+ num_layers = config.num_layers
118
+ if num_layers > 1:
119
+ dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
120
+ else:
121
+ dpr = [0.0]
122
+
123
+ self.layers = nn.ModuleList()
124
+ for i in range(config.num_layers):
125
+ self.layers.append(
126
+ TransformerEncoderLayer(
127
+ embed_dims=config.embed_dims,
128
+ num_heads=config.num_heads,
129
+ feedforward_channels=config.mlp_ratio * config.embed_dims,
130
+ attn_drop_rate=config.attn_drop_rate,
131
+ drop_rate=config.drop_rate,
132
+ drop_path_rate=dpr[i],
133
+ num_fcs=2,
134
+ qkv_bias=config.qkv_bias,
135
+ with_cp=config.with_cp,
136
+ )
137
+ )
138
+
139
+ self.final_norm = config.final_norm
140
+ if config.final_norm:
141
+ self.norm = nn.LayerNorm(config.embed_dims)
142
+
143
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
144
+ self.vocabulary_token = nn.Parameter(
145
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
146
+ )
147
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
148
+
149
+ if self.use_attn:
150
+ self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
151
+ self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
152
+ self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
153
+ self.norm_attn = nn.LayerNorm(config.embed_dims)
154
+
155
+ self.post_init()
156
+
157
+ @staticmethod
158
+ def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"):
159
+ pos_h, pos_w = pos_shape
160
+ pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :]
161
+ pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
162
+ pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode)
163
+ return torch.flatten(pos_embed_weight, 2).transpose(1, 2)
164
+
165
+ def _pos_embedding(self, patched_img, hw_shape, pos_embed):
166
+ x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
167
+ if x_len != pos_len:
168
+ pos_h = self.img_size[0] // self.patch_size
169
+ pos_w = self.img_size[1] // self.patch_size
170
+ pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode)
171
+ return self.drop_after_pos(patched_img + pos_embed)
172
+
173
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
174
+ batch_size, height, width = anno_img.shape
175
+ ann_token = torch.index_select(
176
+ self.vocabulary_token, 0, anno_img.reshape(-1)
177
+ ).reshape(batch_size, height, width, -1)
178
+
179
+ num_patch_h = height // self.patch_size
180
+ num_patch_w = width // self.patch_size
181
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
182
+ weight = (
183
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
184
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
185
+ .reshape(1, height, width, 1)
186
+ )
187
+ ann_token = ann_token * weight
188
+ ann_token = F.avg_pool2d(
189
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
190
+ )
191
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
192
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
193
+ )
194
+
195
+ def forward(
196
+ self,
197
+ pixel_values: torch.Tensor,
198
+ annotation: torch.Tensor,
199
+ mask: Optional[torch.Tensor] = None,
200
+ output_hidden_states: Optional[bool] = None,
201
+ return_dict: Optional[bool] = None,
202
+ ) -> Union[Tuple, BaseModelOutput]:
203
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
204
+
205
+ x, hw_shape = self.patch_embed(pixel_values)
206
+ y = self.create_ann_token(annotation)
207
+ batch_size, num_tokens, channels = y.shape
208
+
209
+ if mask is not None:
210
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
211
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
212
+ y = y * (1.0 - weight) + mask_tokens * weight
213
+
214
+ if self.merge_stage == 0:
215
+ x = (x + y) * 0.5
216
+ else:
217
+ x = x.reshape(batch_size, *hw_shape, channels)
218
+ y = y.reshape(batch_size, *hw_shape, channels)
219
+ x = torch.cat((x, y), dim=2)
220
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
221
+ x = x.reshape(batch_size, -1, channels)
222
+
223
+ x = self._pos_embedding(x, hw_shape, self.pos_embed)
224
+
225
+ all_hidden_states = () if output_hidden_states else None
226
+ feature_maps = []
227
+ merge_idx = self.merge_stage - 1
228
+
229
+ for i, layer in enumerate(self.layers):
230
+ x = layer(x)
231
+
232
+ if i == merge_idx:
233
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
234
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
235
+ x = x.reshape(batch_size, -1, x.shape[-1])
236
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
237
+
238
+ if self.use_attn:
239
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
240
+ if i <= len(attention_blocks) - 1:
241
+ attn_out, _ = attention_blocks[i](x, x, x)
242
+ x = x + attn_out
243
+ if i == len(attention_blocks) - 1:
244
+ x = self.norm_attn(x)
245
+
246
+ if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm:
247
+ x = self.norm(x)
248
+
249
+ if output_hidden_states:
250
+ all_hidden_states = all_hidden_states + (x,)
251
+
252
+ if i in self.out_indices:
253
+ out = x
254
+ out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous()
255
+ if self.output_cls_token:
256
+ out = [out, x[:, 0]]
257
+ feature_maps.append(out)
258
+
259
+ if not return_dict:
260
+ return tuple(feature_maps)
261
+
262
+ return BaseModelOutput(
263
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
264
+ hidden_states=all_hidden_states,
265
+ )
skysensepp-fewshot-release/modeling_utils.py ADDED
@@ -0,0 +1,557 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense: Pure PyTorch + HuggingFace Transformers implementation.
2
+
3
+ Shared utility modules used across SkySense model implementations.
4
+ """
5
+
6
+ import math
7
+ from typing import Optional, Tuple
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+
14
+ def to_2tuple(x):
15
+ """Convert to a 2-tuple."""
16
+ if isinstance(x, (list, tuple)):
17
+ return tuple(x)
18
+ return (x, x)
19
+
20
+
21
+ class DropPath(nn.Module):
22
+ """Drop paths (stochastic depth) per sample.
23
+
24
+ Args:
25
+ drop_prob (float): Probability of dropping a path. Default: 0.0.
26
+ """
27
+
28
+ def __init__(self, drop_prob: float = 0.0):
29
+ super().__init__()
30
+ self.drop_prob = drop_prob
31
+
32
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
33
+ if self.drop_prob == 0.0 or not self.training:
34
+ return x
35
+ keep_prob = 1 - self.drop_prob
36
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1)
37
+ random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
38
+ random_tensor = torch.floor(random_tensor + keep_prob)
39
+ output = x / keep_prob * random_tensor
40
+ return output
41
+
42
+
43
+ class PatchEmbed(nn.Module):
44
+ """Image to Patch Embedding using Conv2d.
45
+
46
+ Args:
47
+ in_channels (int): Number of input channels. Default: 3.
48
+ embed_dims (int): Embedding dimension. Default: 96.
49
+ kernel_size (int): Kernel size of the projection. Default: 4.
50
+ stride (int): Stride of the projection. Default: 4.
51
+ padding (int): Padding of the projection. Default: 0.
52
+ norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
53
+ input_size (int or tuple or None): Input resolution for calculating output size.
54
+ """
55
+
56
+ def __init__(
57
+ self,
58
+ in_channels: int = 3,
59
+ embed_dims: int = 96,
60
+ kernel_size: int = 4,
61
+ stride: int = 4,
62
+ padding: int = 0,
63
+ norm_layer: Optional[type] = nn.LayerNorm,
64
+ input_size: Optional[int] = None,
65
+ ):
66
+ super().__init__()
67
+ self.projection = nn.Conv2d(
68
+ in_channels, embed_dims,
69
+ kernel_size=kernel_size, stride=stride, padding=padding,
70
+ )
71
+ self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
72
+
73
+ # Compute init output size if input_size is given
74
+ if input_size is not None:
75
+ input_size = to_2tuple(input_size)
76
+ self.init_out_size = (
77
+ (input_size[0] - kernel_size + 2 * padding) // stride + 1,
78
+ (input_size[1] - kernel_size + 2 * padding) // stride + 1,
79
+ )
80
+ else:
81
+ self.init_out_size = None
82
+
83
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
84
+ x = self.projection(x) # (B, C, H, W)
85
+ out_size = (x.shape[2], x.shape[3])
86
+ x = x.flatten(2).transpose(1, 2) # (B, H*W, C)
87
+ x = self.norm(x)
88
+ return x, out_size
89
+
90
+
91
+ class FFN(nn.Module):
92
+ """Feed-Forward Network.
93
+
94
+ Args:
95
+ embed_dims (int): Input dimension.
96
+ feedforward_channels (int): Hidden dimension.
97
+ num_fcs (int): Number of FC layers. Default: 2.
98
+ ffn_drop (float): Dropout rate. Default: 0.0.
99
+ drop_path (float): Drop path rate. Default: 0.0.
100
+ act_layer (nn.Module): Activation layer class. Default: nn.GELU.
101
+ add_identity (bool): Whether to add identity connection. Default: True.
102
+ """
103
+
104
+ def __init__(
105
+ self,
106
+ embed_dims: int,
107
+ feedforward_channels: int,
108
+ num_fcs: int = 2,
109
+ ffn_drop: float = 0.0,
110
+ drop_path: float = 0.0,
111
+ act_layer: type = nn.GELU,
112
+ add_identity: bool = True,
113
+ ):
114
+ super().__init__()
115
+ assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
116
+ self.embed_dims = embed_dims
117
+ self.feedforward_channels = feedforward_channels
118
+ self.add_identity = add_identity
119
+
120
+ layers = []
121
+ in_channels = embed_dims
122
+ for i in range(num_fcs - 1):
123
+ layers.append(nn.Linear(in_channels, feedforward_channels))
124
+ layers.append(act_layer())
125
+ layers.append(nn.Dropout(ffn_drop))
126
+ in_channels = feedforward_channels
127
+ layers.append(nn.Linear(feedforward_channels, embed_dims))
128
+ layers.append(nn.Dropout(ffn_drop))
129
+ self.layers = nn.Sequential(*layers)
130
+
131
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
132
+
133
+ def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
134
+ out = self.layers(x)
135
+ out = self.drop_path(out)
136
+ if self.add_identity:
137
+ if identity is None:
138
+ identity = x
139
+ out = out + identity
140
+ return out
141
+
142
+
143
+ class WindowMSAV2(nn.Module):
144
+ """Window-based Multi-head Self-Attention for Swin Transformer V2.
145
+
146
+ Uses cosine attention and log-spaced continuous position bias (log-CPB).
147
+
148
+ Args:
149
+ embed_dims (int): Number of input channels.
150
+ num_heads (int): Number of attention heads.
151
+ window_size (tuple[int]): Window size (Wh, Ww).
152
+ pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
153
+ qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
154
+ attn_drop (float): Attention dropout rate. Default: 0.0.
155
+ proj_drop (float): Output projection dropout rate. Default: 0.0.
156
+ """
157
+
158
+ def __init__(
159
+ self,
160
+ embed_dims: int,
161
+ num_heads: int,
162
+ window_size: Tuple[int, int],
163
+ pretrained_window_size: Tuple[int, int] = (0, 0),
164
+ qkv_bias: bool = True,
165
+ attn_drop: float = 0.0,
166
+ proj_drop: float = 0.0,
167
+ ):
168
+ super().__init__()
169
+ self.embed_dims = embed_dims
170
+ self.num_heads = num_heads
171
+ self.window_size = window_size
172
+ self.pretrained_window_size = pretrained_window_size
173
+
174
+ self.logit_scale = nn.Parameter(
175
+ torch.log(10 * torch.ones((num_heads, 1, 1))))
176
+
177
+ # MLP for continuous relative position bias (log-CPB)
178
+ self.cpb_mlp = nn.Sequential(
179
+ nn.Linear(2, 512, bias=True),
180
+ nn.ReLU(inplace=True),
181
+ nn.Linear(512, num_heads, bias=False),
182
+ )
183
+
184
+ # Build relative coords table
185
+ self._build_relative_coords_table()
186
+ # Build relative position index
187
+ self._build_relative_position_index()
188
+
189
+ self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
190
+ if qkv_bias:
191
+ self.q_bias = nn.Parameter(torch.zeros(embed_dims))
192
+ self.v_bias = nn.Parameter(torch.zeros(embed_dims))
193
+ else:
194
+ self.q_bias = None
195
+ self.v_bias = None
196
+
197
+ self.attn_drop = nn.Dropout(attn_drop)
198
+ self.proj = nn.Linear(embed_dims, embed_dims)
199
+ self.proj_drop = nn.Dropout(proj_drop)
200
+ self.softmax = nn.Softmax(dim=-1)
201
+
202
+ def _build_relative_coords_table(self):
203
+ """Build the relative coordinates table for log-CPB."""
204
+ Wh, Ww = self.window_size
205
+ # Table of relative coordinates
206
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
207
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
208
+ coords_table = torch.stack(
209
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
210
+ ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2)
211
+
212
+ # Normalize to [-1, 1] and apply log-scale
213
+ if self.pretrained_window_size[0] > 0:
214
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
215
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
216
+ else:
217
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
218
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
219
+ coords_table *= 8 # normalize to -8, 8
220
+ coords_table = (
221
+ torch.sign(coords_table)
222
+ * torch.log2(torch.abs(coords_table) + 1.0)
223
+ / math.log2(8)
224
+ )
225
+ self.register_buffer("relative_coords_table", coords_table)
226
+
227
+ def _build_relative_position_index(self):
228
+ """Build the pairwise relative position index for each window token."""
229
+ Wh, Ww = self.window_size
230
+ coords_h = torch.arange(Wh)
231
+ coords_w = torch.arange(Ww)
232
+ coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
233
+ coords_flatten = coords.view(2, -1)
234
+
235
+ relative_coords = (
236
+ coords_flatten[:, :, None] - coords_flatten[:, None, :]
237
+ ) # (2, Wh*Ww, Wh*Ww)
238
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
239
+ relative_coords[:, :, 0] += Wh - 1
240
+ relative_coords[:, :, 1] += Ww - 1
241
+ relative_coords[:, :, 0] *= 2 * Ww - 1
242
+ relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww)
243
+ self.register_buffer("relative_position_index", relative_position_index)
244
+
245
+ def _compute_position_bias(self, N):
246
+ """Compute relative position bias, supporting dynamic window sizes.
247
+
248
+ The log-CPB (Continuous Position Bias) MLP can generalize to any window
249
+ size by computing bias from normalized relative coordinates.
250
+ """
251
+ init_N = self.window_size[0] * self.window_size[1]
252
+ if N == init_N:
253
+ # Use pre-built tables
254
+ relative_position_bias_table = self.cpb_mlp(
255
+ self.relative_coords_table
256
+ ).view(-1, self.num_heads)
257
+ relative_position_bias = relative_position_bias_table[
258
+ self.relative_position_index.view(-1)
259
+ ].view(N, N, -1)
260
+ else:
261
+ # Dynamic: compute for actual window size on-the-fly
262
+ Wh = Ww = int(math.sqrt(N))
263
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
264
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
265
+ coords_table = torch.stack(
266
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
267
+ ).flatten(1).transpose(0, 1).unsqueeze(0)
268
+ if self.pretrained_window_size[0] > 0:
269
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
270
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
271
+ else:
272
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
273
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
274
+ coords_table *= 8
275
+ coords_table = (
276
+ torch.sign(coords_table)
277
+ * torch.log2(torch.abs(coords_table) + 1.0)
278
+ / math.log2(8)
279
+ )
280
+ # Build position index for actual window size
281
+ ch = torch.arange(Wh, device=self.logit_scale.device)
282
+ cw = torch.arange(Ww, device=self.logit_scale.device)
283
+ coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
284
+ coords_flat = coords.view(2, -1)
285
+ rel = coords_flat[:, :, None] - coords_flat[:, None, :]
286
+ rel = rel.permute(1, 2, 0).contiguous()
287
+ rel[:, :, 0] += Wh - 1
288
+ rel[:, :, 1] += Ww - 1
289
+ rel[:, :, 0] *= 2 * Ww - 1
290
+ pos_index = rel.sum(-1)
291
+
292
+ bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
293
+ relative_position_bias = bias_table[
294
+ pos_index.view(-1)
295
+ ].view(N, N, -1)
296
+
297
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
298
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
299
+ return relative_position_bias
300
+
301
+ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
302
+ """
303
+ Args:
304
+ x: (num_windows*B, N, C) where N = Wh*Ww
305
+ mask: (num_windows, N, N) or None
306
+ """
307
+ B_, N, C = x.shape
308
+
309
+ # Compute QKV with bias
310
+ if self.q_bias is not None:
311
+ qkv_bias = torch.cat(
312
+ (self.q_bias,
313
+ torch.zeros_like(self.v_bias, requires_grad=False),
314
+ self.v_bias))
315
+ qkv = F.linear(x, self.qkv.weight, qkv_bias)
316
+ else:
317
+ qkv = self.qkv(x)
318
+
319
+ qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
320
+ qkv = qkv.permute(2, 0, 3, 1, 4)
321
+ q, k, v = qkv.unbind(0)
322
+
323
+ # Cosine attention
324
+ attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
325
+ logit_scale = torch.clamp(
326
+ self.logit_scale, max=math.log(1.0 / 0.01)
327
+ ).exp()
328
+ attn = attn * logit_scale
329
+
330
+ # Log-CPB relative position bias (supports dynamic window sizes)
331
+ relative_position_bias = self._compute_position_bias(N)
332
+ attn = attn + relative_position_bias.unsqueeze(0)
333
+
334
+ if mask is not None:
335
+ nW = mask.shape[0]
336
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
337
+ attn = attn + mask.unsqueeze(1).unsqueeze(0)
338
+ attn = attn.view(-1, self.num_heads, N, N)
339
+
340
+ attn = self.softmax(attn)
341
+ attn = self.attn_drop(attn)
342
+
343
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
344
+ x = self.proj(x)
345
+ x = self.proj_drop(x)
346
+ return x
347
+
348
+
349
+ class ShiftWindowMSA(nn.Module):
350
+ """Shifted Window Multi-head Self-Attention.
351
+
352
+ Args:
353
+ embed_dims (int): Number of input channels.
354
+ num_heads (int): Number of attention heads.
355
+ window_size (int): Window size.
356
+ shift_size (int): Shift size for SW-MSA. Default: 0.
357
+ attn_drop (float): Attention dropout rate. Default: 0.0.
358
+ proj_drop (float): Projection dropout rate. Default: 0.0.
359
+ drop_path (float): Drop path rate. Default: 0.0.
360
+ pad_small_map (bool): Pad small feature maps to window size. Default: False.
361
+ pretrained_window_size (int): Pretrained window size. Default: 0.
362
+ """
363
+
364
+ def __init__(
365
+ self,
366
+ embed_dims: int,
367
+ num_heads: int,
368
+ window_size: int,
369
+ shift_size: int = 0,
370
+ attn_drop: float = 0.0,
371
+ proj_drop: float = 0.0,
372
+ drop_path: float = 0.0,
373
+ pad_small_map: bool = False,
374
+ pretrained_window_size: int = 0,
375
+ ):
376
+ super().__init__()
377
+ self.window_size = window_size
378
+ self.shift_size = shift_size
379
+ self.pad_small_map = pad_small_map
380
+
381
+ self.w_msa = WindowMSAV2(
382
+ embed_dims=embed_dims,
383
+ num_heads=num_heads,
384
+ window_size=to_2tuple(window_size),
385
+ pretrained_window_size=to_2tuple(pretrained_window_size),
386
+ attn_drop=attn_drop,
387
+ proj_drop=proj_drop,
388
+ )
389
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
390
+
391
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
392
+ B, L, C = x.shape
393
+ H, W = hw_shape
394
+ assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
395
+
396
+ x = x.view(B, H, W, C)
397
+
398
+ window_size = self.window_size
399
+ shift_size = self.shift_size
400
+
401
+ # Pad or shrink window
402
+ if self.pad_small_map:
403
+ pad_r = (window_size - W % window_size) % window_size
404
+ pad_b = (window_size - H % window_size) % window_size
405
+ x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
406
+ _, Hp, Wp, _ = x.shape
407
+ else:
408
+ Hp, Wp = H, W
409
+ if window_size > Hp:
410
+ window_size = Hp
411
+ shift_size = 0
412
+ if window_size > Wp:
413
+ window_size = Wp
414
+ shift_size = 0
415
+
416
+ # Compute attention mask for SW-MSA
417
+ attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
418
+
419
+ # Cyclic shift
420
+ if shift_size > 0:
421
+ x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
422
+
423
+ # Partition windows
424
+ x_windows = self._window_partition(x, window_size)
425
+ # (num_windows*B, window_size*window_size, C)
426
+
427
+ # W-MSA/SW-MSA
428
+ attn_windows = self.w_msa(x_windows, mask=attn_mask)
429
+
430
+ # Merge windows
431
+ x = self._window_reverse(attn_windows, window_size, Hp, Wp)
432
+
433
+ # Reverse cyclic shift
434
+ if shift_size > 0:
435
+ x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
436
+
437
+ if self.pad_small_map and (pad_r > 0 or pad_b > 0):
438
+ x = x[:, :H, :W, :].contiguous()
439
+
440
+ x = x.view(B, H * W, C)
441
+ x = self.drop_path(x)
442
+ return x
443
+
444
+ @staticmethod
445
+ def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
446
+ """Partition into non-overlapping windows."""
447
+ B, H, W, C = x.shape
448
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
449
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
450
+ windows = windows.view(-1, window_size * window_size, C)
451
+ return windows
452
+
453
+ @staticmethod
454
+ def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
455
+ """Reverse window partition."""
456
+ B_nW = windows.shape[0]
457
+ nH = H // window_size
458
+ nW = W // window_size
459
+ B = B_nW // (nH * nW)
460
+ x = windows.view(B, nH, nW, window_size, window_size, -1)
461
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
462
+ x = x.view(B, H, W, -1)
463
+ return x
464
+
465
+ @staticmethod
466
+ def _compute_attn_mask(H, W, window_size, shift_size, device):
467
+ """Compute attention mask for shifted window attention."""
468
+ if shift_size <= 0:
469
+ return None
470
+ img_mask = torch.zeros((1, H, W, 1), device=device)
471
+ h_slices = (
472
+ slice(0, -window_size),
473
+ slice(-window_size, -shift_size),
474
+ slice(-shift_size, None),
475
+ )
476
+ w_slices = (
477
+ slice(0, -window_size),
478
+ slice(-window_size, -shift_size),
479
+ slice(-shift_size, None),
480
+ )
481
+ cnt = 0
482
+ for h in h_slices:
483
+ for w in w_slices:
484
+ img_mask[:, h, w, :] = cnt
485
+ cnt += 1
486
+
487
+ # Partition mask
488
+ mask_windows = img_mask.view(
489
+ 1, H // window_size, window_size, W // window_size, window_size, 1
490
+ )
491
+ mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
492
+ mask_windows = mask_windows.view(-1, window_size * window_size)
493
+
494
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
495
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
496
+ attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
497
+ return attn_mask
498
+
499
+
500
+ class PatchMerging(nn.Module):
501
+ """Patch Merging Layer for downsampling (2x).
502
+
503
+ Args:
504
+ in_channels (int): Input channels.
505
+ out_channels (int): Output channels.
506
+ norm_layer (type): Normalization layer. Default: nn.LayerNorm.
507
+ is_post_norm (bool): Apply norm after linear. Default: True.
508
+ """
509
+
510
+ def __init__(
511
+ self,
512
+ in_channels: int,
513
+ out_channels: int,
514
+ norm_layer: type = nn.LayerNorm,
515
+ is_post_norm: bool = True,
516
+ ):
517
+ super().__init__()
518
+ self.in_channels = in_channels
519
+ self.out_channels = out_channels
520
+ self.is_post_norm = is_post_norm
521
+ self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
522
+ if is_post_norm:
523
+ self.norm = norm_layer(out_channels)
524
+ else:
525
+ self.norm = norm_layer(4 * in_channels)
526
+
527
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
528
+ B, L, C = x.shape
529
+ H, W = hw_shape
530
+ assert L == H * W
531
+
532
+ x = x.view(B, H, W, C)
533
+
534
+ # Pad if needed
535
+ pad_h = H % 2
536
+ pad_w = W % 2
537
+ if pad_h or pad_w:
538
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
539
+
540
+ x0 = x[:, 0::2, 0::2, :]
541
+ x1 = x[:, 1::2, 0::2, :]
542
+ x2 = x[:, 0::2, 1::2, :]
543
+ x3 = x[:, 1::2, 1::2, :]
544
+ x = torch.cat([x0, x1, x2, x3], dim=-1)
545
+
546
+ out_h = (H + pad_h) // 2
547
+ out_w = (W + pad_w) // 2
548
+ x = x.view(B, out_h * out_w, 4 * C)
549
+
550
+ if self.is_post_norm:
551
+ x = self.reduction(x)
552
+ x = self.norm(x)
553
+ else:
554
+ x = self.norm(x)
555
+ x = self.reduction(x)
556
+
557
+ return x, (out_h, out_w)
skysensepp-fewshot-release/pipeline_skysensepp.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Custom HuggingFace pipeline for SkySense++ MSL feature extraction."""
2
+
3
+ from typing import Any, Dict, Optional, Union
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline):
11
+ """Pipeline for SkySense++ MSL backbones.
12
+
13
+ Expects image tensors plus semantic annotation maps (class indices).
14
+ """
15
+
16
+ def _sanitize_parameters(
17
+ self,
18
+ annotation=None,
19
+ mask=None,
20
+ output_hidden_states=None,
21
+ **kwargs,
22
+ ):
23
+ preprocess_params = {}
24
+ forward_params = {}
25
+ postprocess_params = {}
26
+
27
+ if annotation is not None:
28
+ preprocess_params["annotation"] = annotation
29
+ if mask is not None:
30
+ forward_params["mask"] = mask
31
+ if output_hidden_states is not None:
32
+ forward_params["output_hidden_states"] = output_hidden_states
33
+
34
+ return preprocess_params, forward_params, postprocess_params
35
+
36
+ def preprocess(
37
+ self,
38
+ pixel_values: Any,
39
+ annotation: Optional[Any] = None,
40
+ **kwargs,
41
+ ) -> Dict[str, torch.Tensor]:
42
+ if isinstance(pixel_values, dict):
43
+ annotation = pixel_values.get("annotation", annotation)
44
+ pixel_values = pixel_values.get("pixel_values", pixel_values)
45
+
46
+ if isinstance(pixel_values, np.ndarray):
47
+ pixel_values = torch.from_numpy(pixel_values).float()
48
+ elif not isinstance(pixel_values, torch.Tensor):
49
+ raise TypeError(
50
+ f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}"
51
+ )
52
+
53
+ if annotation is None:
54
+ raise ValueError("SkySense++ MSL models require an `annotation` semantic map.")
55
+
56
+ if isinstance(annotation, np.ndarray):
57
+ annotation = torch.from_numpy(annotation).long()
58
+ elif not isinstance(annotation, torch.Tensor):
59
+ raise TypeError(
60
+ f"Expected tensor or ndarray for annotation, got {type(annotation)}"
61
+ )
62
+
63
+ if pixel_values.ndim == 3:
64
+ pixel_values = pixel_values.unsqueeze(0)
65
+ if annotation.ndim == 2:
66
+ annotation = annotation.unsqueeze(0)
67
+
68
+ return {"pixel_values": pixel_values, "annotation": annotation}
69
+
70
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
71
+ with torch.no_grad():
72
+ outputs = self.model(
73
+ pixel_values=model_inputs["pixel_values"],
74
+ annotation=model_inputs["annotation"],
75
+ mask=kwargs.get("mask"),
76
+ output_hidden_states=kwargs.get("output_hidden_states", False),
77
+ return_dict=True,
78
+ )
79
+ return {"outputs": outputs}
80
+
81
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
82
+ outputs = model_outputs["outputs"]
83
+ result = {"last_hidden_state": outputs.last_hidden_state}
84
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
85
+ result["hidden_states"] = outputs.hidden_states
86
+ return result
skysensepp-fewshot-release/pipeline_skysensepp_fewshot.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """HuggingFace pipeline for SkySense++ few-shot / 1-shot segmentation."""
2
+
3
+ from typing import Any, Dict, Optional, Union
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusFewShotPipeline(Pipeline):
11
+ """1-shot segmentation pipeline for the full SkySense++ release model.
12
+
13
+ Expects vertically stacked prompt+query tensors:
14
+ - ``hr_img``: (B, 3, 1024, 512) — prompt on top, query on bottom
15
+ - ``s2_img`` / ``s1_img``: (B, C, seq, 32, 32) — stacked along height
16
+ - ``targets``: RGB annotation map (B, 3, 1024, 512), ImageNet-normalized
17
+ - ``anno_mask``: (B, 8, 4) with bottom half = 1 (query region)
18
+ """
19
+
20
+ def _sanitize_parameters(
21
+ self,
22
+ s2_img=None,
23
+ s1_img=None,
24
+ targets=None,
25
+ anno_mask=None,
26
+ modality_flags=None,
27
+ extract_query_only=None,
28
+ **kwargs,
29
+ ):
30
+ preprocess_params = {}
31
+ forward_params = {}
32
+ postprocess_params = {"extract_query_only": True if extract_query_only is None else extract_query_only}
33
+
34
+ if s2_img is not None:
35
+ preprocess_params["s2_img"] = s2_img
36
+ if s1_img is not None:
37
+ preprocess_params["s1_img"] = s1_img
38
+ if targets is not None:
39
+ preprocess_params["targets"] = targets
40
+ if anno_mask is not None:
41
+ preprocess_params["anno_mask"] = anno_mask
42
+ if modality_flags is not None:
43
+ forward_params["modality_flags"] = modality_flags
44
+
45
+ return preprocess_params, forward_params, postprocess_params
46
+
47
+ def _to_tensor(self, value: Any, dtype: torch.dtype) -> torch.Tensor:
48
+ if isinstance(value, np.ndarray):
49
+ return torch.from_numpy(value).to(dtype=dtype)
50
+ if isinstance(value, torch.Tensor):
51
+ return value.to(dtype=dtype)
52
+ raise TypeError(f"Expected tensor or ndarray, got {type(value)}")
53
+
54
+ def preprocess(
55
+ self,
56
+ hr_img: Any,
57
+ s2_img: Optional[Any] = None,
58
+ s1_img: Optional[Any] = None,
59
+ targets: Optional[Any] = None,
60
+ anno_mask: Optional[Any] = None,
61
+ **kwargs,
62
+ ) -> Dict[str, torch.Tensor]:
63
+ if isinstance(hr_img, dict):
64
+ payload = hr_img
65
+ hr_img = payload.get("hr_img", payload.get("pixel_values"))
66
+ s2_img = payload.get("s2_img", s2_img)
67
+ s1_img = payload.get("s1_img", s1_img)
68
+ targets = payload.get("targets", targets)
69
+ anno_mask = payload.get("anno_mask", anno_mask)
70
+
71
+ hr_img = self._to_tensor(hr_img, torch.float32)
72
+ if hr_img.ndim == 3:
73
+ hr_img = hr_img.unsqueeze(0)
74
+
75
+ if s2_img is None or s1_img is None or targets is None:
76
+ raise ValueError("Few-shot pipeline requires hr_img, s2_img, s1_img, and targets.")
77
+
78
+ s2_img = self._to_tensor(s2_img, torch.float32)
79
+ s1_img = self._to_tensor(s1_img, torch.float32)
80
+ targets = self._to_tensor(targets, torch.float32)
81
+
82
+ if s2_img.ndim == 4:
83
+ s2_img = s2_img.unsqueeze(0)
84
+ if s1_img.ndim == 4:
85
+ s1_img = s1_img.unsqueeze(0)
86
+ if targets.ndim == 3:
87
+ targets = targets.unsqueeze(0)
88
+
89
+ if anno_mask is None:
90
+ batch_size = hr_img.shape[0]
91
+ anno_mask = torch.zeros(batch_size, 8, 4, dtype=torch.long)
92
+ anno_mask[:, 4:, :] = 1
93
+ else:
94
+ anno_mask = self._to_tensor(anno_mask, torch.long)
95
+ if anno_mask.ndim == 2:
96
+ anno_mask = anno_mask.unsqueeze(0)
97
+
98
+ return {
99
+ "hr_img": hr_img,
100
+ "s2_img": s2_img,
101
+ "s1_img": s1_img,
102
+ "targets": targets,
103
+ "anno_mask": anno_mask,
104
+ }
105
+
106
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
107
+ with torch.no_grad():
108
+ outputs = self.model(
109
+ hr_img=model_inputs["hr_img"],
110
+ s2_img=model_inputs["s2_img"],
111
+ s1_img=model_inputs["s1_img"],
112
+ targets=model_inputs["targets"],
113
+ anno_mask=model_inputs["anno_mask"],
114
+ modality_flags=kwargs.get("modality_flags"),
115
+ return_dict=True,
116
+ )
117
+ return {"outputs": outputs}
118
+
119
+ def postprocess(
120
+ self,
121
+ model_outputs: Dict[str, Any],
122
+ extract_query_only: bool = True,
123
+ ) -> Dict[str, Any]:
124
+ outputs = model_outputs["outputs"]
125
+ logits = outputs.logits
126
+ if extract_query_only and logits is not None:
127
+ logits = logits[:, :, logits.shape[2] // 2 :, :]
128
+ return {
129
+ "logits": logits,
130
+ "mapped_targets": outputs.mapped_targets,
131
+ "idx_2_color": outputs.idx_2_color,
132
+ }
skysensepp-fewshot-release/pipeline_skysensepp_fusion.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Optional pipeline for SkySense++ fusion neck."""
2
+
3
+ from typing import Any, Dict
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusFusionNeckPipeline(Pipeline):
11
+ """Pipeline for the optional SkySense++ fusion neck module.
12
+
13
+ Expects concatenated multi-modal tokens per spatial location:
14
+ ``(batch, num_modalities, input_dims)``.
15
+ """
16
+
17
+ def _sanitize_parameters(self, output_hidden_states=None, **kwargs):
18
+ preprocess_params = {}
19
+ forward_params = {}
20
+ postprocess_params = {}
21
+ if output_hidden_states is not None:
22
+ forward_params["output_hidden_states"] = output_hidden_states
23
+ return preprocess_params, forward_params, postprocess_params
24
+
25
+ def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]:
26
+ if isinstance(hidden_states, dict):
27
+ hidden_states = hidden_states["hidden_states"]
28
+
29
+ if isinstance(hidden_states, np.ndarray):
30
+ hidden_states = torch.from_numpy(hidden_states).float()
31
+ elif not isinstance(hidden_states, torch.Tensor):
32
+ raise TypeError(
33
+ f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}"
34
+ )
35
+ if hidden_states.ndim == 2:
36
+ hidden_states = hidden_states.unsqueeze(0)
37
+ return {"hidden_states": hidden_states}
38
+
39
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
40
+ with torch.no_grad():
41
+ outputs = self.model(
42
+ hidden_states=model_inputs["hidden_states"],
43
+ output_hidden_states=kwargs.get("output_hidden_states", False),
44
+ return_dict=True,
45
+ )
46
+ return {"outputs": outputs}
47
+
48
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
49
+ outputs = model_outputs["outputs"]
50
+ result = {"pooler_output": outputs.pooler_output}
51
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
52
+ result["hidden_states"] = outputs.hidden_states
53
+ return result
skysensepp-fusion-neck/config.json ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "return_dict": true,
3
+ "output_hidden_states": false,
4
+ "dtype": "float32",
5
+ "chunk_size_feed_forward": 0,
6
+ "is_encoder_decoder": false,
7
+ "architectures": [
8
+ "SkySensePlusPlusFusionNeckModel"
9
+ ],
10
+ "id2label": {
11
+ "0": "LABEL_0",
12
+ "1": "LABEL_1"
13
+ },
14
+ "label2id": {
15
+ "LABEL_0": 0,
16
+ "LABEL_1": 1
17
+ },
18
+ "problem_type": null,
19
+ "_name_or_path": "",
20
+ "transformers_version": "5.0.0",
21
+ "input_dims": 2816,
22
+ "embed_dims": 1024,
23
+ "num_layers": 24,
24
+ "num_heads": 16,
25
+ "mlp_ratio": 4,
26
+ "qkv_bias": true,
27
+ "drop_rate": 0.0,
28
+ "attn_drop_rate": 0.0,
29
+ "drop_path_rate": 0.3,
30
+ "with_cls_token": true,
31
+ "output_cls_token": true,
32
+ "with_cp": false,
33
+ "model_type": "skysensepp_fusion_neck",
34
+ "output_attentions": false,
35
+ "auto_map": {
36
+ "AutoConfig": "configuration_skysensepp.SkySensePlusPlusFusionNeckConfig",
37
+ "AutoModel": "modeling_skysensepp_fusion_neck.SkySensePlusPlusFusionNeckModel"
38
+ },
39
+ "custom_pipelines": {
40
+ "skysensepp-fusion": {
41
+ "impl": "pipeline_skysensepp_fusion.SkySensePlusPlusFusionNeckPipeline",
42
+ "pt": [
43
+ "AutoModel"
44
+ ]
45
+ }
46
+ }
47
+ }
skysensepp-fusion-neck/configuration_skysensepp.py ADDED
@@ -0,0 +1,165 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration classes for SkySense++ MSL backbones."""
2
+
3
+ from transformers import PretrainedConfig
4
+
5
+
6
+ class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig):
7
+ """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical)."""
8
+
9
+ model_type = "skysensepp_swinv2_msl"
10
+
11
+ arch_zoo = {
12
+ "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
13
+ "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
14
+ "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0},
15
+ "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0},
16
+ "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6},
17
+ "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6},
18
+ }
19
+
20
+ def __init__(
21
+ self,
22
+ arch="huge",
23
+ img_size=512,
24
+ patch_size=4,
25
+ in_channels=3,
26
+ window_size=8,
27
+ drop_rate=0.0,
28
+ drop_path_rate=0.2,
29
+ out_indices=(0, 1, 2, 3),
30
+ use_abs_pos_embed=False,
31
+ with_cp=False,
32
+ pad_small_map=False,
33
+ pretrained_window_sizes=(0, 0, 0, 0),
34
+ is_post_norm_downsample=True,
35
+ vocabulary_size=64,
36
+ merge_stage=2,
37
+ use_attn=True,
38
+ **kwargs,
39
+ ):
40
+ super().__init__(**kwargs)
41
+
42
+ arch = arch.lower()
43
+ if arch not in self.arch_zoo:
44
+ raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
45
+ arch_settings = self.arch_zoo[arch]
46
+
47
+ self.arch = arch
48
+ self.embed_dims = arch_settings["embed_dims"]
49
+ self.depths = arch_settings["depths"]
50
+ self.num_heads = arch_settings["num_heads"]
51
+ self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"]
52
+
53
+ self.img_size = img_size
54
+ self.patch_size = patch_size
55
+ self.in_channels = in_channels
56
+ self.window_size = window_size
57
+ self.drop_rate = drop_rate
58
+ self.drop_path_rate = drop_path_rate
59
+ self.out_indices = list(out_indices)
60
+ self.use_abs_pos_embed = use_abs_pos_embed
61
+ self.with_cp = with_cp
62
+ self.pad_small_map = pad_small_map
63
+ self.pretrained_window_sizes = list(pretrained_window_sizes)
64
+ self.is_post_norm_downsample = is_post_norm_downsample
65
+
66
+ self.vocabulary_size = vocabulary_size
67
+ self.num_vocabulary_tokens = vocabulary_size + 1
68
+ self.merge_stage = merge_stage
69
+ self.use_attn = use_attn
70
+
71
+
72
+ class SkySensePlusPlusViTMSLConfig(PretrainedConfig):
73
+ """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1)."""
74
+
75
+ model_type = "skysensepp_vit_msl"
76
+
77
+ def __init__(
78
+ self,
79
+ img_size=16,
80
+ patch_size=4,
81
+ in_channels=10,
82
+ embed_dims=1024,
83
+ num_layers=24,
84
+ num_heads=16,
85
+ mlp_ratio=4,
86
+ out_indices=(5, 11, 17, 23),
87
+ qkv_bias=True,
88
+ drop_rate=0.0,
89
+ attn_drop_rate=0.0,
90
+ drop_path_rate=0.3,
91
+ with_cls_token=False,
92
+ output_cls_token=False,
93
+ patch_norm=False,
94
+ final_norm=False,
95
+ with_cp=False,
96
+ vocabulary_size=64,
97
+ merge_stage=4,
98
+ use_attn=False,
99
+ modality="s2",
100
+ **kwargs,
101
+ ):
102
+ super().__init__(**kwargs)
103
+ self.img_size = img_size
104
+ self.patch_size = patch_size
105
+ self.in_channels = in_channels
106
+ self.embed_dims = embed_dims
107
+ self.num_layers = num_layers
108
+ self.num_heads = num_heads
109
+ self.mlp_ratio = mlp_ratio
110
+ self.out_indices = list(out_indices)
111
+ self.qkv_bias = qkv_bias
112
+ self.drop_rate = drop_rate
113
+ self.attn_drop_rate = attn_drop_rate
114
+ self.drop_path_rate = drop_path_rate
115
+ self.with_cls_token = with_cls_token
116
+ self.output_cls_token = output_cls_token
117
+ self.patch_norm = patch_norm
118
+ self.final_norm = final_norm
119
+ self.with_cp = with_cp
120
+ self.vocabulary_size = vocabulary_size
121
+ self.num_vocabulary_tokens = vocabulary_size + 1
122
+ self.merge_stage = merge_stage
123
+ self.use_attn = use_attn
124
+ self.modality = modality
125
+
126
+
127
+ class SkySensePlusPlusFusionNeckConfig(PretrainedConfig):
128
+ """Configuration for SkySense++ multi-modal fusion neck (TransformerEncoder).
129
+
130
+ Optional component — not used by default backbone checkpoints.
131
+ Fuses concatenated HR/S2/S1 stage-3 features (2816-dim) via a ViT encoder
132
+ with cls token output (1024-dim).
133
+ """
134
+
135
+ model_type = "skysensepp_fusion_neck"
136
+
137
+ def __init__(
138
+ self,
139
+ input_dims=2816,
140
+ embed_dims=1024,
141
+ num_layers=24,
142
+ num_heads=16,
143
+ mlp_ratio=4,
144
+ qkv_bias=True,
145
+ drop_rate=0.0,
146
+ attn_drop_rate=0.0,
147
+ drop_path_rate=0.3,
148
+ with_cls_token=True,
149
+ output_cls_token=True,
150
+ with_cp=False,
151
+ **kwargs,
152
+ ):
153
+ super().__init__(**kwargs)
154
+ self.input_dims = input_dims
155
+ self.embed_dims = embed_dims
156
+ self.num_layers = num_layers
157
+ self.num_heads = num_heads
158
+ self.mlp_ratio = mlp_ratio
159
+ self.qkv_bias = qkv_bias
160
+ self.drop_rate = drop_rate
161
+ self.attn_drop_rate = attn_drop_rate
162
+ self.drop_path_rate = drop_path_rate
163
+ self.with_cls_token = with_cls_token
164
+ self.output_cls_token = output_cls_token
165
+ self.with_cp = with_cp
skysensepp-fusion-neck/conversion_manifest.json ADDED
@@ -0,0 +1,301 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release.ckpt",
3
+ "modality": "fusion",
4
+ "model_class": "SkySensePlusPlusFusionNeckModel",
5
+ "num_tensors": 291,
6
+ "missing_keys": [],
7
+ "unexpected_keys": [],
8
+ "tensor_names": [
9
+ "cls_token",
10
+ "layers.0.attn.in_proj_bias",
11
+ "layers.0.attn.in_proj_weight",
12
+ "layers.0.attn.out_proj.bias",
13
+ "layers.0.attn.out_proj.weight",
14
+ "layers.0.ffn.layers.0.bias",
15
+ "layers.0.ffn.layers.0.weight",
16
+ "layers.0.ffn.layers.3.bias",
17
+ "layers.0.ffn.layers.3.weight",
18
+ "layers.0.norm1.bias",
19
+ "layers.0.norm1.weight",
20
+ "layers.0.norm2.bias",
21
+ "layers.0.norm2.weight",
22
+ "layers.1.attn.in_proj_bias",
23
+ "layers.1.attn.in_proj_weight",
24
+ "layers.1.attn.out_proj.bias",
25
+ "layers.1.attn.out_proj.weight",
26
+ "layers.1.ffn.layers.0.bias",
27
+ "layers.1.ffn.layers.0.weight",
28
+ "layers.1.ffn.layers.3.bias",
29
+ "layers.1.ffn.layers.3.weight",
30
+ "layers.1.norm1.bias",
31
+ "layers.1.norm1.weight",
32
+ "layers.1.norm2.bias",
33
+ "layers.1.norm2.weight",
34
+ "layers.10.attn.in_proj_bias",
35
+ "layers.10.attn.in_proj_weight",
36
+ "layers.10.attn.out_proj.bias",
37
+ "layers.10.attn.out_proj.weight",
38
+ "layers.10.ffn.layers.0.bias",
39
+ "layers.10.ffn.layers.0.weight",
40
+ "layers.10.ffn.layers.3.bias",
41
+ "layers.10.ffn.layers.3.weight",
42
+ "layers.10.norm1.bias",
43
+ "layers.10.norm1.weight",
44
+ "layers.10.norm2.bias",
45
+ "layers.10.norm2.weight",
46
+ "layers.11.attn.in_proj_bias",
47
+ "layers.11.attn.in_proj_weight",
48
+ "layers.11.attn.out_proj.bias",
49
+ "layers.11.attn.out_proj.weight",
50
+ "layers.11.ffn.layers.0.bias",
51
+ "layers.11.ffn.layers.0.weight",
52
+ "layers.11.ffn.layers.3.bias",
53
+ "layers.11.ffn.layers.3.weight",
54
+ "layers.11.norm1.bias",
55
+ "layers.11.norm1.weight",
56
+ "layers.11.norm2.bias",
57
+ "layers.11.norm2.weight",
58
+ "layers.12.attn.in_proj_bias",
59
+ "layers.12.attn.in_proj_weight",
60
+ "layers.12.attn.out_proj.bias",
61
+ "layers.12.attn.out_proj.weight",
62
+ "layers.12.ffn.layers.0.bias",
63
+ "layers.12.ffn.layers.0.weight",
64
+ "layers.12.ffn.layers.3.bias",
65
+ "layers.12.ffn.layers.3.weight",
66
+ "layers.12.norm1.bias",
67
+ "layers.12.norm1.weight",
68
+ "layers.12.norm2.bias",
69
+ "layers.12.norm2.weight",
70
+ "layers.13.attn.in_proj_bias",
71
+ "layers.13.attn.in_proj_weight",
72
+ "layers.13.attn.out_proj.bias",
73
+ "layers.13.attn.out_proj.weight",
74
+ "layers.13.ffn.layers.0.bias",
75
+ "layers.13.ffn.layers.0.weight",
76
+ "layers.13.ffn.layers.3.bias",
77
+ "layers.13.ffn.layers.3.weight",
78
+ "layers.13.norm1.bias",
79
+ "layers.13.norm1.weight",
80
+ "layers.13.norm2.bias",
81
+ "layers.13.norm2.weight",
82
+ "layers.14.attn.in_proj_bias",
83
+ "layers.14.attn.in_proj_weight",
84
+ "layers.14.attn.out_proj.bias",
85
+ "layers.14.attn.out_proj.weight",
86
+ "layers.14.ffn.layers.0.bias",
87
+ "layers.14.ffn.layers.0.weight",
88
+ "layers.14.ffn.layers.3.bias",
89
+ "layers.14.ffn.layers.3.weight",
90
+ "layers.14.norm1.bias",
91
+ "layers.14.norm1.weight",
92
+ "layers.14.norm2.bias",
93
+ "layers.14.norm2.weight",
94
+ "layers.15.attn.in_proj_bias",
95
+ "layers.15.attn.in_proj_weight",
96
+ "layers.15.attn.out_proj.bias",
97
+ "layers.15.attn.out_proj.weight",
98
+ "layers.15.ffn.layers.0.bias",
99
+ "layers.15.ffn.layers.0.weight",
100
+ "layers.15.ffn.layers.3.bias",
101
+ "layers.15.ffn.layers.3.weight",
102
+ "layers.15.norm1.bias",
103
+ "layers.15.norm1.weight",
104
+ "layers.15.norm2.bias",
105
+ "layers.15.norm2.weight",
106
+ "layers.16.attn.in_proj_bias",
107
+ "layers.16.attn.in_proj_weight",
108
+ "layers.16.attn.out_proj.bias",
109
+ "layers.16.attn.out_proj.weight",
110
+ "layers.16.ffn.layers.0.bias",
111
+ "layers.16.ffn.layers.0.weight",
112
+ "layers.16.ffn.layers.3.bias",
113
+ "layers.16.ffn.layers.3.weight",
114
+ "layers.16.norm1.bias",
115
+ "layers.16.norm1.weight",
116
+ "layers.16.norm2.bias",
117
+ "layers.16.norm2.weight",
118
+ "layers.17.attn.in_proj_bias",
119
+ "layers.17.attn.in_proj_weight",
120
+ "layers.17.attn.out_proj.bias",
121
+ "layers.17.attn.out_proj.weight",
122
+ "layers.17.ffn.layers.0.bias",
123
+ "layers.17.ffn.layers.0.weight",
124
+ "layers.17.ffn.layers.3.bias",
125
+ "layers.17.ffn.layers.3.weight",
126
+ "layers.17.norm1.bias",
127
+ "layers.17.norm1.weight",
128
+ "layers.17.norm2.bias",
129
+ "layers.17.norm2.weight",
130
+ "layers.18.attn.in_proj_bias",
131
+ "layers.18.attn.in_proj_weight",
132
+ "layers.18.attn.out_proj.bias",
133
+ "layers.18.attn.out_proj.weight",
134
+ "layers.18.ffn.layers.0.bias",
135
+ "layers.18.ffn.layers.0.weight",
136
+ "layers.18.ffn.layers.3.bias",
137
+ "layers.18.ffn.layers.3.weight",
138
+ "layers.18.norm1.bias",
139
+ "layers.18.norm1.weight",
140
+ "layers.18.norm2.bias",
141
+ "layers.18.norm2.weight",
142
+ "layers.19.attn.in_proj_bias",
143
+ "layers.19.attn.in_proj_weight",
144
+ "layers.19.attn.out_proj.bias",
145
+ "layers.19.attn.out_proj.weight",
146
+ "layers.19.ffn.layers.0.bias",
147
+ "layers.19.ffn.layers.0.weight",
148
+ "layers.19.ffn.layers.3.bias",
149
+ "layers.19.ffn.layers.3.weight",
150
+ "layers.19.norm1.bias",
151
+ "layers.19.norm1.weight",
152
+ "layers.19.norm2.bias",
153
+ "layers.19.norm2.weight",
154
+ "layers.2.attn.in_proj_bias",
155
+ "layers.2.attn.in_proj_weight",
156
+ "layers.2.attn.out_proj.bias",
157
+ "layers.2.attn.out_proj.weight",
158
+ "layers.2.ffn.layers.0.bias",
159
+ "layers.2.ffn.layers.0.weight",
160
+ "layers.2.ffn.layers.3.bias",
161
+ "layers.2.ffn.layers.3.weight",
162
+ "layers.2.norm1.bias",
163
+ "layers.2.norm1.weight",
164
+ "layers.2.norm2.bias",
165
+ "layers.2.norm2.weight",
166
+ "layers.20.attn.in_proj_bias",
167
+ "layers.20.attn.in_proj_weight",
168
+ "layers.20.attn.out_proj.bias",
169
+ "layers.20.attn.out_proj.weight",
170
+ "layers.20.ffn.layers.0.bias",
171
+ "layers.20.ffn.layers.0.weight",
172
+ "layers.20.ffn.layers.3.bias",
173
+ "layers.20.ffn.layers.3.weight",
174
+ "layers.20.norm1.bias",
175
+ "layers.20.norm1.weight",
176
+ "layers.20.norm2.bias",
177
+ "layers.20.norm2.weight",
178
+ "layers.21.attn.in_proj_bias",
179
+ "layers.21.attn.in_proj_weight",
180
+ "layers.21.attn.out_proj.bias",
181
+ "layers.21.attn.out_proj.weight",
182
+ "layers.21.ffn.layers.0.bias",
183
+ "layers.21.ffn.layers.0.weight",
184
+ "layers.21.ffn.layers.3.bias",
185
+ "layers.21.ffn.layers.3.weight",
186
+ "layers.21.norm1.bias",
187
+ "layers.21.norm1.weight",
188
+ "layers.21.norm2.bias",
189
+ "layers.21.norm2.weight",
190
+ "layers.22.attn.in_proj_bias",
191
+ "layers.22.attn.in_proj_weight",
192
+ "layers.22.attn.out_proj.bias",
193
+ "layers.22.attn.out_proj.weight",
194
+ "layers.22.ffn.layers.0.bias",
195
+ "layers.22.ffn.layers.0.weight",
196
+ "layers.22.ffn.layers.3.bias",
197
+ "layers.22.ffn.layers.3.weight",
198
+ "layers.22.norm1.bias",
199
+ "layers.22.norm1.weight",
200
+ "layers.22.norm2.bias",
201
+ "layers.22.norm2.weight",
202
+ "layers.23.attn.in_proj_bias",
203
+ "layers.23.attn.in_proj_weight",
204
+ "layers.23.attn.out_proj.bias",
205
+ "layers.23.attn.out_proj.weight",
206
+ "layers.23.ffn.layers.0.bias",
207
+ "layers.23.ffn.layers.0.weight",
208
+ "layers.23.ffn.layers.3.bias",
209
+ "layers.23.ffn.layers.3.weight",
210
+ "layers.23.norm1.bias",
211
+ "layers.23.norm1.weight",
212
+ "layers.23.norm2.bias",
213
+ "layers.23.norm2.weight",
214
+ "layers.3.attn.in_proj_bias",
215
+ "layers.3.attn.in_proj_weight",
216
+ "layers.3.attn.out_proj.bias",
217
+ "layers.3.attn.out_proj.weight",
218
+ "layers.3.ffn.layers.0.bias",
219
+ "layers.3.ffn.layers.0.weight",
220
+ "layers.3.ffn.layers.3.bias",
221
+ "layers.3.ffn.layers.3.weight",
222
+ "layers.3.norm1.bias",
223
+ "layers.3.norm1.weight",
224
+ "layers.3.norm2.bias",
225
+ "layers.3.norm2.weight",
226
+ "layers.4.attn.in_proj_bias",
227
+ "layers.4.attn.in_proj_weight",
228
+ "layers.4.attn.out_proj.bias",
229
+ "layers.4.attn.out_proj.weight",
230
+ "layers.4.ffn.layers.0.bias",
231
+ "layers.4.ffn.layers.0.weight",
232
+ "layers.4.ffn.layers.3.bias",
233
+ "layers.4.ffn.layers.3.weight",
234
+ "layers.4.norm1.bias",
235
+ "layers.4.norm1.weight",
236
+ "layers.4.norm2.bias",
237
+ "layers.4.norm2.weight",
238
+ "layers.5.attn.in_proj_bias",
239
+ "layers.5.attn.in_proj_weight",
240
+ "layers.5.attn.out_proj.bias",
241
+ "layers.5.attn.out_proj.weight",
242
+ "layers.5.ffn.layers.0.bias",
243
+ "layers.5.ffn.layers.0.weight",
244
+ "layers.5.ffn.layers.3.bias",
245
+ "layers.5.ffn.layers.3.weight",
246
+ "layers.5.norm1.bias",
247
+ "layers.5.norm1.weight",
248
+ "layers.5.norm2.bias",
249
+ "layers.5.norm2.weight",
250
+ "layers.6.attn.in_proj_bias",
251
+ "layers.6.attn.in_proj_weight",
252
+ "layers.6.attn.out_proj.bias",
253
+ "layers.6.attn.out_proj.weight",
254
+ "layers.6.ffn.layers.0.bias",
255
+ "layers.6.ffn.layers.0.weight",
256
+ "layers.6.ffn.layers.3.bias",
257
+ "layers.6.ffn.layers.3.weight",
258
+ "layers.6.norm1.bias",
259
+ "layers.6.norm1.weight",
260
+ "layers.6.norm2.bias",
261
+ "layers.6.norm2.weight",
262
+ "layers.7.attn.in_proj_bias",
263
+ "layers.7.attn.in_proj_weight",
264
+ "layers.7.attn.out_proj.bias",
265
+ "layers.7.attn.out_proj.weight",
266
+ "layers.7.ffn.layers.0.bias",
267
+ "layers.7.ffn.layers.0.weight",
268
+ "layers.7.ffn.layers.3.bias",
269
+ "layers.7.ffn.layers.3.weight",
270
+ "layers.7.norm1.bias",
271
+ "layers.7.norm1.weight",
272
+ "layers.7.norm2.bias",
273
+ "layers.7.norm2.weight",
274
+ "layers.8.attn.in_proj_bias",
275
+ "layers.8.attn.in_proj_weight",
276
+ "layers.8.attn.out_proj.bias",
277
+ "layers.8.attn.out_proj.weight",
278
+ "layers.8.ffn.layers.0.bias",
279
+ "layers.8.ffn.layers.0.weight",
280
+ "layers.8.ffn.layers.3.bias",
281
+ "layers.8.ffn.layers.3.weight",
282
+ "layers.8.norm1.bias",
283
+ "layers.8.norm1.weight",
284
+ "layers.8.norm2.bias",
285
+ "layers.8.norm2.weight",
286
+ "layers.9.attn.in_proj_bias",
287
+ "layers.9.attn.in_proj_weight",
288
+ "layers.9.attn.out_proj.bias",
289
+ "layers.9.attn.out_proj.weight",
290
+ "layers.9.ffn.layers.0.bias",
291
+ "layers.9.ffn.layers.0.weight",
292
+ "layers.9.ffn.layers.3.bias",
293
+ "layers.9.ffn.layers.3.weight",
294
+ "layers.9.norm1.bias",
295
+ "layers.9.norm1.weight",
296
+ "layers.9.norm2.bias",
297
+ "layers.9.norm2.weight",
298
+ "porj_linear.bias",
299
+ "porj_linear.weight"
300
+ ]
301
+ }
skysensepp-fusion-neck/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:18605dbf0082ad045575b7d115604ab048506e785585e4fbdef7dcf180ac1cb7
3
+ size 1220808632
skysensepp-fusion-neck/modeling_skysensepp_fusion_neck.py ADDED
@@ -0,0 +1,164 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ fusion neck (TransformerEncoder) — optional multi-modal fusion module."""
2
+
3
+ from typing import Optional, Tuple, Union
4
+
5
+ import torch
6
+ import torch.nn as nn
7
+ import torch.utils.checkpoint as cp
8
+ from transformers import PreTrainedModel
9
+ from transformers.modeling_outputs import BaseModelOutputWithPooling
10
+
11
+ from .configuration_skysensepp import SkySensePlusPlusFusionNeckConfig
12
+ from .modeling_utils import DropPath, FFN
13
+
14
+
15
+ class FusionEncoderLayer(nn.Module):
16
+ def __init__(
17
+ self,
18
+ embed_dims: int,
19
+ num_heads: int,
20
+ feedforward_channels: int,
21
+ drop_rate: float = 0.0,
22
+ attn_drop_rate: float = 0.0,
23
+ drop_path_rate: float = 0.0,
24
+ qkv_bias: bool = True,
25
+ with_cp: bool = False,
26
+ ):
27
+ super().__init__()
28
+ self.with_cp = with_cp
29
+ self.norm1 = nn.LayerNorm(embed_dims)
30
+ self.attn = nn.MultiheadAttention(
31
+ embed_dim=embed_dims,
32
+ num_heads=num_heads,
33
+ dropout=attn_drop_rate,
34
+ bias=qkv_bias,
35
+ batch_first=True,
36
+ )
37
+ self.proj_drop = nn.Dropout(drop_rate)
38
+ self.norm2 = nn.LayerNorm(embed_dims)
39
+ self.ffn = FFN(
40
+ embed_dims=embed_dims,
41
+ feedforward_channels=feedforward_channels,
42
+ num_fcs=2,
43
+ ffn_drop=drop_rate,
44
+ drop_path=drop_path_rate,
45
+ act_layer=nn.GELU,
46
+ add_identity=True,
47
+ )
48
+ self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
49
+
50
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
51
+ def _inner_forward(x):
52
+ residual = x
53
+ x_norm = self.norm1(x)
54
+ attn_out, _ = self.attn(x_norm, x_norm, x_norm)
55
+ attn_out = self.proj_drop(attn_out)
56
+ x = residual + self.drop_path(attn_out)
57
+ return self.ffn(self.norm2(x), identity=x)
58
+
59
+ if self.with_cp and x.requires_grad:
60
+ return cp.checkpoint(_inner_forward, x, use_reentrant=False)
61
+ return _inner_forward(x)
62
+
63
+
64
+ class SkySensePlusPlusFusionNeckPreTrainedModel(PreTrainedModel):
65
+ config_class = SkySensePlusPlusFusionNeckConfig
66
+ base_model_prefix = "skysensepp_fusion_neck"
67
+ supports_gradient_checkpointing = True
68
+
69
+ def _init_weights(self, module):
70
+ if isinstance(module, nn.Linear):
71
+ nn.init.trunc_normal_(module.weight, std=0.02)
72
+ if module.bias is not None:
73
+ nn.init.zeros_(module.bias)
74
+ elif isinstance(module, nn.LayerNorm):
75
+ nn.init.ones_(module.weight)
76
+ nn.init.zeros_(module.bias)
77
+
78
+
79
+ class SkySensePlusPlusFusionNeckModel(SkySensePlusPlusFusionNeckPreTrainedModel):
80
+ """Fuses per-location multi-modal tokens into a cls-token representation.
81
+
82
+ Input shape: ``(batch, num_modalities, input_dims)`` — e.g. concatenated
83
+ HR + S2 + S1 stage-3 features with ``input_dims=2816``.
84
+
85
+ Output: cls token embedding ``(batch, embed_dims)`` when
86
+ ``output_cls_token=True`` (default).
87
+ """
88
+
89
+ def __init__(self, config: SkySensePlusPlusFusionNeckConfig):
90
+ super().__init__(config)
91
+
92
+ # Original checkpoint uses the typo `porj_linear`.
93
+ self.porj_linear = nn.Linear(config.input_dims, config.embed_dims)
94
+ self.with_cls_token = config.with_cls_token
95
+ self.output_cls_token = config.output_cls_token
96
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
97
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
98
+
99
+ num_layers = config.num_layers
100
+ if num_layers > 1:
101
+ dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
102
+ else:
103
+ dpr = [0.0]
104
+
105
+ self.layers = nn.ModuleList()
106
+ for i in range(config.num_layers):
107
+ self.layers.append(
108
+ FusionEncoderLayer(
109
+ embed_dims=config.embed_dims,
110
+ num_heads=config.num_heads,
111
+ feedforward_channels=config.mlp_ratio * config.embed_dims,
112
+ attn_drop_rate=config.attn_drop_rate,
113
+ drop_rate=config.drop_rate,
114
+ drop_path_rate=dpr[i],
115
+ qkv_bias=config.qkv_bias,
116
+ with_cp=config.with_cp,
117
+ )
118
+ )
119
+
120
+ self.post_init()
121
+
122
+ def forward(
123
+ self,
124
+ hidden_states: torch.Tensor,
125
+ output_hidden_states: Optional[bool] = None,
126
+ return_dict: Optional[bool] = None,
127
+ ) -> Union[Tuple, BaseModelOutputWithPooling]:
128
+ """Forward pass.
129
+
130
+ Args:
131
+ hidden_states: ``(batch, seq_len, input_dims)`` fused modality tokens.
132
+ """
133
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
134
+
135
+ x = self.porj_linear(hidden_states)
136
+ cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
137
+ x = torch.cat((cls_tokens, x), dim=1)
138
+ if not self.with_cls_token:
139
+ x = x[:, 1:]
140
+
141
+ all_hidden_states = () if output_hidden_states else None
142
+ for layer in self.layers:
143
+ x = layer(x)
144
+ if output_hidden_states:
145
+ all_hidden_states = all_hidden_states + (x,)
146
+
147
+ if self.output_cls_token:
148
+ pooler = x[:, 0]
149
+ last_hidden = pooler.unsqueeze(1)
150
+ elif self.with_cls_token:
151
+ pooler = None
152
+ last_hidden = x[:, 1:]
153
+ else:
154
+ pooler = None
155
+ last_hidden = x
156
+
157
+ if not return_dict:
158
+ return (last_hidden, pooler) if pooler is not None else (last_hidden,)
159
+
160
+ return BaseModelOutputWithPooling(
161
+ last_hidden_state=last_hidden,
162
+ pooler_output=pooler,
163
+ hidden_states=all_hidden_states,
164
+ )
skysensepp-fusion-neck/modeling_utils.py ADDED
@@ -0,0 +1,557 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense: Pure PyTorch + HuggingFace Transformers implementation.
2
+
3
+ Shared utility modules used across SkySense model implementations.
4
+ """
5
+
6
+ import math
7
+ from typing import Optional, Tuple
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+
14
+ def to_2tuple(x):
15
+ """Convert to a 2-tuple."""
16
+ if isinstance(x, (list, tuple)):
17
+ return tuple(x)
18
+ return (x, x)
19
+
20
+
21
+ class DropPath(nn.Module):
22
+ """Drop paths (stochastic depth) per sample.
23
+
24
+ Args:
25
+ drop_prob (float): Probability of dropping a path. Default: 0.0.
26
+ """
27
+
28
+ def __init__(self, drop_prob: float = 0.0):
29
+ super().__init__()
30
+ self.drop_prob = drop_prob
31
+
32
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
33
+ if self.drop_prob == 0.0 or not self.training:
34
+ return x
35
+ keep_prob = 1 - self.drop_prob
36
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1)
37
+ random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
38
+ random_tensor = torch.floor(random_tensor + keep_prob)
39
+ output = x / keep_prob * random_tensor
40
+ return output
41
+
42
+
43
+ class PatchEmbed(nn.Module):
44
+ """Image to Patch Embedding using Conv2d.
45
+
46
+ Args:
47
+ in_channels (int): Number of input channels. Default: 3.
48
+ embed_dims (int): Embedding dimension. Default: 96.
49
+ kernel_size (int): Kernel size of the projection. Default: 4.
50
+ stride (int): Stride of the projection. Default: 4.
51
+ padding (int): Padding of the projection. Default: 0.
52
+ norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
53
+ input_size (int or tuple or None): Input resolution for calculating output size.
54
+ """
55
+
56
+ def __init__(
57
+ self,
58
+ in_channels: int = 3,
59
+ embed_dims: int = 96,
60
+ kernel_size: int = 4,
61
+ stride: int = 4,
62
+ padding: int = 0,
63
+ norm_layer: Optional[type] = nn.LayerNorm,
64
+ input_size: Optional[int] = None,
65
+ ):
66
+ super().__init__()
67
+ self.projection = nn.Conv2d(
68
+ in_channels, embed_dims,
69
+ kernel_size=kernel_size, stride=stride, padding=padding,
70
+ )
71
+ self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
72
+
73
+ # Compute init output size if input_size is given
74
+ if input_size is not None:
75
+ input_size = to_2tuple(input_size)
76
+ self.init_out_size = (
77
+ (input_size[0] - kernel_size + 2 * padding) // stride + 1,
78
+ (input_size[1] - kernel_size + 2 * padding) // stride + 1,
79
+ )
80
+ else:
81
+ self.init_out_size = None
82
+
83
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
84
+ x = self.projection(x) # (B, C, H, W)
85
+ out_size = (x.shape[2], x.shape[3])
86
+ x = x.flatten(2).transpose(1, 2) # (B, H*W, C)
87
+ x = self.norm(x)
88
+ return x, out_size
89
+
90
+
91
+ class FFN(nn.Module):
92
+ """Feed-Forward Network.
93
+
94
+ Args:
95
+ embed_dims (int): Input dimension.
96
+ feedforward_channels (int): Hidden dimension.
97
+ num_fcs (int): Number of FC layers. Default: 2.
98
+ ffn_drop (float): Dropout rate. Default: 0.0.
99
+ drop_path (float): Drop path rate. Default: 0.0.
100
+ act_layer (nn.Module): Activation layer class. Default: nn.GELU.
101
+ add_identity (bool): Whether to add identity connection. Default: True.
102
+ """
103
+
104
+ def __init__(
105
+ self,
106
+ embed_dims: int,
107
+ feedforward_channels: int,
108
+ num_fcs: int = 2,
109
+ ffn_drop: float = 0.0,
110
+ drop_path: float = 0.0,
111
+ act_layer: type = nn.GELU,
112
+ add_identity: bool = True,
113
+ ):
114
+ super().__init__()
115
+ assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
116
+ self.embed_dims = embed_dims
117
+ self.feedforward_channels = feedforward_channels
118
+ self.add_identity = add_identity
119
+
120
+ layers = []
121
+ in_channels = embed_dims
122
+ for i in range(num_fcs - 1):
123
+ layers.append(nn.Linear(in_channels, feedforward_channels))
124
+ layers.append(act_layer())
125
+ layers.append(nn.Dropout(ffn_drop))
126
+ in_channels = feedforward_channels
127
+ layers.append(nn.Linear(feedforward_channels, embed_dims))
128
+ layers.append(nn.Dropout(ffn_drop))
129
+ self.layers = nn.Sequential(*layers)
130
+
131
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
132
+
133
+ def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
134
+ out = self.layers(x)
135
+ out = self.drop_path(out)
136
+ if self.add_identity:
137
+ if identity is None:
138
+ identity = x
139
+ out = out + identity
140
+ return out
141
+
142
+
143
+ class WindowMSAV2(nn.Module):
144
+ """Window-based Multi-head Self-Attention for Swin Transformer V2.
145
+
146
+ Uses cosine attention and log-spaced continuous position bias (log-CPB).
147
+
148
+ Args:
149
+ embed_dims (int): Number of input channels.
150
+ num_heads (int): Number of attention heads.
151
+ window_size (tuple[int]): Window size (Wh, Ww).
152
+ pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
153
+ qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
154
+ attn_drop (float): Attention dropout rate. Default: 0.0.
155
+ proj_drop (float): Output projection dropout rate. Default: 0.0.
156
+ """
157
+
158
+ def __init__(
159
+ self,
160
+ embed_dims: int,
161
+ num_heads: int,
162
+ window_size: Tuple[int, int],
163
+ pretrained_window_size: Tuple[int, int] = (0, 0),
164
+ qkv_bias: bool = True,
165
+ attn_drop: float = 0.0,
166
+ proj_drop: float = 0.0,
167
+ ):
168
+ super().__init__()
169
+ self.embed_dims = embed_dims
170
+ self.num_heads = num_heads
171
+ self.window_size = window_size
172
+ self.pretrained_window_size = pretrained_window_size
173
+
174
+ self.logit_scale = nn.Parameter(
175
+ torch.log(10 * torch.ones((num_heads, 1, 1))))
176
+
177
+ # MLP for continuous relative position bias (log-CPB)
178
+ self.cpb_mlp = nn.Sequential(
179
+ nn.Linear(2, 512, bias=True),
180
+ nn.ReLU(inplace=True),
181
+ nn.Linear(512, num_heads, bias=False),
182
+ )
183
+
184
+ # Build relative coords table
185
+ self._build_relative_coords_table()
186
+ # Build relative position index
187
+ self._build_relative_position_index()
188
+
189
+ self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
190
+ if qkv_bias:
191
+ self.q_bias = nn.Parameter(torch.zeros(embed_dims))
192
+ self.v_bias = nn.Parameter(torch.zeros(embed_dims))
193
+ else:
194
+ self.q_bias = None
195
+ self.v_bias = None
196
+
197
+ self.attn_drop = nn.Dropout(attn_drop)
198
+ self.proj = nn.Linear(embed_dims, embed_dims)
199
+ self.proj_drop = nn.Dropout(proj_drop)
200
+ self.softmax = nn.Softmax(dim=-1)
201
+
202
+ def _build_relative_coords_table(self):
203
+ """Build the relative coordinates table for log-CPB."""
204
+ Wh, Ww = self.window_size
205
+ # Table of relative coordinates
206
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
207
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
208
+ coords_table = torch.stack(
209
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
210
+ ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2)
211
+
212
+ # Normalize to [-1, 1] and apply log-scale
213
+ if self.pretrained_window_size[0] > 0:
214
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
215
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
216
+ else:
217
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
218
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
219
+ coords_table *= 8 # normalize to -8, 8
220
+ coords_table = (
221
+ torch.sign(coords_table)
222
+ * torch.log2(torch.abs(coords_table) + 1.0)
223
+ / math.log2(8)
224
+ )
225
+ self.register_buffer("relative_coords_table", coords_table)
226
+
227
+ def _build_relative_position_index(self):
228
+ """Build the pairwise relative position index for each window token."""
229
+ Wh, Ww = self.window_size
230
+ coords_h = torch.arange(Wh)
231
+ coords_w = torch.arange(Ww)
232
+ coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
233
+ coords_flatten = coords.view(2, -1)
234
+
235
+ relative_coords = (
236
+ coords_flatten[:, :, None] - coords_flatten[:, None, :]
237
+ ) # (2, Wh*Ww, Wh*Ww)
238
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
239
+ relative_coords[:, :, 0] += Wh - 1
240
+ relative_coords[:, :, 1] += Ww - 1
241
+ relative_coords[:, :, 0] *= 2 * Ww - 1
242
+ relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww)
243
+ self.register_buffer("relative_position_index", relative_position_index)
244
+
245
+ def _compute_position_bias(self, N):
246
+ """Compute relative position bias, supporting dynamic window sizes.
247
+
248
+ The log-CPB (Continuous Position Bias) MLP can generalize to any window
249
+ size by computing bias from normalized relative coordinates.
250
+ """
251
+ init_N = self.window_size[0] * self.window_size[1]
252
+ if N == init_N:
253
+ # Use pre-built tables
254
+ relative_position_bias_table = self.cpb_mlp(
255
+ self.relative_coords_table
256
+ ).view(-1, self.num_heads)
257
+ relative_position_bias = relative_position_bias_table[
258
+ self.relative_position_index.view(-1)
259
+ ].view(N, N, -1)
260
+ else:
261
+ # Dynamic: compute for actual window size on-the-fly
262
+ Wh = Ww = int(math.sqrt(N))
263
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
264
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
265
+ coords_table = torch.stack(
266
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
267
+ ).flatten(1).transpose(0, 1).unsqueeze(0)
268
+ if self.pretrained_window_size[0] > 0:
269
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
270
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
271
+ else:
272
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
273
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
274
+ coords_table *= 8
275
+ coords_table = (
276
+ torch.sign(coords_table)
277
+ * torch.log2(torch.abs(coords_table) + 1.0)
278
+ / math.log2(8)
279
+ )
280
+ # Build position index for actual window size
281
+ ch = torch.arange(Wh, device=self.logit_scale.device)
282
+ cw = torch.arange(Ww, device=self.logit_scale.device)
283
+ coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
284
+ coords_flat = coords.view(2, -1)
285
+ rel = coords_flat[:, :, None] - coords_flat[:, None, :]
286
+ rel = rel.permute(1, 2, 0).contiguous()
287
+ rel[:, :, 0] += Wh - 1
288
+ rel[:, :, 1] += Ww - 1
289
+ rel[:, :, 0] *= 2 * Ww - 1
290
+ pos_index = rel.sum(-1)
291
+
292
+ bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
293
+ relative_position_bias = bias_table[
294
+ pos_index.view(-1)
295
+ ].view(N, N, -1)
296
+
297
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
298
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
299
+ return relative_position_bias
300
+
301
+ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
302
+ """
303
+ Args:
304
+ x: (num_windows*B, N, C) where N = Wh*Ww
305
+ mask: (num_windows, N, N) or None
306
+ """
307
+ B_, N, C = x.shape
308
+
309
+ # Compute QKV with bias
310
+ if self.q_bias is not None:
311
+ qkv_bias = torch.cat(
312
+ (self.q_bias,
313
+ torch.zeros_like(self.v_bias, requires_grad=False),
314
+ self.v_bias))
315
+ qkv = F.linear(x, self.qkv.weight, qkv_bias)
316
+ else:
317
+ qkv = self.qkv(x)
318
+
319
+ qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
320
+ qkv = qkv.permute(2, 0, 3, 1, 4)
321
+ q, k, v = qkv.unbind(0)
322
+
323
+ # Cosine attention
324
+ attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
325
+ logit_scale = torch.clamp(
326
+ self.logit_scale, max=math.log(1.0 / 0.01)
327
+ ).exp()
328
+ attn = attn * logit_scale
329
+
330
+ # Log-CPB relative position bias (supports dynamic window sizes)
331
+ relative_position_bias = self._compute_position_bias(N)
332
+ attn = attn + relative_position_bias.unsqueeze(0)
333
+
334
+ if mask is not None:
335
+ nW = mask.shape[0]
336
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
337
+ attn = attn + mask.unsqueeze(1).unsqueeze(0)
338
+ attn = attn.view(-1, self.num_heads, N, N)
339
+
340
+ attn = self.softmax(attn)
341
+ attn = self.attn_drop(attn)
342
+
343
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
344
+ x = self.proj(x)
345
+ x = self.proj_drop(x)
346
+ return x
347
+
348
+
349
+ class ShiftWindowMSA(nn.Module):
350
+ """Shifted Window Multi-head Self-Attention.
351
+
352
+ Args:
353
+ embed_dims (int): Number of input channels.
354
+ num_heads (int): Number of attention heads.
355
+ window_size (int): Window size.
356
+ shift_size (int): Shift size for SW-MSA. Default: 0.
357
+ attn_drop (float): Attention dropout rate. Default: 0.0.
358
+ proj_drop (float): Projection dropout rate. Default: 0.0.
359
+ drop_path (float): Drop path rate. Default: 0.0.
360
+ pad_small_map (bool): Pad small feature maps to window size. Default: False.
361
+ pretrained_window_size (int): Pretrained window size. Default: 0.
362
+ """
363
+
364
+ def __init__(
365
+ self,
366
+ embed_dims: int,
367
+ num_heads: int,
368
+ window_size: int,
369
+ shift_size: int = 0,
370
+ attn_drop: float = 0.0,
371
+ proj_drop: float = 0.0,
372
+ drop_path: float = 0.0,
373
+ pad_small_map: bool = False,
374
+ pretrained_window_size: int = 0,
375
+ ):
376
+ super().__init__()
377
+ self.window_size = window_size
378
+ self.shift_size = shift_size
379
+ self.pad_small_map = pad_small_map
380
+
381
+ self.w_msa = WindowMSAV2(
382
+ embed_dims=embed_dims,
383
+ num_heads=num_heads,
384
+ window_size=to_2tuple(window_size),
385
+ pretrained_window_size=to_2tuple(pretrained_window_size),
386
+ attn_drop=attn_drop,
387
+ proj_drop=proj_drop,
388
+ )
389
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
390
+
391
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
392
+ B, L, C = x.shape
393
+ H, W = hw_shape
394
+ assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
395
+
396
+ x = x.view(B, H, W, C)
397
+
398
+ window_size = self.window_size
399
+ shift_size = self.shift_size
400
+
401
+ # Pad or shrink window
402
+ if self.pad_small_map:
403
+ pad_r = (window_size - W % window_size) % window_size
404
+ pad_b = (window_size - H % window_size) % window_size
405
+ x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
406
+ _, Hp, Wp, _ = x.shape
407
+ else:
408
+ Hp, Wp = H, W
409
+ if window_size > Hp:
410
+ window_size = Hp
411
+ shift_size = 0
412
+ if window_size > Wp:
413
+ window_size = Wp
414
+ shift_size = 0
415
+
416
+ # Compute attention mask for SW-MSA
417
+ attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
418
+
419
+ # Cyclic shift
420
+ if shift_size > 0:
421
+ x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
422
+
423
+ # Partition windows
424
+ x_windows = self._window_partition(x, window_size)
425
+ # (num_windows*B, window_size*window_size, C)
426
+
427
+ # W-MSA/SW-MSA
428
+ attn_windows = self.w_msa(x_windows, mask=attn_mask)
429
+
430
+ # Merge windows
431
+ x = self._window_reverse(attn_windows, window_size, Hp, Wp)
432
+
433
+ # Reverse cyclic shift
434
+ if shift_size > 0:
435
+ x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
436
+
437
+ if self.pad_small_map and (pad_r > 0 or pad_b > 0):
438
+ x = x[:, :H, :W, :].contiguous()
439
+
440
+ x = x.view(B, H * W, C)
441
+ x = self.drop_path(x)
442
+ return x
443
+
444
+ @staticmethod
445
+ def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
446
+ """Partition into non-overlapping windows."""
447
+ B, H, W, C = x.shape
448
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
449
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
450
+ windows = windows.view(-1, window_size * window_size, C)
451
+ return windows
452
+
453
+ @staticmethod
454
+ def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
455
+ """Reverse window partition."""
456
+ B_nW = windows.shape[0]
457
+ nH = H // window_size
458
+ nW = W // window_size
459
+ B = B_nW // (nH * nW)
460
+ x = windows.view(B, nH, nW, window_size, window_size, -1)
461
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
462
+ x = x.view(B, H, W, -1)
463
+ return x
464
+
465
+ @staticmethod
466
+ def _compute_attn_mask(H, W, window_size, shift_size, device):
467
+ """Compute attention mask for shifted window attention."""
468
+ if shift_size <= 0:
469
+ return None
470
+ img_mask = torch.zeros((1, H, W, 1), device=device)
471
+ h_slices = (
472
+ slice(0, -window_size),
473
+ slice(-window_size, -shift_size),
474
+ slice(-shift_size, None),
475
+ )
476
+ w_slices = (
477
+ slice(0, -window_size),
478
+ slice(-window_size, -shift_size),
479
+ slice(-shift_size, None),
480
+ )
481
+ cnt = 0
482
+ for h in h_slices:
483
+ for w in w_slices:
484
+ img_mask[:, h, w, :] = cnt
485
+ cnt += 1
486
+
487
+ # Partition mask
488
+ mask_windows = img_mask.view(
489
+ 1, H // window_size, window_size, W // window_size, window_size, 1
490
+ )
491
+ mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
492
+ mask_windows = mask_windows.view(-1, window_size * window_size)
493
+
494
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
495
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
496
+ attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
497
+ return attn_mask
498
+
499
+
500
+ class PatchMerging(nn.Module):
501
+ """Patch Merging Layer for downsampling (2x).
502
+
503
+ Args:
504
+ in_channels (int): Input channels.
505
+ out_channels (int): Output channels.
506
+ norm_layer (type): Normalization layer. Default: nn.LayerNorm.
507
+ is_post_norm (bool): Apply norm after linear. Default: True.
508
+ """
509
+
510
+ def __init__(
511
+ self,
512
+ in_channels: int,
513
+ out_channels: int,
514
+ norm_layer: type = nn.LayerNorm,
515
+ is_post_norm: bool = True,
516
+ ):
517
+ super().__init__()
518
+ self.in_channels = in_channels
519
+ self.out_channels = out_channels
520
+ self.is_post_norm = is_post_norm
521
+ self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
522
+ if is_post_norm:
523
+ self.norm = norm_layer(out_channels)
524
+ else:
525
+ self.norm = norm_layer(4 * in_channels)
526
+
527
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
528
+ B, L, C = x.shape
529
+ H, W = hw_shape
530
+ assert L == H * W
531
+
532
+ x = x.view(B, H, W, C)
533
+
534
+ # Pad if needed
535
+ pad_h = H % 2
536
+ pad_w = W % 2
537
+ if pad_h or pad_w:
538
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
539
+
540
+ x0 = x[:, 0::2, 0::2, :]
541
+ x1 = x[:, 1::2, 0::2, :]
542
+ x2 = x[:, 0::2, 1::2, :]
543
+ x3 = x[:, 1::2, 1::2, :]
544
+ x = torch.cat([x0, x1, x2, x3], dim=-1)
545
+
546
+ out_h = (H + pad_h) // 2
547
+ out_w = (W + pad_w) // 2
548
+ x = x.view(B, out_h * out_w, 4 * C)
549
+
550
+ if self.is_post_norm:
551
+ x = self.reduction(x)
552
+ x = self.norm(x)
553
+ else:
554
+ x = self.norm(x)
555
+ x = self.reduction(x)
556
+
557
+ return x, (out_h, out_w)
skysensepp-fusion-neck/pipeline_skysensepp.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Custom HuggingFace pipeline for SkySense++ MSL feature extraction."""
2
+
3
+ from typing import Any, Dict, Optional, Union
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline):
11
+ """Pipeline for SkySense++ MSL backbones.
12
+
13
+ Expects image tensors plus semantic annotation maps (class indices).
14
+ """
15
+
16
+ def _sanitize_parameters(
17
+ self,
18
+ annotation=None,
19
+ mask=None,
20
+ output_hidden_states=None,
21
+ **kwargs,
22
+ ):
23
+ preprocess_params = {}
24
+ forward_params = {}
25
+ postprocess_params = {}
26
+
27
+ if annotation is not None:
28
+ preprocess_params["annotation"] = annotation
29
+ if mask is not None:
30
+ forward_params["mask"] = mask
31
+ if output_hidden_states is not None:
32
+ forward_params["output_hidden_states"] = output_hidden_states
33
+
34
+ return preprocess_params, forward_params, postprocess_params
35
+
36
+ def preprocess(
37
+ self,
38
+ pixel_values: Any,
39
+ annotation: Optional[Any] = None,
40
+ **kwargs,
41
+ ) -> Dict[str, torch.Tensor]:
42
+ if isinstance(pixel_values, dict):
43
+ annotation = pixel_values.get("annotation", annotation)
44
+ pixel_values = pixel_values.get("pixel_values", pixel_values)
45
+
46
+ if isinstance(pixel_values, np.ndarray):
47
+ pixel_values = torch.from_numpy(pixel_values).float()
48
+ elif not isinstance(pixel_values, torch.Tensor):
49
+ raise TypeError(
50
+ f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}"
51
+ )
52
+
53
+ if annotation is None:
54
+ raise ValueError("SkySense++ MSL models require an `annotation` semantic map.")
55
+
56
+ if isinstance(annotation, np.ndarray):
57
+ annotation = torch.from_numpy(annotation).long()
58
+ elif not isinstance(annotation, torch.Tensor):
59
+ raise TypeError(
60
+ f"Expected tensor or ndarray for annotation, got {type(annotation)}"
61
+ )
62
+
63
+ if pixel_values.ndim == 3:
64
+ pixel_values = pixel_values.unsqueeze(0)
65
+ if annotation.ndim == 2:
66
+ annotation = annotation.unsqueeze(0)
67
+
68
+ return {"pixel_values": pixel_values, "annotation": annotation}
69
+
70
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
71
+ with torch.no_grad():
72
+ outputs = self.model(
73
+ pixel_values=model_inputs["pixel_values"],
74
+ annotation=model_inputs["annotation"],
75
+ mask=kwargs.get("mask"),
76
+ output_hidden_states=kwargs.get("output_hidden_states", False),
77
+ return_dict=True,
78
+ )
79
+ return {"outputs": outputs}
80
+
81
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
82
+ outputs = model_outputs["outputs"]
83
+ result = {"last_hidden_state": outputs.last_hidden_state}
84
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
85
+ result["hidden_states"] = outputs.hidden_states
86
+ return result
skysensepp-fusion-neck/pipeline_skysensepp_fusion.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Optional pipeline for SkySense++ fusion neck."""
2
+
3
+ from typing import Any, Dict
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusFusionNeckPipeline(Pipeline):
11
+ """Pipeline for the optional SkySense++ fusion neck module.
12
+
13
+ Expects concatenated multi-modal tokens per spatial location:
14
+ ``(batch, num_modalities, input_dims)``.
15
+ """
16
+
17
+ def _sanitize_parameters(self, output_hidden_states=None, **kwargs):
18
+ preprocess_params = {}
19
+ forward_params = {}
20
+ postprocess_params = {}
21
+ if output_hidden_states is not None:
22
+ forward_params["output_hidden_states"] = output_hidden_states
23
+ return preprocess_params, forward_params, postprocess_params
24
+
25
+ def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]:
26
+ if isinstance(hidden_states, dict):
27
+ hidden_states = hidden_states["hidden_states"]
28
+
29
+ if isinstance(hidden_states, np.ndarray):
30
+ hidden_states = torch.from_numpy(hidden_states).float()
31
+ elif not isinstance(hidden_states, torch.Tensor):
32
+ raise TypeError(
33
+ f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}"
34
+ )
35
+ if hidden_states.ndim == 2:
36
+ hidden_states = hidden_states.unsqueeze(0)
37
+ return {"hidden_states": hidden_states}
38
+
39
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
40
+ with torch.no_grad():
41
+ outputs = self.model(
42
+ hidden_states=model_inputs["hidden_states"],
43
+ output_hidden_states=kwargs.get("output_hidden_states", False),
44
+ return_dict=True,
45
+ )
46
+ return {"outputs": outputs}
47
+
48
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
49
+ outputs = model_outputs["outputs"]
50
+ result = {"pooler_output": outputs.pooler_output}
51
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
52
+ result["hidden_states"] = outputs.hidden_states
53
+ return result
skysensepp-swinv2-msl-hr/__init__.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++: Multi-Modal Remote Sensing Foundation Model (HuggingFace)."""
2
+
3
+ from .configuration_skysensepp import (
4
+ SkySensePlusPlusSwinV2MSLConfig,
5
+ SkySensePlusPlusViTMSLConfig,
6
+ )
7
+ from .modeling_skysensepp_swinv2_msl import (
8
+ SkySensePlusPlusSwinV2MSLModel,
9
+ SkySensePlusPlusSwinV2MSLPreTrainedModel,
10
+ )
11
+ from .modeling_skysensepp_vit_msl import (
12
+ SkySensePlusPlusViTMSLModel,
13
+ SkySensePlusPlusViTMSLPreTrainedModel,
14
+ )
15
+ from .pipeline_skysensepp import SkySensePlusPlusMSLFeatureExtractionPipeline
16
+
17
+ __all__ = [
18
+ "SkySensePlusPlusSwinV2MSLConfig",
19
+ "SkySensePlusPlusViTMSLConfig",
20
+ "SkySensePlusPlusSwinV2MSLModel",
21
+ "SkySensePlusPlusSwinV2MSLPreTrainedModel",
22
+ "SkySensePlusPlusViTMSLModel",
23
+ "SkySensePlusPlusViTMSLPreTrainedModel",
24
+ "SkySensePlusPlusMSLFeatureExtractionPipeline",
25
+ ]
skysensepp-swinv2-msl-hr/config.json ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "return_dict": true,
3
+ "output_hidden_states": false,
4
+ "dtype": "float32",
5
+ "chunk_size_feed_forward": 0,
6
+ "is_encoder_decoder": false,
7
+ "architectures": [
8
+ "SkySensePlusPlusSwinV2MSLModel"
9
+ ],
10
+ "id2label": {
11
+ "0": "LABEL_0",
12
+ "1": "LABEL_1"
13
+ },
14
+ "label2id": {
15
+ "LABEL_0": 0,
16
+ "LABEL_1": 1
17
+ },
18
+ "problem_type": null,
19
+ "_name_or_path": "",
20
+ "transformers_version": "5.0.0",
21
+ "arch": "huge",
22
+ "embed_dims": 352,
23
+ "depths": [
24
+ 2,
25
+ 2,
26
+ 18,
27
+ 2
28
+ ],
29
+ "num_heads": [
30
+ 8,
31
+ 16,
32
+ 32,
33
+ 64
34
+ ],
35
+ "extra_norm_every_n_blocks": 6,
36
+ "img_size": 512,
37
+ "patch_size": 4,
38
+ "in_channels": 3,
39
+ "window_size": 8,
40
+ "drop_rate": 0.0,
41
+ "drop_path_rate": 0.2,
42
+ "out_indices": [
43
+ 0,
44
+ 1,
45
+ 2,
46
+ 3
47
+ ],
48
+ "use_abs_pos_embed": false,
49
+ "with_cp": false,
50
+ "pad_small_map": false,
51
+ "pretrained_window_sizes": [
52
+ 0,
53
+ 0,
54
+ 0,
55
+ 0
56
+ ],
57
+ "is_post_norm_downsample": true,
58
+ "vocabulary_size": 64,
59
+ "num_vocabulary_tokens": 65,
60
+ "merge_stage": 2,
61
+ "use_attn": true,
62
+ "model_type": "skysensepp_swinv2_msl",
63
+ "output_attentions": false,
64
+ "auto_map": {
65
+ "AutoConfig": "configuration_skysensepp.SkySensePlusPlusSwinV2MSLConfig",
66
+ "AutoModel": "modeling_skysensepp_swinv2_msl.SkySensePlusPlusSwinV2MSLModel"
67
+ },
68
+ "custom_pipelines": {
69
+ "skysensepp-feature-extraction": {
70
+ "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline",
71
+ "pt": [
72
+ "AutoModel"
73
+ ]
74
+ },
75
+ "image-feature-extraction": {
76
+ "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline",
77
+ "pt": [
78
+ "AutoModel"
79
+ ]
80
+ }
81
+ }
82
+ }
skysensepp-swinv2-msl-hr/configuration_skysensepp.py ADDED
@@ -0,0 +1,124 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration classes for SkySense++ MSL backbones."""
2
+
3
+ from transformers import PretrainedConfig
4
+
5
+
6
+ class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig):
7
+ """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical)."""
8
+
9
+ model_type = "skysensepp_swinv2_msl"
10
+
11
+ arch_zoo = {
12
+ "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
13
+ "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
14
+ "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0},
15
+ "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0},
16
+ "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6},
17
+ "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6},
18
+ }
19
+
20
+ def __init__(
21
+ self,
22
+ arch="huge",
23
+ img_size=224,
24
+ patch_size=4,
25
+ in_channels=3,
26
+ window_size=8,
27
+ drop_rate=0.0,
28
+ drop_path_rate=0.2,
29
+ out_indices=(0, 1, 2, 3),
30
+ use_abs_pos_embed=False,
31
+ with_cp=False,
32
+ pad_small_map=False,
33
+ pretrained_window_sizes=(0, 0, 0, 0),
34
+ is_post_norm_downsample=True,
35
+ vocabulary_size=64,
36
+ merge_stage=2,
37
+ use_attn=True,
38
+ **kwargs,
39
+ ):
40
+ super().__init__(**kwargs)
41
+
42
+ arch = arch.lower()
43
+ if arch not in self.arch_zoo:
44
+ raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
45
+ arch_settings = self.arch_zoo[arch]
46
+
47
+ self.arch = arch
48
+ self.embed_dims = arch_settings["embed_dims"]
49
+ self.depths = arch_settings["depths"]
50
+ self.num_heads = arch_settings["num_heads"]
51
+ self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"]
52
+
53
+ self.img_size = img_size
54
+ self.patch_size = patch_size
55
+ self.in_channels = in_channels
56
+ self.window_size = window_size
57
+ self.drop_rate = drop_rate
58
+ self.drop_path_rate = drop_path_rate
59
+ self.out_indices = list(out_indices)
60
+ self.use_abs_pos_embed = use_abs_pos_embed
61
+ self.with_cp = with_cp
62
+ self.pad_small_map = pad_small_map
63
+ self.pretrained_window_sizes = list(pretrained_window_sizes)
64
+ self.is_post_norm_downsample = is_post_norm_downsample
65
+
66
+ self.vocabulary_size = vocabulary_size
67
+ self.num_vocabulary_tokens = vocabulary_size + 1
68
+ self.merge_stage = merge_stage
69
+ self.use_attn = use_attn
70
+
71
+
72
+ class SkySensePlusPlusViTMSLConfig(PretrainedConfig):
73
+ """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1)."""
74
+
75
+ model_type = "skysensepp_vit_msl"
76
+
77
+ def __init__(
78
+ self,
79
+ img_size=16,
80
+ patch_size=4,
81
+ in_channels=10,
82
+ embed_dims=1024,
83
+ num_layers=24,
84
+ num_heads=16,
85
+ mlp_ratio=4,
86
+ out_indices=(5, 11, 17, 23),
87
+ qkv_bias=True,
88
+ drop_rate=0.0,
89
+ attn_drop_rate=0.0,
90
+ drop_path_rate=0.3,
91
+ with_cls_token=False,
92
+ output_cls_token=False,
93
+ patch_norm=False,
94
+ final_norm=False,
95
+ with_cp=False,
96
+ vocabulary_size=64,
97
+ merge_stage=4,
98
+ use_attn=False,
99
+ modality="s2",
100
+ **kwargs,
101
+ ):
102
+ super().__init__(**kwargs)
103
+ self.img_size = img_size
104
+ self.patch_size = patch_size
105
+ self.in_channels = in_channels
106
+ self.embed_dims = embed_dims
107
+ self.num_layers = num_layers
108
+ self.num_heads = num_heads
109
+ self.mlp_ratio = mlp_ratio
110
+ self.out_indices = list(out_indices)
111
+ self.qkv_bias = qkv_bias
112
+ self.drop_rate = drop_rate
113
+ self.attn_drop_rate = attn_drop_rate
114
+ self.drop_path_rate = drop_path_rate
115
+ self.with_cls_token = with_cls_token
116
+ self.output_cls_token = output_cls_token
117
+ self.patch_norm = patch_norm
118
+ self.final_norm = final_norm
119
+ self.with_cp = with_cp
120
+ self.vocabulary_size = vocabulary_size
121
+ self.num_vocabulary_tokens = vocabulary_size + 1
122
+ self.merge_stage = merge_stage
123
+ self.use_attn = use_attn
124
+ self.modality = modality
skysensepp-swinv2-msl-hr/conversion_manifest.json ADDED
@@ -0,0 +1,523 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release_hr.pth",
3
+ "modality": "hr",
4
+ "model_class": "SkySensePlusPlusSwinV2MSLModel",
5
+ "num_tensors": 464,
6
+ "missing_keys": [
7
+ "stages.0.blocks.0.attn.w_msa.relative_coords_table",
8
+ "stages.0.blocks.0.attn.w_msa.relative_position_index",
9
+ "stages.0.blocks.1.attn.w_msa.relative_coords_table",
10
+ "stages.0.blocks.1.attn.w_msa.relative_position_index",
11
+ "stages.1.blocks.0.attn.w_msa.relative_coords_table",
12
+ "stages.1.blocks.0.attn.w_msa.relative_position_index",
13
+ "stages.1.blocks.1.attn.w_msa.relative_coords_table",
14
+ "stages.1.blocks.1.attn.w_msa.relative_position_index",
15
+ "stages.2.blocks.0.attn.w_msa.relative_coords_table",
16
+ "stages.2.blocks.0.attn.w_msa.relative_position_index",
17
+ "stages.2.blocks.1.attn.w_msa.relative_coords_table",
18
+ "stages.2.blocks.1.attn.w_msa.relative_position_index",
19
+ "stages.2.blocks.2.attn.w_msa.relative_coords_table",
20
+ "stages.2.blocks.2.attn.w_msa.relative_position_index",
21
+ "stages.2.blocks.3.attn.w_msa.relative_coords_table",
22
+ "stages.2.blocks.3.attn.w_msa.relative_position_index",
23
+ "stages.2.blocks.4.attn.w_msa.relative_coords_table",
24
+ "stages.2.blocks.4.attn.w_msa.relative_position_index",
25
+ "stages.2.blocks.5.attn.w_msa.relative_coords_table",
26
+ "stages.2.blocks.5.attn.w_msa.relative_position_index",
27
+ "stages.2.blocks.6.attn.w_msa.relative_coords_table",
28
+ "stages.2.blocks.6.attn.w_msa.relative_position_index",
29
+ "stages.2.blocks.7.attn.w_msa.relative_coords_table",
30
+ "stages.2.blocks.7.attn.w_msa.relative_position_index",
31
+ "stages.2.blocks.8.attn.w_msa.relative_coords_table",
32
+ "stages.2.blocks.8.attn.w_msa.relative_position_index",
33
+ "stages.2.blocks.9.attn.w_msa.relative_coords_table",
34
+ "stages.2.blocks.9.attn.w_msa.relative_position_index",
35
+ "stages.2.blocks.10.attn.w_msa.relative_coords_table",
36
+ "stages.2.blocks.10.attn.w_msa.relative_position_index",
37
+ "stages.2.blocks.11.attn.w_msa.relative_coords_table",
38
+ "stages.2.blocks.11.attn.w_msa.relative_position_index",
39
+ "stages.2.blocks.12.attn.w_msa.relative_coords_table",
40
+ "stages.2.blocks.12.attn.w_msa.relative_position_index",
41
+ "stages.2.blocks.13.attn.w_msa.relative_coords_table",
42
+ "stages.2.blocks.13.attn.w_msa.relative_position_index",
43
+ "stages.2.blocks.14.attn.w_msa.relative_coords_table",
44
+ "stages.2.blocks.14.attn.w_msa.relative_position_index",
45
+ "stages.2.blocks.15.attn.w_msa.relative_coords_table",
46
+ "stages.2.blocks.15.attn.w_msa.relative_position_index",
47
+ "stages.2.blocks.16.attn.w_msa.relative_coords_table",
48
+ "stages.2.blocks.16.attn.w_msa.relative_position_index",
49
+ "stages.2.blocks.17.attn.w_msa.relative_coords_table",
50
+ "stages.2.blocks.17.attn.w_msa.relative_position_index",
51
+ "stages.3.blocks.0.attn.w_msa.relative_coords_table",
52
+ "stages.3.blocks.0.attn.w_msa.relative_position_index",
53
+ "stages.3.blocks.1.attn.w_msa.relative_coords_table",
54
+ "stages.3.blocks.1.attn.w_msa.relative_position_index"
55
+ ],
56
+ "unexpected_keys": [],
57
+ "tensor_names": [
58
+ "attn1.attn.in_proj_bias",
59
+ "attn1.attn.in_proj_weight",
60
+ "attn1.attn.out_proj.bias",
61
+ "attn1.attn.out_proj.weight",
62
+ "attn1.proj_in.bias",
63
+ "attn1.proj_in.weight",
64
+ "attn1.proj_out.bias",
65
+ "attn1.proj_out.weight",
66
+ "attn2.attn.in_proj_bias",
67
+ "attn2.attn.in_proj_weight",
68
+ "attn2.attn.out_proj.bias",
69
+ "attn2.attn.out_proj.weight",
70
+ "attn2.proj_in.bias",
71
+ "attn2.proj_in.weight",
72
+ "attn2.proj_out.bias",
73
+ "attn2.proj_out.weight",
74
+ "attn3.attn.in_proj_bias",
75
+ "attn3.attn.in_proj_weight",
76
+ "attn3.attn.out_proj.bias",
77
+ "attn3.attn.out_proj.weight",
78
+ "attn3.proj_in.bias",
79
+ "attn3.proj_in.weight",
80
+ "attn3.proj_out.bias",
81
+ "attn3.proj_out.weight",
82
+ "mask_token",
83
+ "norm0.bias",
84
+ "norm0.weight",
85
+ "norm1.bias",
86
+ "norm1.weight",
87
+ "norm2.bias",
88
+ "norm2.weight",
89
+ "norm3.bias",
90
+ "norm3.weight",
91
+ "norm_attn.bias",
92
+ "norm_attn.weight",
93
+ "patch_embed.norm.bias",
94
+ "patch_embed.norm.weight",
95
+ "patch_embed.projection.bias",
96
+ "patch_embed.projection.weight",
97
+ "stages.0.blocks.0.attn.w_msa.cpb_mlp.0.bias",
98
+ "stages.0.blocks.0.attn.w_msa.cpb_mlp.0.weight",
99
+ "stages.0.blocks.0.attn.w_msa.cpb_mlp.2.weight",
100
+ "stages.0.blocks.0.attn.w_msa.logit_scale",
101
+ "stages.0.blocks.0.attn.w_msa.proj.bias",
102
+ "stages.0.blocks.0.attn.w_msa.proj.weight",
103
+ "stages.0.blocks.0.attn.w_msa.q_bias",
104
+ "stages.0.blocks.0.attn.w_msa.qkv.weight",
105
+ "stages.0.blocks.0.attn.w_msa.v_bias",
106
+ "stages.0.blocks.0.ffn.layers.0.bias",
107
+ "stages.0.blocks.0.ffn.layers.0.weight",
108
+ "stages.0.blocks.0.ffn.layers.3.bias",
109
+ "stages.0.blocks.0.ffn.layers.3.weight",
110
+ "stages.0.blocks.0.norm1.bias",
111
+ "stages.0.blocks.0.norm1.weight",
112
+ "stages.0.blocks.0.norm2.bias",
113
+ "stages.0.blocks.0.norm2.weight",
114
+ "stages.0.blocks.1.attn.w_msa.cpb_mlp.0.bias",
115
+ "stages.0.blocks.1.attn.w_msa.cpb_mlp.0.weight",
116
+ "stages.0.blocks.1.attn.w_msa.cpb_mlp.2.weight",
117
+ "stages.0.blocks.1.attn.w_msa.logit_scale",
118
+ "stages.0.blocks.1.attn.w_msa.proj.bias",
119
+ "stages.0.blocks.1.attn.w_msa.proj.weight",
120
+ "stages.0.blocks.1.attn.w_msa.q_bias",
121
+ "stages.0.blocks.1.attn.w_msa.qkv.weight",
122
+ "stages.0.blocks.1.attn.w_msa.v_bias",
123
+ "stages.0.blocks.1.ffn.layers.0.bias",
124
+ "stages.0.blocks.1.ffn.layers.0.weight",
125
+ "stages.0.blocks.1.ffn.layers.3.bias",
126
+ "stages.0.blocks.1.ffn.layers.3.weight",
127
+ "stages.0.blocks.1.norm1.bias",
128
+ "stages.0.blocks.1.norm1.weight",
129
+ "stages.0.blocks.1.norm2.bias",
130
+ "stages.0.blocks.1.norm2.weight",
131
+ "stages.1.blocks.0.attn.w_msa.cpb_mlp.0.bias",
132
+ "stages.1.blocks.0.attn.w_msa.cpb_mlp.0.weight",
133
+ "stages.1.blocks.0.attn.w_msa.cpb_mlp.2.weight",
134
+ "stages.1.blocks.0.attn.w_msa.logit_scale",
135
+ "stages.1.blocks.0.attn.w_msa.proj.bias",
136
+ "stages.1.blocks.0.attn.w_msa.proj.weight",
137
+ "stages.1.blocks.0.attn.w_msa.q_bias",
138
+ "stages.1.blocks.0.attn.w_msa.qkv.weight",
139
+ "stages.1.blocks.0.attn.w_msa.v_bias",
140
+ "stages.1.blocks.0.ffn.layers.0.bias",
141
+ "stages.1.blocks.0.ffn.layers.0.weight",
142
+ "stages.1.blocks.0.ffn.layers.3.bias",
143
+ "stages.1.blocks.0.ffn.layers.3.weight",
144
+ "stages.1.blocks.0.norm1.bias",
145
+ "stages.1.blocks.0.norm1.weight",
146
+ "stages.1.blocks.0.norm2.bias",
147
+ "stages.1.blocks.0.norm2.weight",
148
+ "stages.1.blocks.1.attn.w_msa.cpb_mlp.0.bias",
149
+ "stages.1.blocks.1.attn.w_msa.cpb_mlp.0.weight",
150
+ "stages.1.blocks.1.attn.w_msa.cpb_mlp.2.weight",
151
+ "stages.1.blocks.1.attn.w_msa.logit_scale",
152
+ "stages.1.blocks.1.attn.w_msa.proj.bias",
153
+ "stages.1.blocks.1.attn.w_msa.proj.weight",
154
+ "stages.1.blocks.1.attn.w_msa.q_bias",
155
+ "stages.1.blocks.1.attn.w_msa.qkv.weight",
156
+ "stages.1.blocks.1.attn.w_msa.v_bias",
157
+ "stages.1.blocks.1.ffn.layers.0.bias",
158
+ "stages.1.blocks.1.ffn.layers.0.weight",
159
+ "stages.1.blocks.1.ffn.layers.3.bias",
160
+ "stages.1.blocks.1.ffn.layers.3.weight",
161
+ "stages.1.blocks.1.norm1.bias",
162
+ "stages.1.blocks.1.norm1.weight",
163
+ "stages.1.blocks.1.norm2.bias",
164
+ "stages.1.blocks.1.norm2.weight",
165
+ "stages.1.downsample.norm.bias",
166
+ "stages.1.downsample.norm.weight",
167
+ "stages.1.downsample.reduction.weight",
168
+ "stages.2.blocks.0.attn.w_msa.cpb_mlp.0.bias",
169
+ "stages.2.blocks.0.attn.w_msa.cpb_mlp.0.weight",
170
+ "stages.2.blocks.0.attn.w_msa.cpb_mlp.2.weight",
171
+ "stages.2.blocks.0.attn.w_msa.logit_scale",
172
+ "stages.2.blocks.0.attn.w_msa.proj.bias",
173
+ "stages.2.blocks.0.attn.w_msa.proj.weight",
174
+ "stages.2.blocks.0.attn.w_msa.q_bias",
175
+ "stages.2.blocks.0.attn.w_msa.qkv.weight",
176
+ "stages.2.blocks.0.attn.w_msa.v_bias",
177
+ "stages.2.blocks.0.ffn.layers.0.bias",
178
+ "stages.2.blocks.0.ffn.layers.0.weight",
179
+ "stages.2.blocks.0.ffn.layers.3.bias",
180
+ "stages.2.blocks.0.ffn.layers.3.weight",
181
+ "stages.2.blocks.0.norm1.bias",
182
+ "stages.2.blocks.0.norm1.weight",
183
+ "stages.2.blocks.0.norm2.bias",
184
+ "stages.2.blocks.0.norm2.weight",
185
+ "stages.2.blocks.1.attn.w_msa.cpb_mlp.0.bias",
186
+ "stages.2.blocks.1.attn.w_msa.cpb_mlp.0.weight",
187
+ "stages.2.blocks.1.attn.w_msa.cpb_mlp.2.weight",
188
+ "stages.2.blocks.1.attn.w_msa.logit_scale",
189
+ "stages.2.blocks.1.attn.w_msa.proj.bias",
190
+ "stages.2.blocks.1.attn.w_msa.proj.weight",
191
+ "stages.2.blocks.1.attn.w_msa.q_bias",
192
+ "stages.2.blocks.1.attn.w_msa.qkv.weight",
193
+ "stages.2.blocks.1.attn.w_msa.v_bias",
194
+ "stages.2.blocks.1.ffn.layers.0.bias",
195
+ "stages.2.blocks.1.ffn.layers.0.weight",
196
+ "stages.2.blocks.1.ffn.layers.3.bias",
197
+ "stages.2.blocks.1.ffn.layers.3.weight",
198
+ "stages.2.blocks.1.norm1.bias",
199
+ "stages.2.blocks.1.norm1.weight",
200
+ "stages.2.blocks.1.norm2.bias",
201
+ "stages.2.blocks.1.norm2.weight",
202
+ "stages.2.blocks.10.attn.w_msa.cpb_mlp.0.bias",
203
+ "stages.2.blocks.10.attn.w_msa.cpb_mlp.0.weight",
204
+ "stages.2.blocks.10.attn.w_msa.cpb_mlp.2.weight",
205
+ "stages.2.blocks.10.attn.w_msa.logit_scale",
206
+ "stages.2.blocks.10.attn.w_msa.proj.bias",
207
+ "stages.2.blocks.10.attn.w_msa.proj.weight",
208
+ "stages.2.blocks.10.attn.w_msa.q_bias",
209
+ "stages.2.blocks.10.attn.w_msa.qkv.weight",
210
+ "stages.2.blocks.10.attn.w_msa.v_bias",
211
+ "stages.2.blocks.10.ffn.layers.0.bias",
212
+ "stages.2.blocks.10.ffn.layers.0.weight",
213
+ "stages.2.blocks.10.ffn.layers.3.bias",
214
+ "stages.2.blocks.10.ffn.layers.3.weight",
215
+ "stages.2.blocks.10.norm1.bias",
216
+ "stages.2.blocks.10.norm1.weight",
217
+ "stages.2.blocks.10.norm2.bias",
218
+ "stages.2.blocks.10.norm2.weight",
219
+ "stages.2.blocks.11.attn.w_msa.cpb_mlp.0.bias",
220
+ "stages.2.blocks.11.attn.w_msa.cpb_mlp.0.weight",
221
+ "stages.2.blocks.11.attn.w_msa.cpb_mlp.2.weight",
222
+ "stages.2.blocks.11.attn.w_msa.logit_scale",
223
+ "stages.2.blocks.11.attn.w_msa.proj.bias",
224
+ "stages.2.blocks.11.attn.w_msa.proj.weight",
225
+ "stages.2.blocks.11.attn.w_msa.q_bias",
226
+ "stages.2.blocks.11.attn.w_msa.qkv.weight",
227
+ "stages.2.blocks.11.attn.w_msa.v_bias",
228
+ "stages.2.blocks.11.ffn.layers.0.bias",
229
+ "stages.2.blocks.11.ffn.layers.0.weight",
230
+ "stages.2.blocks.11.ffn.layers.3.bias",
231
+ "stages.2.blocks.11.ffn.layers.3.weight",
232
+ "stages.2.blocks.11.norm1.bias",
233
+ "stages.2.blocks.11.norm1.weight",
234
+ "stages.2.blocks.11.norm2.bias",
235
+ "stages.2.blocks.11.norm2.weight",
236
+ "stages.2.blocks.11.norm3.bias",
237
+ "stages.2.blocks.11.norm3.weight",
238
+ "stages.2.blocks.12.attn.w_msa.cpb_mlp.0.bias",
239
+ "stages.2.blocks.12.attn.w_msa.cpb_mlp.0.weight",
240
+ "stages.2.blocks.12.attn.w_msa.cpb_mlp.2.weight",
241
+ "stages.2.blocks.12.attn.w_msa.logit_scale",
242
+ "stages.2.blocks.12.attn.w_msa.proj.bias",
243
+ "stages.2.blocks.12.attn.w_msa.proj.weight",
244
+ "stages.2.blocks.12.attn.w_msa.q_bias",
245
+ "stages.2.blocks.12.attn.w_msa.qkv.weight",
246
+ "stages.2.blocks.12.attn.w_msa.v_bias",
247
+ "stages.2.blocks.12.ffn.layers.0.bias",
248
+ "stages.2.blocks.12.ffn.layers.0.weight",
249
+ "stages.2.blocks.12.ffn.layers.3.bias",
250
+ "stages.2.blocks.12.ffn.layers.3.weight",
251
+ "stages.2.blocks.12.norm1.bias",
252
+ "stages.2.blocks.12.norm1.weight",
253
+ "stages.2.blocks.12.norm2.bias",
254
+ "stages.2.blocks.12.norm2.weight",
255
+ "stages.2.blocks.13.attn.w_msa.cpb_mlp.0.bias",
256
+ "stages.2.blocks.13.attn.w_msa.cpb_mlp.0.weight",
257
+ "stages.2.blocks.13.attn.w_msa.cpb_mlp.2.weight",
258
+ "stages.2.blocks.13.attn.w_msa.logit_scale",
259
+ "stages.2.blocks.13.attn.w_msa.proj.bias",
260
+ "stages.2.blocks.13.attn.w_msa.proj.weight",
261
+ "stages.2.blocks.13.attn.w_msa.q_bias",
262
+ "stages.2.blocks.13.attn.w_msa.qkv.weight",
263
+ "stages.2.blocks.13.attn.w_msa.v_bias",
264
+ "stages.2.blocks.13.ffn.layers.0.bias",
265
+ "stages.2.blocks.13.ffn.layers.0.weight",
266
+ "stages.2.blocks.13.ffn.layers.3.bias",
267
+ "stages.2.blocks.13.ffn.layers.3.weight",
268
+ "stages.2.blocks.13.norm1.bias",
269
+ "stages.2.blocks.13.norm1.weight",
270
+ "stages.2.blocks.13.norm2.bias",
271
+ "stages.2.blocks.13.norm2.weight",
272
+ "stages.2.blocks.14.attn.w_msa.cpb_mlp.0.bias",
273
+ "stages.2.blocks.14.attn.w_msa.cpb_mlp.0.weight",
274
+ "stages.2.blocks.14.attn.w_msa.cpb_mlp.2.weight",
275
+ "stages.2.blocks.14.attn.w_msa.logit_scale",
276
+ "stages.2.blocks.14.attn.w_msa.proj.bias",
277
+ "stages.2.blocks.14.attn.w_msa.proj.weight",
278
+ "stages.2.blocks.14.attn.w_msa.q_bias",
279
+ "stages.2.blocks.14.attn.w_msa.qkv.weight",
280
+ "stages.2.blocks.14.attn.w_msa.v_bias",
281
+ "stages.2.blocks.14.ffn.layers.0.bias",
282
+ "stages.2.blocks.14.ffn.layers.0.weight",
283
+ "stages.2.blocks.14.ffn.layers.3.bias",
284
+ "stages.2.blocks.14.ffn.layers.3.weight",
285
+ "stages.2.blocks.14.norm1.bias",
286
+ "stages.2.blocks.14.norm1.weight",
287
+ "stages.2.blocks.14.norm2.bias",
288
+ "stages.2.blocks.14.norm2.weight",
289
+ "stages.2.blocks.15.attn.w_msa.cpb_mlp.0.bias",
290
+ "stages.2.blocks.15.attn.w_msa.cpb_mlp.0.weight",
291
+ "stages.2.blocks.15.attn.w_msa.cpb_mlp.2.weight",
292
+ "stages.2.blocks.15.attn.w_msa.logit_scale",
293
+ "stages.2.blocks.15.attn.w_msa.proj.bias",
294
+ "stages.2.blocks.15.attn.w_msa.proj.weight",
295
+ "stages.2.blocks.15.attn.w_msa.q_bias",
296
+ "stages.2.blocks.15.attn.w_msa.qkv.weight",
297
+ "stages.2.blocks.15.attn.w_msa.v_bias",
298
+ "stages.2.blocks.15.ffn.layers.0.bias",
299
+ "stages.2.blocks.15.ffn.layers.0.weight",
300
+ "stages.2.blocks.15.ffn.layers.3.bias",
301
+ "stages.2.blocks.15.ffn.layers.3.weight",
302
+ "stages.2.blocks.15.norm1.bias",
303
+ "stages.2.blocks.15.norm1.weight",
304
+ "stages.2.blocks.15.norm2.bias",
305
+ "stages.2.blocks.15.norm2.weight",
306
+ "stages.2.blocks.16.attn.w_msa.cpb_mlp.0.bias",
307
+ "stages.2.blocks.16.attn.w_msa.cpb_mlp.0.weight",
308
+ "stages.2.blocks.16.attn.w_msa.cpb_mlp.2.weight",
309
+ "stages.2.blocks.16.attn.w_msa.logit_scale",
310
+ "stages.2.blocks.16.attn.w_msa.proj.bias",
311
+ "stages.2.blocks.16.attn.w_msa.proj.weight",
312
+ "stages.2.blocks.16.attn.w_msa.q_bias",
313
+ "stages.2.blocks.16.attn.w_msa.qkv.weight",
314
+ "stages.2.blocks.16.attn.w_msa.v_bias",
315
+ "stages.2.blocks.16.ffn.layers.0.bias",
316
+ "stages.2.blocks.16.ffn.layers.0.weight",
317
+ "stages.2.blocks.16.ffn.layers.3.bias",
318
+ "stages.2.blocks.16.ffn.layers.3.weight",
319
+ "stages.2.blocks.16.norm1.bias",
320
+ "stages.2.blocks.16.norm1.weight",
321
+ "stages.2.blocks.16.norm2.bias",
322
+ "stages.2.blocks.16.norm2.weight",
323
+ "stages.2.blocks.17.attn.w_msa.cpb_mlp.0.bias",
324
+ "stages.2.blocks.17.attn.w_msa.cpb_mlp.0.weight",
325
+ "stages.2.blocks.17.attn.w_msa.cpb_mlp.2.weight",
326
+ "stages.2.blocks.17.attn.w_msa.logit_scale",
327
+ "stages.2.blocks.17.attn.w_msa.proj.bias",
328
+ "stages.2.blocks.17.attn.w_msa.proj.weight",
329
+ "stages.2.blocks.17.attn.w_msa.q_bias",
330
+ "stages.2.blocks.17.attn.w_msa.qkv.weight",
331
+ "stages.2.blocks.17.attn.w_msa.v_bias",
332
+ "stages.2.blocks.17.ffn.layers.0.bias",
333
+ "stages.2.blocks.17.ffn.layers.0.weight",
334
+ "stages.2.blocks.17.ffn.layers.3.bias",
335
+ "stages.2.blocks.17.ffn.layers.3.weight",
336
+ "stages.2.blocks.17.norm1.bias",
337
+ "stages.2.blocks.17.norm1.weight",
338
+ "stages.2.blocks.17.norm2.bias",
339
+ "stages.2.blocks.17.norm2.weight",
340
+ "stages.2.blocks.17.norm3.bias",
341
+ "stages.2.blocks.17.norm3.weight",
342
+ "stages.2.blocks.2.attn.w_msa.cpb_mlp.0.bias",
343
+ "stages.2.blocks.2.attn.w_msa.cpb_mlp.0.weight",
344
+ "stages.2.blocks.2.attn.w_msa.cpb_mlp.2.weight",
345
+ "stages.2.blocks.2.attn.w_msa.logit_scale",
346
+ "stages.2.blocks.2.attn.w_msa.proj.bias",
347
+ "stages.2.blocks.2.attn.w_msa.proj.weight",
348
+ "stages.2.blocks.2.attn.w_msa.q_bias",
349
+ "stages.2.blocks.2.attn.w_msa.qkv.weight",
350
+ "stages.2.blocks.2.attn.w_msa.v_bias",
351
+ "stages.2.blocks.2.ffn.layers.0.bias",
352
+ "stages.2.blocks.2.ffn.layers.0.weight",
353
+ "stages.2.blocks.2.ffn.layers.3.bias",
354
+ "stages.2.blocks.2.ffn.layers.3.weight",
355
+ "stages.2.blocks.2.norm1.bias",
356
+ "stages.2.blocks.2.norm1.weight",
357
+ "stages.2.blocks.2.norm2.bias",
358
+ "stages.2.blocks.2.norm2.weight",
359
+ "stages.2.blocks.3.attn.w_msa.cpb_mlp.0.bias",
360
+ "stages.2.blocks.3.attn.w_msa.cpb_mlp.0.weight",
361
+ "stages.2.blocks.3.attn.w_msa.cpb_mlp.2.weight",
362
+ "stages.2.blocks.3.attn.w_msa.logit_scale",
363
+ "stages.2.blocks.3.attn.w_msa.proj.bias",
364
+ "stages.2.blocks.3.attn.w_msa.proj.weight",
365
+ "stages.2.blocks.3.attn.w_msa.q_bias",
366
+ "stages.2.blocks.3.attn.w_msa.qkv.weight",
367
+ "stages.2.blocks.3.attn.w_msa.v_bias",
368
+ "stages.2.blocks.3.ffn.layers.0.bias",
369
+ "stages.2.blocks.3.ffn.layers.0.weight",
370
+ "stages.2.blocks.3.ffn.layers.3.bias",
371
+ "stages.2.blocks.3.ffn.layers.3.weight",
372
+ "stages.2.blocks.3.norm1.bias",
373
+ "stages.2.blocks.3.norm1.weight",
374
+ "stages.2.blocks.3.norm2.bias",
375
+ "stages.2.blocks.3.norm2.weight",
376
+ "stages.2.blocks.4.attn.w_msa.cpb_mlp.0.bias",
377
+ "stages.2.blocks.4.attn.w_msa.cpb_mlp.0.weight",
378
+ "stages.2.blocks.4.attn.w_msa.cpb_mlp.2.weight",
379
+ "stages.2.blocks.4.attn.w_msa.logit_scale",
380
+ "stages.2.blocks.4.attn.w_msa.proj.bias",
381
+ "stages.2.blocks.4.attn.w_msa.proj.weight",
382
+ "stages.2.blocks.4.attn.w_msa.q_bias",
383
+ "stages.2.blocks.4.attn.w_msa.qkv.weight",
384
+ "stages.2.blocks.4.attn.w_msa.v_bias",
385
+ "stages.2.blocks.4.ffn.layers.0.bias",
386
+ "stages.2.blocks.4.ffn.layers.0.weight",
387
+ "stages.2.blocks.4.ffn.layers.3.bias",
388
+ "stages.2.blocks.4.ffn.layers.3.weight",
389
+ "stages.2.blocks.4.norm1.bias",
390
+ "stages.2.blocks.4.norm1.weight",
391
+ "stages.2.blocks.4.norm2.bias",
392
+ "stages.2.blocks.4.norm2.weight",
393
+ "stages.2.blocks.5.attn.w_msa.cpb_mlp.0.bias",
394
+ "stages.2.blocks.5.attn.w_msa.cpb_mlp.0.weight",
395
+ "stages.2.blocks.5.attn.w_msa.cpb_mlp.2.weight",
396
+ "stages.2.blocks.5.attn.w_msa.logit_scale",
397
+ "stages.2.blocks.5.attn.w_msa.proj.bias",
398
+ "stages.2.blocks.5.attn.w_msa.proj.weight",
399
+ "stages.2.blocks.5.attn.w_msa.q_bias",
400
+ "stages.2.blocks.5.attn.w_msa.qkv.weight",
401
+ "stages.2.blocks.5.attn.w_msa.v_bias",
402
+ "stages.2.blocks.5.ffn.layers.0.bias",
403
+ "stages.2.blocks.5.ffn.layers.0.weight",
404
+ "stages.2.blocks.5.ffn.layers.3.bias",
405
+ "stages.2.blocks.5.ffn.layers.3.weight",
406
+ "stages.2.blocks.5.norm1.bias",
407
+ "stages.2.blocks.5.norm1.weight",
408
+ "stages.2.blocks.5.norm2.bias",
409
+ "stages.2.blocks.5.norm2.weight",
410
+ "stages.2.blocks.5.norm3.bias",
411
+ "stages.2.blocks.5.norm3.weight",
412
+ "stages.2.blocks.6.attn.w_msa.cpb_mlp.0.bias",
413
+ "stages.2.blocks.6.attn.w_msa.cpb_mlp.0.weight",
414
+ "stages.2.blocks.6.attn.w_msa.cpb_mlp.2.weight",
415
+ "stages.2.blocks.6.attn.w_msa.logit_scale",
416
+ "stages.2.blocks.6.attn.w_msa.proj.bias",
417
+ "stages.2.blocks.6.attn.w_msa.proj.weight",
418
+ "stages.2.blocks.6.attn.w_msa.q_bias",
419
+ "stages.2.blocks.6.attn.w_msa.qkv.weight",
420
+ "stages.2.blocks.6.attn.w_msa.v_bias",
421
+ "stages.2.blocks.6.ffn.layers.0.bias",
422
+ "stages.2.blocks.6.ffn.layers.0.weight",
423
+ "stages.2.blocks.6.ffn.layers.3.bias",
424
+ "stages.2.blocks.6.ffn.layers.3.weight",
425
+ "stages.2.blocks.6.norm1.bias",
426
+ "stages.2.blocks.6.norm1.weight",
427
+ "stages.2.blocks.6.norm2.bias",
428
+ "stages.2.blocks.6.norm2.weight",
429
+ "stages.2.blocks.7.attn.w_msa.cpb_mlp.0.bias",
430
+ "stages.2.blocks.7.attn.w_msa.cpb_mlp.0.weight",
431
+ "stages.2.blocks.7.attn.w_msa.cpb_mlp.2.weight",
432
+ "stages.2.blocks.7.attn.w_msa.logit_scale",
433
+ "stages.2.blocks.7.attn.w_msa.proj.bias",
434
+ "stages.2.blocks.7.attn.w_msa.proj.weight",
435
+ "stages.2.blocks.7.attn.w_msa.q_bias",
436
+ "stages.2.blocks.7.attn.w_msa.qkv.weight",
437
+ "stages.2.blocks.7.attn.w_msa.v_bias",
438
+ "stages.2.blocks.7.ffn.layers.0.bias",
439
+ "stages.2.blocks.7.ffn.layers.0.weight",
440
+ "stages.2.blocks.7.ffn.layers.3.bias",
441
+ "stages.2.blocks.7.ffn.layers.3.weight",
442
+ "stages.2.blocks.7.norm1.bias",
443
+ "stages.2.blocks.7.norm1.weight",
444
+ "stages.2.blocks.7.norm2.bias",
445
+ "stages.2.blocks.7.norm2.weight",
446
+ "stages.2.blocks.8.attn.w_msa.cpb_mlp.0.bias",
447
+ "stages.2.blocks.8.attn.w_msa.cpb_mlp.0.weight",
448
+ "stages.2.blocks.8.attn.w_msa.cpb_mlp.2.weight",
449
+ "stages.2.blocks.8.attn.w_msa.logit_scale",
450
+ "stages.2.blocks.8.attn.w_msa.proj.bias",
451
+ "stages.2.blocks.8.attn.w_msa.proj.weight",
452
+ "stages.2.blocks.8.attn.w_msa.q_bias",
453
+ "stages.2.blocks.8.attn.w_msa.qkv.weight",
454
+ "stages.2.blocks.8.attn.w_msa.v_bias",
455
+ "stages.2.blocks.8.ffn.layers.0.bias",
456
+ "stages.2.blocks.8.ffn.layers.0.weight",
457
+ "stages.2.blocks.8.ffn.layers.3.bias",
458
+ "stages.2.blocks.8.ffn.layers.3.weight",
459
+ "stages.2.blocks.8.norm1.bias",
460
+ "stages.2.blocks.8.norm1.weight",
461
+ "stages.2.blocks.8.norm2.bias",
462
+ "stages.2.blocks.8.norm2.weight",
463
+ "stages.2.blocks.9.attn.w_msa.cpb_mlp.0.bias",
464
+ "stages.2.blocks.9.attn.w_msa.cpb_mlp.0.weight",
465
+ "stages.2.blocks.9.attn.w_msa.cpb_mlp.2.weight",
466
+ "stages.2.blocks.9.attn.w_msa.logit_scale",
467
+ "stages.2.blocks.9.attn.w_msa.proj.bias",
468
+ "stages.2.blocks.9.attn.w_msa.proj.weight",
469
+ "stages.2.blocks.9.attn.w_msa.q_bias",
470
+ "stages.2.blocks.9.attn.w_msa.qkv.weight",
471
+ "stages.2.blocks.9.attn.w_msa.v_bias",
472
+ "stages.2.blocks.9.ffn.layers.0.bias",
473
+ "stages.2.blocks.9.ffn.layers.0.weight",
474
+ "stages.2.blocks.9.ffn.layers.3.bias",
475
+ "stages.2.blocks.9.ffn.layers.3.weight",
476
+ "stages.2.blocks.9.norm1.bias",
477
+ "stages.2.blocks.9.norm1.weight",
478
+ "stages.2.blocks.9.norm2.bias",
479
+ "stages.2.blocks.9.norm2.weight",
480
+ "stages.2.downsample.norm.bias",
481
+ "stages.2.downsample.norm.weight",
482
+ "stages.2.downsample.reduction.weight",
483
+ "stages.3.blocks.0.attn.w_msa.cpb_mlp.0.bias",
484
+ "stages.3.blocks.0.attn.w_msa.cpb_mlp.0.weight",
485
+ "stages.3.blocks.0.attn.w_msa.cpb_mlp.2.weight",
486
+ "stages.3.blocks.0.attn.w_msa.logit_scale",
487
+ "stages.3.blocks.0.attn.w_msa.proj.bias",
488
+ "stages.3.blocks.0.attn.w_msa.proj.weight",
489
+ "stages.3.blocks.0.attn.w_msa.q_bias",
490
+ "stages.3.blocks.0.attn.w_msa.qkv.weight",
491
+ "stages.3.blocks.0.attn.w_msa.v_bias",
492
+ "stages.3.blocks.0.ffn.layers.0.bias",
493
+ "stages.3.blocks.0.ffn.layers.0.weight",
494
+ "stages.3.blocks.0.ffn.layers.3.bias",
495
+ "stages.3.blocks.0.ffn.layers.3.weight",
496
+ "stages.3.blocks.0.norm1.bias",
497
+ "stages.3.blocks.0.norm1.weight",
498
+ "stages.3.blocks.0.norm2.bias",
499
+ "stages.3.blocks.0.norm2.weight",
500
+ "stages.3.blocks.1.attn.w_msa.cpb_mlp.0.bias",
501
+ "stages.3.blocks.1.attn.w_msa.cpb_mlp.0.weight",
502
+ "stages.3.blocks.1.attn.w_msa.cpb_mlp.2.weight",
503
+ "stages.3.blocks.1.attn.w_msa.logit_scale",
504
+ "stages.3.blocks.1.attn.w_msa.proj.bias",
505
+ "stages.3.blocks.1.attn.w_msa.proj.weight",
506
+ "stages.3.blocks.1.attn.w_msa.q_bias",
507
+ "stages.3.blocks.1.attn.w_msa.qkv.weight",
508
+ "stages.3.blocks.1.attn.w_msa.v_bias",
509
+ "stages.3.blocks.1.ffn.layers.0.bias",
510
+ "stages.3.blocks.1.ffn.layers.0.weight",
511
+ "stages.3.blocks.1.ffn.layers.3.bias",
512
+ "stages.3.blocks.1.ffn.layers.3.weight",
513
+ "stages.3.blocks.1.norm1.bias",
514
+ "stages.3.blocks.1.norm1.weight",
515
+ "stages.3.blocks.1.norm2.bias",
516
+ "stages.3.blocks.1.norm2.weight",
517
+ "stages.3.downsample.norm.bias",
518
+ "stages.3.downsample.norm.weight",
519
+ "stages.3.downsample.reduction.weight",
520
+ "vocabulary_token",
521
+ "vocabulary_weight"
522
+ ]
523
+ }
skysensepp-swinv2-msl-hr/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2727a8049e5cc0fe0e16a1c7c4ccb6f4d9cb1df2ed3eae2dba857b507dd49d9d
3
+ size 2658512808
skysensepp-swinv2-msl-hr/modeling_skysensepp_swinv2_msl.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ from copy import deepcopy
4
+ from typing import Optional, Sequence, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig
14
+ from .modeling_utils import (
15
+ DropPath,
16
+ FFN,
17
+ PatchEmbed,
18
+ PatchMerging,
19
+ ShiftWindowMSA,
20
+ to_2tuple,
21
+ )
22
+
23
+
24
+ class SwinBlockV2(nn.Module):
25
+ def __init__(
26
+ self,
27
+ embed_dims: int,
28
+ num_heads: int,
29
+ window_size: int = 8,
30
+ shift: bool = False,
31
+ extra_norm: bool = False,
32
+ ffn_ratio: float = 4.0,
33
+ drop_path: float = 0.0,
34
+ pad_small_map: bool = False,
35
+ with_cp: bool = False,
36
+ pretrained_window_size: int = 0,
37
+ ):
38
+ super().__init__()
39
+ self.with_cp = with_cp
40
+ self.extra_norm = extra_norm
41
+ self.attn = ShiftWindowMSA(
42
+ embed_dims=embed_dims,
43
+ num_heads=num_heads,
44
+ window_size=window_size,
45
+ shift_size=window_size // 2 if shift else 0,
46
+ drop_path=drop_path,
47
+ pad_small_map=pad_small_map,
48
+ pretrained_window_size=pretrained_window_size,
49
+ )
50
+ self.norm1 = nn.LayerNorm(embed_dims)
51
+ self.ffn = FFN(
52
+ embed_dims=embed_dims,
53
+ feedforward_channels=int(embed_dims * ffn_ratio),
54
+ num_fcs=2,
55
+ drop_path=drop_path,
56
+ act_layer=nn.GELU,
57
+ add_identity=False,
58
+ )
59
+ self.norm2 = nn.LayerNorm(embed_dims)
60
+ if self.extra_norm:
61
+ self.norm3 = nn.LayerNorm(embed_dims)
62
+
63
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
64
+ def _inner_forward(x):
65
+ identity = x
66
+ x = self.attn(x, hw_shape)
67
+ x = self.norm1(x)
68
+ x = x + identity
69
+
70
+ identity = x
71
+ x = self.ffn(x)
72
+ x = self.norm2(x)
73
+ x = x + identity
74
+
75
+ if self.extra_norm:
76
+ x = self.norm3(x)
77
+ return x
78
+
79
+ if self.with_cp and x.requires_grad:
80
+ x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
81
+ else:
82
+ x = _inner_forward(x)
83
+ return x
84
+
85
+
86
+ class SwinBlockV2Sequence(nn.Module):
87
+ def __init__(
88
+ self,
89
+ embed_dims: int,
90
+ depth: int,
91
+ num_heads: int,
92
+ window_size: int = 8,
93
+ downsample: bool = False,
94
+ drop_paths: Union[Sequence[float], float] = 0.0,
95
+ with_cp: bool = False,
96
+ pad_small_map: bool = False,
97
+ extra_norm_every_n_blocks: int = 0,
98
+ pretrained_window_size: int = 0,
99
+ is_post_norm_downsample: bool = True,
100
+ ):
101
+ super().__init__()
102
+ if not isinstance(drop_paths, Sequence):
103
+ drop_paths = [drop_paths] * depth
104
+
105
+ if downsample:
106
+ self.out_channels = 2 * embed_dims
107
+ self.downsample = PatchMerging(
108
+ in_channels=embed_dims,
109
+ out_channels=self.out_channels,
110
+ is_post_norm=is_post_norm_downsample,
111
+ )
112
+ else:
113
+ self.out_channels = embed_dims
114
+ self.downsample = None
115
+
116
+ self.blocks = nn.ModuleList()
117
+ for i in range(depth):
118
+ extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0
119
+ self.blocks.append(
120
+ SwinBlockV2(
121
+ embed_dims=self.out_channels,
122
+ num_heads=num_heads,
123
+ window_size=window_size,
124
+ shift=(i % 2 == 1),
125
+ extra_norm=extra_norm,
126
+ drop_path=drop_paths[i],
127
+ with_cp=with_cp,
128
+ pad_small_map=pad_small_map,
129
+ pretrained_window_size=pretrained_window_size,
130
+ )
131
+ )
132
+
133
+ def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
134
+ if self.downsample is not None:
135
+ x, out_shape = self.downsample(x, in_shape)
136
+ else:
137
+ out_shape = in_shape
138
+
139
+ for block in self.blocks:
140
+ x = block(x, out_shape)
141
+ return x, out_shape
142
+
143
+
144
+ class ProjMHSA(nn.Module):
145
+ """Projected multi-head self-attention used in SkySense++ HR backbone."""
146
+
147
+ def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True):
148
+ super().__init__()
149
+ self.proj_in = nn.Linear(embed_dims, proj_dims)
150
+ self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias)
151
+ self.proj_out = nn.Linear(proj_dims, embed_dims)
152
+
153
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
154
+ x = self.proj_in(x)
155
+ x, _ = self.attn(x, x, x)
156
+ return self.proj_out(x)
157
+
158
+
159
+ class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel):
160
+ config_class = SkySensePlusPlusSwinV2MSLConfig
161
+ base_model_prefix = "skysensepp_swinv2_msl"
162
+ supports_gradient_checkpointing = True
163
+
164
+ def _init_weights(self, module):
165
+ if isinstance(module, nn.Linear):
166
+ nn.init.trunc_normal_(module.weight, std=0.02)
167
+ if module.bias is not None:
168
+ nn.init.zeros_(module.bias)
169
+ elif isinstance(module, nn.LayerNorm):
170
+ nn.init.ones_(module.weight)
171
+ nn.init.zeros_(module.bias)
172
+ elif isinstance(module, nn.Conv2d):
173
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
174
+ if module.bias is not None:
175
+ nn.init.zeros_(module.bias)
176
+
177
+
178
+ class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel):
179
+ """SkySense++ HR backbone with semantic vocabulary and annotation conditioning."""
180
+
181
+ def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig):
182
+ super().__init__(config)
183
+
184
+ self.num_layers = len(config.depths)
185
+ self.out_indices = config.out_indices
186
+ self.merge_stage = config.merge_stage
187
+ self.use_attn = config.use_attn
188
+ self.patch_size = config.patch_size
189
+
190
+ if isinstance(config.window_size, int):
191
+ window_sizes = [config.window_size] * self.num_layers
192
+ else:
193
+ window_sizes = list(config.window_size)
194
+
195
+ self.patch_embed = PatchEmbed(
196
+ in_channels=config.in_channels,
197
+ embed_dims=config.embed_dims,
198
+ kernel_size=config.patch_size,
199
+ stride=config.patch_size,
200
+ norm_layer=nn.LayerNorm,
201
+ input_size=config.img_size,
202
+ )
203
+
204
+ self.use_abs_pos_embed = config.use_abs_pos_embed
205
+ if self.use_abs_pos_embed:
206
+ patch_resolution = self.patch_embed.init_out_size
207
+ num_patches = patch_resolution[0] * patch_resolution[1]
208
+ self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
209
+
210
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
211
+
212
+ total_depth = sum(config.depths)
213
+ if total_depth > 1:
214
+ dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
215
+ else:
216
+ dpr = [0.0]
217
+
218
+ self.stages = nn.ModuleList()
219
+ embed_dims_list = [config.embed_dims]
220
+ for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)):
221
+ stage = SwinBlockV2Sequence(
222
+ embed_dims=embed_dims_list[-1],
223
+ depth=depth,
224
+ num_heads=num_heads,
225
+ window_size=window_sizes[i],
226
+ downsample=(i > 0),
227
+ drop_paths=dpr[:depth],
228
+ with_cp=config.with_cp,
229
+ pad_small_map=config.pad_small_map,
230
+ extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
231
+ pretrained_window_size=config.pretrained_window_sizes[i],
232
+ is_post_norm_downsample=config.is_post_norm_downsample,
233
+ )
234
+ self.stages.append(stage)
235
+ dpr = dpr[depth:]
236
+ embed_dims_list.append(stage.out_channels)
237
+
238
+ for i in self.out_indices:
239
+ self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1]))
240
+
241
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
242
+ self.vocabulary_token = nn.Parameter(
243
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
244
+ )
245
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
246
+
247
+ if self.use_attn:
248
+ self.attn1 = ProjMHSA(352, 256, num_heads=16)
249
+ self.attn2 = ProjMHSA(704, 512, num_heads=16)
250
+ self.attn3 = ProjMHSA(1408, 1024, num_heads=16)
251
+ self.norm_attn = nn.LayerNorm(1408)
252
+
253
+ self.post_init()
254
+
255
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
256
+ batch_size, height, width = anno_img.shape
257
+ ann_token = torch.index_select(
258
+ self.vocabulary_token, 0, anno_img.reshape(-1)
259
+ ).reshape(batch_size, height, width, -1)
260
+
261
+ num_patch_h = height // self.patch_size
262
+ num_patch_w = width // self.patch_size
263
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
264
+ weight = (
265
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
266
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
267
+ .reshape(1, height, width, 1)
268
+ )
269
+ ann_token = ann_token * weight
270
+ ann_token = F.avg_pool2d(
271
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
272
+ )
273
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
274
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
275
+ )
276
+
277
+ def forward(
278
+ self,
279
+ pixel_values: torch.Tensor,
280
+ annotation: torch.Tensor,
281
+ mask: Optional[torch.Tensor] = None,
282
+ output_hidden_states: Optional[bool] = None,
283
+ return_dict: Optional[bool] = None,
284
+ ) -> Union[Tuple, BaseModelOutput]:
285
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
286
+
287
+ x, hw_shape = self.patch_embed(pixel_values)
288
+ y = self.create_ann_token(annotation)
289
+ batch_size, num_tokens, channels = y.shape
290
+
291
+ if mask is not None:
292
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
293
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
294
+ y = y * (1.0 - weight) + mask_tokens * weight
295
+
296
+ if self.merge_stage == 0:
297
+ x = (x + y) * 0.5
298
+ else:
299
+ x = x.reshape(batch_size, *hw_shape, channels)
300
+ y = y.reshape(batch_size, *hw_shape, channels)
301
+ x = torch.cat((x, y), dim=2)
302
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
303
+ x = x.reshape(batch_size, -1, channels)
304
+
305
+ if self.use_abs_pos_embed:
306
+ x = x + self.absolute_pos_embed
307
+ x = self.drop_after_pos(x)
308
+
309
+ all_hidden_states = () if output_hidden_states else None
310
+ feature_maps = []
311
+ merge_idx = self.merge_stage - 1
312
+
313
+ for i, stage in enumerate(self.stages):
314
+ x, hw_shape = stage(x, hw_shape)
315
+ if i == merge_idx:
316
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
317
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
318
+ x = x.reshape(batch_size, -1, x.shape[-1])
319
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
320
+
321
+ if self.use_attn:
322
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
323
+ if i <= len(attention_blocks) - 1:
324
+ x = x + attention_blocks[i](x)
325
+ if i == len(attention_blocks) - 1:
326
+ x = self.norm_attn(x)
327
+
328
+ if output_hidden_states:
329
+ all_hidden_states = all_hidden_states + (x,)
330
+
331
+ if i in self.out_indices:
332
+ norm_layer = getattr(self, f"norm{i}")
333
+ out = norm_layer(x)
334
+ out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous()
335
+ feature_maps.append(out)
336
+
337
+ if not return_dict:
338
+ return tuple(feature_maps)
339
+
340
+ return BaseModelOutput(
341
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
342
+ hidden_states=all_hidden_states,
343
+ )
skysensepp-swinv2-msl-hr/modeling_skysensepp_vit_msl.py ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ import math
4
+ from typing import Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig
14
+ from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
15
+
16
+
17
+ class TransformerEncoderLayer(nn.Module):
18
+ def __init__(
19
+ self,
20
+ embed_dims: int,
21
+ num_heads: int,
22
+ feedforward_channels: int,
23
+ drop_rate: float = 0.0,
24
+ attn_drop_rate: float = 0.0,
25
+ drop_path_rate: float = 0.0,
26
+ num_fcs: int = 2,
27
+ qkv_bias: bool = True,
28
+ with_cp: bool = False,
29
+ ):
30
+ super().__init__()
31
+ self.with_cp = with_cp
32
+ self.norm1 = nn.LayerNorm(embed_dims)
33
+ self.attn = nn.MultiheadAttention(
34
+ embed_dim=embed_dims,
35
+ num_heads=num_heads,
36
+ dropout=attn_drop_rate,
37
+ bias=qkv_bias,
38
+ batch_first=True,
39
+ )
40
+ self.proj_drop = nn.Dropout(drop_rate)
41
+ self.norm2 = nn.LayerNorm(embed_dims)
42
+ self.ffn = FFN(
43
+ embed_dims=embed_dims,
44
+ feedforward_channels=feedforward_channels,
45
+ num_fcs=num_fcs,
46
+ ffn_drop=drop_rate,
47
+ drop_path=drop_path_rate,
48
+ act_layer=nn.GELU,
49
+ add_identity=True,
50
+ )
51
+ self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
52
+
53
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
54
+ def _inner_forward(x):
55
+ residual = x
56
+ x_norm = self.norm1(x)
57
+ attn_out, _ = self.attn(x_norm, x_norm, x_norm)
58
+ attn_out = self.proj_drop(attn_out)
59
+ x = residual + self.drop_path(attn_out)
60
+ return self.ffn(self.norm2(x), identity=x)
61
+
62
+ if self.with_cp and x.requires_grad:
63
+ return cp.checkpoint(_inner_forward, x, use_reentrant=False)
64
+ return _inner_forward(x)
65
+
66
+
67
+ class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel):
68
+ config_class = SkySensePlusPlusViTMSLConfig
69
+ base_model_prefix = "skysensepp_vit_msl"
70
+ supports_gradient_checkpointing = True
71
+
72
+ def _init_weights(self, module):
73
+ if isinstance(module, nn.Linear):
74
+ nn.init.trunc_normal_(module.weight, std=0.02)
75
+ if module.bias is not None:
76
+ nn.init.zeros_(module.bias)
77
+ elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
78
+ nn.init.ones_(module.weight)
79
+ nn.init.zeros_(module.bias)
80
+ elif isinstance(module, nn.Conv2d):
81
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
82
+ if module.bias is not None:
83
+ nn.init.zeros_(module.bias)
84
+
85
+
86
+ class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel):
87
+ """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning."""
88
+
89
+ def __init__(self, config: SkySensePlusPlusViTMSLConfig):
90
+ super().__init__(config)
91
+
92
+ img_size = to_2tuple(config.img_size)
93
+ self.img_size = img_size
94
+ self.patch_size = config.patch_size
95
+ self.with_cls_token = config.with_cls_token
96
+ self.output_cls_token = config.output_cls_token
97
+ self.merge_stage = config.merge_stage
98
+ self.use_attn = config.use_attn
99
+ self.interpolate_mode = "bicubic"
100
+
101
+ self.patch_embed = PatchEmbed(
102
+ in_channels=config.in_channels,
103
+ embed_dims=config.embed_dims,
104
+ kernel_size=config.patch_size,
105
+ stride=config.patch_size,
106
+ norm_layer=nn.LayerNorm if config.patch_norm else None,
107
+ )
108
+
109
+ num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size)
110
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
111
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
112
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
113
+
114
+ out_indices = list(config.out_indices)
115
+ self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices]
116
+
117
+ num_layers = config.num_layers
118
+ if num_layers > 1:
119
+ dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
120
+ else:
121
+ dpr = [0.0]
122
+
123
+ self.layers = nn.ModuleList()
124
+ for i in range(config.num_layers):
125
+ self.layers.append(
126
+ TransformerEncoderLayer(
127
+ embed_dims=config.embed_dims,
128
+ num_heads=config.num_heads,
129
+ feedforward_channels=config.mlp_ratio * config.embed_dims,
130
+ attn_drop_rate=config.attn_drop_rate,
131
+ drop_rate=config.drop_rate,
132
+ drop_path_rate=dpr[i],
133
+ num_fcs=2,
134
+ qkv_bias=config.qkv_bias,
135
+ with_cp=config.with_cp,
136
+ )
137
+ )
138
+
139
+ self.final_norm = config.final_norm
140
+ if config.final_norm:
141
+ self.norm = nn.LayerNorm(config.embed_dims)
142
+
143
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
144
+ self.vocabulary_token = nn.Parameter(
145
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
146
+ )
147
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
148
+
149
+ if self.use_attn:
150
+ self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
151
+ self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
152
+ self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
153
+ self.norm_attn = nn.LayerNorm(config.embed_dims)
154
+
155
+ self.post_init()
156
+
157
+ @staticmethod
158
+ def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"):
159
+ pos_h, pos_w = pos_shape
160
+ pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :]
161
+ pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
162
+ pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode)
163
+ return torch.flatten(pos_embed_weight, 2).transpose(1, 2)
164
+
165
+ def _pos_embedding(self, patched_img, hw_shape, pos_embed):
166
+ x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
167
+ if x_len != pos_len:
168
+ pos_h = self.img_size[0] // self.patch_size
169
+ pos_w = self.img_size[1] // self.patch_size
170
+ pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode)
171
+ return self.drop_after_pos(patched_img + pos_embed)
172
+
173
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
174
+ batch_size, height, width = anno_img.shape
175
+ ann_token = torch.index_select(
176
+ self.vocabulary_token, 0, anno_img.reshape(-1)
177
+ ).reshape(batch_size, height, width, -1)
178
+
179
+ num_patch_h = height // self.patch_size
180
+ num_patch_w = width // self.patch_size
181
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
182
+ weight = (
183
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
184
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
185
+ .reshape(1, height, width, 1)
186
+ )
187
+ ann_token = ann_token * weight
188
+ ann_token = F.avg_pool2d(
189
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
190
+ )
191
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
192
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
193
+ )
194
+
195
+ def forward(
196
+ self,
197
+ pixel_values: torch.Tensor,
198
+ annotation: torch.Tensor,
199
+ mask: Optional[torch.Tensor] = None,
200
+ output_hidden_states: Optional[bool] = None,
201
+ return_dict: Optional[bool] = None,
202
+ ) -> Union[Tuple, BaseModelOutput]:
203
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
204
+
205
+ x, hw_shape = self.patch_embed(pixel_values)
206
+ y = self.create_ann_token(annotation)
207
+ batch_size, num_tokens, channels = y.shape
208
+
209
+ if mask is not None:
210
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
211
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
212
+ y = y * (1.0 - weight) + mask_tokens * weight
213
+
214
+ if self.merge_stage == 0:
215
+ x = (x + y) * 0.5
216
+ else:
217
+ x = x.reshape(batch_size, *hw_shape, channels)
218
+ y = y.reshape(batch_size, *hw_shape, channels)
219
+ x = torch.cat((x, y), dim=2)
220
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
221
+ x = x.reshape(batch_size, -1, channels)
222
+
223
+ x = self._pos_embedding(x, hw_shape, self.pos_embed)
224
+
225
+ all_hidden_states = () if output_hidden_states else None
226
+ feature_maps = []
227
+ merge_idx = self.merge_stage - 1
228
+
229
+ for i, layer in enumerate(self.layers):
230
+ x = layer(x)
231
+
232
+ if i == merge_idx:
233
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
234
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
235
+ x = x.reshape(batch_size, -1, x.shape[-1])
236
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
237
+
238
+ if self.use_attn:
239
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
240
+ if i <= len(attention_blocks) - 1:
241
+ attn_out, _ = attention_blocks[i](x, x, x)
242
+ x = x + attn_out
243
+ if i == len(attention_blocks) - 1:
244
+ x = self.norm_attn(x)
245
+
246
+ if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm:
247
+ x = self.norm(x)
248
+
249
+ if output_hidden_states:
250
+ all_hidden_states = all_hidden_states + (x,)
251
+
252
+ if i in self.out_indices:
253
+ out = x
254
+ out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous()
255
+ if self.output_cls_token:
256
+ out = [out, x[:, 0]]
257
+ feature_maps.append(out)
258
+
259
+ if not return_dict:
260
+ return tuple(feature_maps)
261
+
262
+ return BaseModelOutput(
263
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
264
+ hidden_states=all_hidden_states,
265
+ )
skysensepp-swinv2-msl-hr/modeling_utils.py ADDED
@@ -0,0 +1,557 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense: Pure PyTorch + HuggingFace Transformers implementation.
2
+
3
+ Shared utility modules used across SkySense model implementations.
4
+ """
5
+
6
+ import math
7
+ from typing import Optional, Tuple
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+
14
+ def to_2tuple(x):
15
+ """Convert to a 2-tuple."""
16
+ if isinstance(x, (list, tuple)):
17
+ return tuple(x)
18
+ return (x, x)
19
+
20
+
21
+ class DropPath(nn.Module):
22
+ """Drop paths (stochastic depth) per sample.
23
+
24
+ Args:
25
+ drop_prob (float): Probability of dropping a path. Default: 0.0.
26
+ """
27
+
28
+ def __init__(self, drop_prob: float = 0.0):
29
+ super().__init__()
30
+ self.drop_prob = drop_prob
31
+
32
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
33
+ if self.drop_prob == 0.0 or not self.training:
34
+ return x
35
+ keep_prob = 1 - self.drop_prob
36
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1)
37
+ random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
38
+ random_tensor = torch.floor(random_tensor + keep_prob)
39
+ output = x / keep_prob * random_tensor
40
+ return output
41
+
42
+
43
+ class PatchEmbed(nn.Module):
44
+ """Image to Patch Embedding using Conv2d.
45
+
46
+ Args:
47
+ in_channels (int): Number of input channels. Default: 3.
48
+ embed_dims (int): Embedding dimension. Default: 96.
49
+ kernel_size (int): Kernel size of the projection. Default: 4.
50
+ stride (int): Stride of the projection. Default: 4.
51
+ padding (int): Padding of the projection. Default: 0.
52
+ norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
53
+ input_size (int or tuple or None): Input resolution for calculating output size.
54
+ """
55
+
56
+ def __init__(
57
+ self,
58
+ in_channels: int = 3,
59
+ embed_dims: int = 96,
60
+ kernel_size: int = 4,
61
+ stride: int = 4,
62
+ padding: int = 0,
63
+ norm_layer: Optional[type] = nn.LayerNorm,
64
+ input_size: Optional[int] = None,
65
+ ):
66
+ super().__init__()
67
+ self.projection = nn.Conv2d(
68
+ in_channels, embed_dims,
69
+ kernel_size=kernel_size, stride=stride, padding=padding,
70
+ )
71
+ self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
72
+
73
+ # Compute init output size if input_size is given
74
+ if input_size is not None:
75
+ input_size = to_2tuple(input_size)
76
+ self.init_out_size = (
77
+ (input_size[0] - kernel_size + 2 * padding) // stride + 1,
78
+ (input_size[1] - kernel_size + 2 * padding) // stride + 1,
79
+ )
80
+ else:
81
+ self.init_out_size = None
82
+
83
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
84
+ x = self.projection(x) # (B, C, H, W)
85
+ out_size = (x.shape[2], x.shape[3])
86
+ x = x.flatten(2).transpose(1, 2) # (B, H*W, C)
87
+ x = self.norm(x)
88
+ return x, out_size
89
+
90
+
91
+ class FFN(nn.Module):
92
+ """Feed-Forward Network.
93
+
94
+ Args:
95
+ embed_dims (int): Input dimension.
96
+ feedforward_channels (int): Hidden dimension.
97
+ num_fcs (int): Number of FC layers. Default: 2.
98
+ ffn_drop (float): Dropout rate. Default: 0.0.
99
+ drop_path (float): Drop path rate. Default: 0.0.
100
+ act_layer (nn.Module): Activation layer class. Default: nn.GELU.
101
+ add_identity (bool): Whether to add identity connection. Default: True.
102
+ """
103
+
104
+ def __init__(
105
+ self,
106
+ embed_dims: int,
107
+ feedforward_channels: int,
108
+ num_fcs: int = 2,
109
+ ffn_drop: float = 0.0,
110
+ drop_path: float = 0.0,
111
+ act_layer: type = nn.GELU,
112
+ add_identity: bool = True,
113
+ ):
114
+ super().__init__()
115
+ assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
116
+ self.embed_dims = embed_dims
117
+ self.feedforward_channels = feedforward_channels
118
+ self.add_identity = add_identity
119
+
120
+ layers = []
121
+ in_channels = embed_dims
122
+ for i in range(num_fcs - 1):
123
+ layers.append(nn.Linear(in_channels, feedforward_channels))
124
+ layers.append(act_layer())
125
+ layers.append(nn.Dropout(ffn_drop))
126
+ in_channels = feedforward_channels
127
+ layers.append(nn.Linear(feedforward_channels, embed_dims))
128
+ layers.append(nn.Dropout(ffn_drop))
129
+ self.layers = nn.Sequential(*layers)
130
+
131
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
132
+
133
+ def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
134
+ out = self.layers(x)
135
+ out = self.drop_path(out)
136
+ if self.add_identity:
137
+ if identity is None:
138
+ identity = x
139
+ out = out + identity
140
+ return out
141
+
142
+
143
+ class WindowMSAV2(nn.Module):
144
+ """Window-based Multi-head Self-Attention for Swin Transformer V2.
145
+
146
+ Uses cosine attention and log-spaced continuous position bias (log-CPB).
147
+
148
+ Args:
149
+ embed_dims (int): Number of input channels.
150
+ num_heads (int): Number of attention heads.
151
+ window_size (tuple[int]): Window size (Wh, Ww).
152
+ pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
153
+ qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
154
+ attn_drop (float): Attention dropout rate. Default: 0.0.
155
+ proj_drop (float): Output projection dropout rate. Default: 0.0.
156
+ """
157
+
158
+ def __init__(
159
+ self,
160
+ embed_dims: int,
161
+ num_heads: int,
162
+ window_size: Tuple[int, int],
163
+ pretrained_window_size: Tuple[int, int] = (0, 0),
164
+ qkv_bias: bool = True,
165
+ attn_drop: float = 0.0,
166
+ proj_drop: float = 0.0,
167
+ ):
168
+ super().__init__()
169
+ self.embed_dims = embed_dims
170
+ self.num_heads = num_heads
171
+ self.window_size = window_size
172
+ self.pretrained_window_size = pretrained_window_size
173
+
174
+ self.logit_scale = nn.Parameter(
175
+ torch.log(10 * torch.ones((num_heads, 1, 1))))
176
+
177
+ # MLP for continuous relative position bias (log-CPB)
178
+ self.cpb_mlp = nn.Sequential(
179
+ nn.Linear(2, 512, bias=True),
180
+ nn.ReLU(inplace=True),
181
+ nn.Linear(512, num_heads, bias=False),
182
+ )
183
+
184
+ # Build relative coords table
185
+ self._build_relative_coords_table()
186
+ # Build relative position index
187
+ self._build_relative_position_index()
188
+
189
+ self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
190
+ if qkv_bias:
191
+ self.q_bias = nn.Parameter(torch.zeros(embed_dims))
192
+ self.v_bias = nn.Parameter(torch.zeros(embed_dims))
193
+ else:
194
+ self.q_bias = None
195
+ self.v_bias = None
196
+
197
+ self.attn_drop = nn.Dropout(attn_drop)
198
+ self.proj = nn.Linear(embed_dims, embed_dims)
199
+ self.proj_drop = nn.Dropout(proj_drop)
200
+ self.softmax = nn.Softmax(dim=-1)
201
+
202
+ def _build_relative_coords_table(self):
203
+ """Build the relative coordinates table for log-CPB."""
204
+ Wh, Ww = self.window_size
205
+ # Table of relative coordinates
206
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
207
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
208
+ coords_table = torch.stack(
209
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
210
+ ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2)
211
+
212
+ # Normalize to [-1, 1] and apply log-scale
213
+ if self.pretrained_window_size[0] > 0:
214
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
215
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
216
+ else:
217
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
218
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
219
+ coords_table *= 8 # normalize to -8, 8
220
+ coords_table = (
221
+ torch.sign(coords_table)
222
+ * torch.log2(torch.abs(coords_table) + 1.0)
223
+ / math.log2(8)
224
+ )
225
+ self.register_buffer("relative_coords_table", coords_table)
226
+
227
+ def _build_relative_position_index(self):
228
+ """Build the pairwise relative position index for each window token."""
229
+ Wh, Ww = self.window_size
230
+ coords_h = torch.arange(Wh)
231
+ coords_w = torch.arange(Ww)
232
+ coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
233
+ coords_flatten = coords.view(2, -1)
234
+
235
+ relative_coords = (
236
+ coords_flatten[:, :, None] - coords_flatten[:, None, :]
237
+ ) # (2, Wh*Ww, Wh*Ww)
238
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
239
+ relative_coords[:, :, 0] += Wh - 1
240
+ relative_coords[:, :, 1] += Ww - 1
241
+ relative_coords[:, :, 0] *= 2 * Ww - 1
242
+ relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww)
243
+ self.register_buffer("relative_position_index", relative_position_index)
244
+
245
+ def _compute_position_bias(self, N):
246
+ """Compute relative position bias, supporting dynamic window sizes.
247
+
248
+ The log-CPB (Continuous Position Bias) MLP can generalize to any window
249
+ size by computing bias from normalized relative coordinates.
250
+ """
251
+ init_N = self.window_size[0] * self.window_size[1]
252
+ if N == init_N:
253
+ # Use pre-built tables
254
+ relative_position_bias_table = self.cpb_mlp(
255
+ self.relative_coords_table
256
+ ).view(-1, self.num_heads)
257
+ relative_position_bias = relative_position_bias_table[
258
+ self.relative_position_index.view(-1)
259
+ ].view(N, N, -1)
260
+ else:
261
+ # Dynamic: compute for actual window size on-the-fly
262
+ Wh = Ww = int(math.sqrt(N))
263
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
264
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
265
+ coords_table = torch.stack(
266
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
267
+ ).flatten(1).transpose(0, 1).unsqueeze(0)
268
+ if self.pretrained_window_size[0] > 0:
269
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
270
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
271
+ else:
272
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
273
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
274
+ coords_table *= 8
275
+ coords_table = (
276
+ torch.sign(coords_table)
277
+ * torch.log2(torch.abs(coords_table) + 1.0)
278
+ / math.log2(8)
279
+ )
280
+ # Build position index for actual window size
281
+ ch = torch.arange(Wh, device=self.logit_scale.device)
282
+ cw = torch.arange(Ww, device=self.logit_scale.device)
283
+ coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
284
+ coords_flat = coords.view(2, -1)
285
+ rel = coords_flat[:, :, None] - coords_flat[:, None, :]
286
+ rel = rel.permute(1, 2, 0).contiguous()
287
+ rel[:, :, 0] += Wh - 1
288
+ rel[:, :, 1] += Ww - 1
289
+ rel[:, :, 0] *= 2 * Ww - 1
290
+ pos_index = rel.sum(-1)
291
+
292
+ bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
293
+ relative_position_bias = bias_table[
294
+ pos_index.view(-1)
295
+ ].view(N, N, -1)
296
+
297
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
298
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
299
+ return relative_position_bias
300
+
301
+ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
302
+ """
303
+ Args:
304
+ x: (num_windows*B, N, C) where N = Wh*Ww
305
+ mask: (num_windows, N, N) or None
306
+ """
307
+ B_, N, C = x.shape
308
+
309
+ # Compute QKV with bias
310
+ if self.q_bias is not None:
311
+ qkv_bias = torch.cat(
312
+ (self.q_bias,
313
+ torch.zeros_like(self.v_bias, requires_grad=False),
314
+ self.v_bias))
315
+ qkv = F.linear(x, self.qkv.weight, qkv_bias)
316
+ else:
317
+ qkv = self.qkv(x)
318
+
319
+ qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
320
+ qkv = qkv.permute(2, 0, 3, 1, 4)
321
+ q, k, v = qkv.unbind(0)
322
+
323
+ # Cosine attention
324
+ attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
325
+ logit_scale = torch.clamp(
326
+ self.logit_scale, max=math.log(1.0 / 0.01)
327
+ ).exp()
328
+ attn = attn * logit_scale
329
+
330
+ # Log-CPB relative position bias (supports dynamic window sizes)
331
+ relative_position_bias = self._compute_position_bias(N)
332
+ attn = attn + relative_position_bias.unsqueeze(0)
333
+
334
+ if mask is not None:
335
+ nW = mask.shape[0]
336
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
337
+ attn = attn + mask.unsqueeze(1).unsqueeze(0)
338
+ attn = attn.view(-1, self.num_heads, N, N)
339
+
340
+ attn = self.softmax(attn)
341
+ attn = self.attn_drop(attn)
342
+
343
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
344
+ x = self.proj(x)
345
+ x = self.proj_drop(x)
346
+ return x
347
+
348
+
349
+ class ShiftWindowMSA(nn.Module):
350
+ """Shifted Window Multi-head Self-Attention.
351
+
352
+ Args:
353
+ embed_dims (int): Number of input channels.
354
+ num_heads (int): Number of attention heads.
355
+ window_size (int): Window size.
356
+ shift_size (int): Shift size for SW-MSA. Default: 0.
357
+ attn_drop (float): Attention dropout rate. Default: 0.0.
358
+ proj_drop (float): Projection dropout rate. Default: 0.0.
359
+ drop_path (float): Drop path rate. Default: 0.0.
360
+ pad_small_map (bool): Pad small feature maps to window size. Default: False.
361
+ pretrained_window_size (int): Pretrained window size. Default: 0.
362
+ """
363
+
364
+ def __init__(
365
+ self,
366
+ embed_dims: int,
367
+ num_heads: int,
368
+ window_size: int,
369
+ shift_size: int = 0,
370
+ attn_drop: float = 0.0,
371
+ proj_drop: float = 0.0,
372
+ drop_path: float = 0.0,
373
+ pad_small_map: bool = False,
374
+ pretrained_window_size: int = 0,
375
+ ):
376
+ super().__init__()
377
+ self.window_size = window_size
378
+ self.shift_size = shift_size
379
+ self.pad_small_map = pad_small_map
380
+
381
+ self.w_msa = WindowMSAV2(
382
+ embed_dims=embed_dims,
383
+ num_heads=num_heads,
384
+ window_size=to_2tuple(window_size),
385
+ pretrained_window_size=to_2tuple(pretrained_window_size),
386
+ attn_drop=attn_drop,
387
+ proj_drop=proj_drop,
388
+ )
389
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
390
+
391
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
392
+ B, L, C = x.shape
393
+ H, W = hw_shape
394
+ assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
395
+
396
+ x = x.view(B, H, W, C)
397
+
398
+ window_size = self.window_size
399
+ shift_size = self.shift_size
400
+
401
+ # Pad or shrink window
402
+ if self.pad_small_map:
403
+ pad_r = (window_size - W % window_size) % window_size
404
+ pad_b = (window_size - H % window_size) % window_size
405
+ x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
406
+ _, Hp, Wp, _ = x.shape
407
+ else:
408
+ Hp, Wp = H, W
409
+ if window_size > Hp:
410
+ window_size = Hp
411
+ shift_size = 0
412
+ if window_size > Wp:
413
+ window_size = Wp
414
+ shift_size = 0
415
+
416
+ # Compute attention mask for SW-MSA
417
+ attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
418
+
419
+ # Cyclic shift
420
+ if shift_size > 0:
421
+ x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
422
+
423
+ # Partition windows
424
+ x_windows = self._window_partition(x, window_size)
425
+ # (num_windows*B, window_size*window_size, C)
426
+
427
+ # W-MSA/SW-MSA
428
+ attn_windows = self.w_msa(x_windows, mask=attn_mask)
429
+
430
+ # Merge windows
431
+ x = self._window_reverse(attn_windows, window_size, Hp, Wp)
432
+
433
+ # Reverse cyclic shift
434
+ if shift_size > 0:
435
+ x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
436
+
437
+ if self.pad_small_map and (pad_r > 0 or pad_b > 0):
438
+ x = x[:, :H, :W, :].contiguous()
439
+
440
+ x = x.view(B, H * W, C)
441
+ x = self.drop_path(x)
442
+ return x
443
+
444
+ @staticmethod
445
+ def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
446
+ """Partition into non-overlapping windows."""
447
+ B, H, W, C = x.shape
448
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
449
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
450
+ windows = windows.view(-1, window_size * window_size, C)
451
+ return windows
452
+
453
+ @staticmethod
454
+ def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
455
+ """Reverse window partition."""
456
+ B_nW = windows.shape[0]
457
+ nH = H // window_size
458
+ nW = W // window_size
459
+ B = B_nW // (nH * nW)
460
+ x = windows.view(B, nH, nW, window_size, window_size, -1)
461
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
462
+ x = x.view(B, H, W, -1)
463
+ return x
464
+
465
+ @staticmethod
466
+ def _compute_attn_mask(H, W, window_size, shift_size, device):
467
+ """Compute attention mask for shifted window attention."""
468
+ if shift_size <= 0:
469
+ return None
470
+ img_mask = torch.zeros((1, H, W, 1), device=device)
471
+ h_slices = (
472
+ slice(0, -window_size),
473
+ slice(-window_size, -shift_size),
474
+ slice(-shift_size, None),
475
+ )
476
+ w_slices = (
477
+ slice(0, -window_size),
478
+ slice(-window_size, -shift_size),
479
+ slice(-shift_size, None),
480
+ )
481
+ cnt = 0
482
+ for h in h_slices:
483
+ for w in w_slices:
484
+ img_mask[:, h, w, :] = cnt
485
+ cnt += 1
486
+
487
+ # Partition mask
488
+ mask_windows = img_mask.view(
489
+ 1, H // window_size, window_size, W // window_size, window_size, 1
490
+ )
491
+ mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
492
+ mask_windows = mask_windows.view(-1, window_size * window_size)
493
+
494
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
495
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
496
+ attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
497
+ return attn_mask
498
+
499
+
500
+ class PatchMerging(nn.Module):
501
+ """Patch Merging Layer for downsampling (2x).
502
+
503
+ Args:
504
+ in_channels (int): Input channels.
505
+ out_channels (int): Output channels.
506
+ norm_layer (type): Normalization layer. Default: nn.LayerNorm.
507
+ is_post_norm (bool): Apply norm after linear. Default: True.
508
+ """
509
+
510
+ def __init__(
511
+ self,
512
+ in_channels: int,
513
+ out_channels: int,
514
+ norm_layer: type = nn.LayerNorm,
515
+ is_post_norm: bool = True,
516
+ ):
517
+ super().__init__()
518
+ self.in_channels = in_channels
519
+ self.out_channels = out_channels
520
+ self.is_post_norm = is_post_norm
521
+ self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
522
+ if is_post_norm:
523
+ self.norm = norm_layer(out_channels)
524
+ else:
525
+ self.norm = norm_layer(4 * in_channels)
526
+
527
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
528
+ B, L, C = x.shape
529
+ H, W = hw_shape
530
+ assert L == H * W
531
+
532
+ x = x.view(B, H, W, C)
533
+
534
+ # Pad if needed
535
+ pad_h = H % 2
536
+ pad_w = W % 2
537
+ if pad_h or pad_w:
538
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
539
+
540
+ x0 = x[:, 0::2, 0::2, :]
541
+ x1 = x[:, 1::2, 0::2, :]
542
+ x2 = x[:, 0::2, 1::2, :]
543
+ x3 = x[:, 1::2, 1::2, :]
544
+ x = torch.cat([x0, x1, x2, x3], dim=-1)
545
+
546
+ out_h = (H + pad_h) // 2
547
+ out_w = (W + pad_w) // 2
548
+ x = x.view(B, out_h * out_w, 4 * C)
549
+
550
+ if self.is_post_norm:
551
+ x = self.reduction(x)
552
+ x = self.norm(x)
553
+ else:
554
+ x = self.norm(x)
555
+ x = self.reduction(x)
556
+
557
+ return x, (out_h, out_w)
skysensepp-swinv2-msl-hr/pipeline_skysensepp.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Custom HuggingFace pipeline for SkySense++ MSL feature extraction."""
2
+
3
+ from typing import Any, Dict, Optional, Union
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline):
11
+ """Pipeline for SkySense++ MSL backbones.
12
+
13
+ Expects image tensors plus semantic annotation maps (class indices).
14
+ """
15
+
16
+ def _sanitize_parameters(
17
+ self,
18
+ annotation=None,
19
+ mask=None,
20
+ output_hidden_states=None,
21
+ **kwargs,
22
+ ):
23
+ preprocess_params = {}
24
+ forward_params = {}
25
+ postprocess_params = {}
26
+
27
+ if annotation is not None:
28
+ preprocess_params["annotation"] = annotation
29
+ if mask is not None:
30
+ forward_params["mask"] = mask
31
+ if output_hidden_states is not None:
32
+ forward_params["output_hidden_states"] = output_hidden_states
33
+
34
+ return preprocess_params, forward_params, postprocess_params
35
+
36
+ def preprocess(
37
+ self,
38
+ pixel_values: Any,
39
+ annotation: Optional[Any] = None,
40
+ **kwargs,
41
+ ) -> Dict[str, torch.Tensor]:
42
+ if isinstance(pixel_values, dict):
43
+ annotation = pixel_values.get("annotation", annotation)
44
+ pixel_values = pixel_values.get("pixel_values", pixel_values)
45
+
46
+ if isinstance(pixel_values, np.ndarray):
47
+ pixel_values = torch.from_numpy(pixel_values).float()
48
+ elif not isinstance(pixel_values, torch.Tensor):
49
+ raise TypeError(
50
+ f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}"
51
+ )
52
+
53
+ if annotation is None:
54
+ raise ValueError("SkySense++ MSL models require an `annotation` semantic map.")
55
+
56
+ if isinstance(annotation, np.ndarray):
57
+ annotation = torch.from_numpy(annotation).long()
58
+ elif not isinstance(annotation, torch.Tensor):
59
+ raise TypeError(
60
+ f"Expected tensor or ndarray for annotation, got {type(annotation)}"
61
+ )
62
+
63
+ if pixel_values.ndim == 3:
64
+ pixel_values = pixel_values.unsqueeze(0)
65
+ if annotation.ndim == 2:
66
+ annotation = annotation.unsqueeze(0)
67
+
68
+ return {"pixel_values": pixel_values, "annotation": annotation}
69
+
70
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
71
+ with torch.no_grad():
72
+ outputs = self.model(
73
+ pixel_values=model_inputs["pixel_values"],
74
+ annotation=model_inputs["annotation"],
75
+ mask=kwargs.get("mask"),
76
+ output_hidden_states=kwargs.get("output_hidden_states", False),
77
+ return_dict=True,
78
+ )
79
+ return {"outputs": outputs}
80
+
81
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
82
+ outputs = model_outputs["outputs"]
83
+ result = {"last_hidden_state": outputs.last_hidden_state}
84
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
85
+ result["hidden_states"] = outputs.hidden_states
86
+ return result
skysensepp-swinv2-msl-hr/pipeline_skysensepp_fusion.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Optional pipeline for SkySense++ fusion neck."""
2
+
3
+ from typing import Any, Dict
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusFusionNeckPipeline(Pipeline):
11
+ """Pipeline for the optional SkySense++ fusion neck module.
12
+
13
+ Expects concatenated multi-modal tokens per spatial location:
14
+ ``(batch, num_modalities, input_dims)``.
15
+ """
16
+
17
+ def _sanitize_parameters(self, output_hidden_states=None, **kwargs):
18
+ preprocess_params = {}
19
+ forward_params = {}
20
+ postprocess_params = {}
21
+ if output_hidden_states is not None:
22
+ forward_params["output_hidden_states"] = output_hidden_states
23
+ return preprocess_params, forward_params, postprocess_params
24
+
25
+ def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]:
26
+ if isinstance(hidden_states, dict):
27
+ hidden_states = hidden_states["hidden_states"]
28
+
29
+ if isinstance(hidden_states, np.ndarray):
30
+ hidden_states = torch.from_numpy(hidden_states).float()
31
+ elif not isinstance(hidden_states, torch.Tensor):
32
+ raise TypeError(
33
+ f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}"
34
+ )
35
+ if hidden_states.ndim == 2:
36
+ hidden_states = hidden_states.unsqueeze(0)
37
+ return {"hidden_states": hidden_states}
38
+
39
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
40
+ with torch.no_grad():
41
+ outputs = self.model(
42
+ hidden_states=model_inputs["hidden_states"],
43
+ output_hidden_states=kwargs.get("output_hidden_states", False),
44
+ return_dict=True,
45
+ )
46
+ return {"outputs": outputs}
47
+
48
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
49
+ outputs = model_outputs["outputs"]
50
+ result = {"pooler_output": outputs.pooler_output}
51
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
52
+ result["hidden_states"] = outputs.hidden_states
53
+ return result
skysensepp-vit-msl-s1/__init__.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++: Multi-Modal Remote Sensing Foundation Model (HuggingFace)."""
2
+
3
+ from .configuration_skysensepp import (
4
+ SkySensePlusPlusSwinV2MSLConfig,
5
+ SkySensePlusPlusViTMSLConfig,
6
+ )
7
+ from .modeling_skysensepp_swinv2_msl import (
8
+ SkySensePlusPlusSwinV2MSLModel,
9
+ SkySensePlusPlusSwinV2MSLPreTrainedModel,
10
+ )
11
+ from .modeling_skysensepp_vit_msl import (
12
+ SkySensePlusPlusViTMSLModel,
13
+ SkySensePlusPlusViTMSLPreTrainedModel,
14
+ )
15
+ from .pipeline_skysensepp import SkySensePlusPlusMSLFeatureExtractionPipeline
16
+
17
+ __all__ = [
18
+ "SkySensePlusPlusSwinV2MSLConfig",
19
+ "SkySensePlusPlusViTMSLConfig",
20
+ "SkySensePlusPlusSwinV2MSLModel",
21
+ "SkySensePlusPlusSwinV2MSLPreTrainedModel",
22
+ "SkySensePlusPlusViTMSLModel",
23
+ "SkySensePlusPlusViTMSLPreTrainedModel",
24
+ "SkySensePlusPlusMSLFeatureExtractionPipeline",
25
+ ]
skysensepp-vit-msl-s1/config.json ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "return_dict": true,
3
+ "output_hidden_states": false,
4
+ "dtype": "float32",
5
+ "chunk_size_feed_forward": 0,
6
+ "is_encoder_decoder": false,
7
+ "architectures": [
8
+ "SkySensePlusPlusViTMSLModel"
9
+ ],
10
+ "id2label": {
11
+ "0": "LABEL_0",
12
+ "1": "LABEL_1"
13
+ },
14
+ "label2id": {
15
+ "LABEL_0": 0,
16
+ "LABEL_1": 1
17
+ },
18
+ "problem_type": null,
19
+ "_name_or_path": "",
20
+ "transformers_version": "5.0.0",
21
+ "img_size": 16,
22
+ "patch_size": 4,
23
+ "in_channels": 2,
24
+ "embed_dims": 1024,
25
+ "num_layers": 24,
26
+ "num_heads": 16,
27
+ "mlp_ratio": 4,
28
+ "out_indices": [
29
+ 5,
30
+ 11,
31
+ 17,
32
+ 23
33
+ ],
34
+ "qkv_bias": true,
35
+ "drop_rate": 0.0,
36
+ "attn_drop_rate": 0.0,
37
+ "drop_path_rate": 0.3,
38
+ "with_cls_token": false,
39
+ "output_cls_token": false,
40
+ "patch_norm": false,
41
+ "final_norm": false,
42
+ "with_cp": false,
43
+ "vocabulary_size": 64,
44
+ "num_vocabulary_tokens": 65,
45
+ "merge_stage": 4,
46
+ "use_attn": false,
47
+ "modality": "s1",
48
+ "model_type": "skysensepp_vit_msl",
49
+ "output_attentions": false,
50
+ "auto_map": {
51
+ "AutoConfig": "configuration_skysensepp.SkySensePlusPlusViTMSLConfig",
52
+ "AutoModel": "modeling_skysensepp_vit_msl.SkySensePlusPlusViTMSLModel"
53
+ },
54
+ "custom_pipelines": {
55
+ "skysensepp-feature-extraction": {
56
+ "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline",
57
+ "pt": [
58
+ "AutoModel"
59
+ ]
60
+ },
61
+ "image-feature-extraction": {
62
+ "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline",
63
+ "pt": [
64
+ "AutoModel"
65
+ ]
66
+ }
67
+ }
68
+ }
skysensepp-vit-msl-s1/configuration_skysensepp.py ADDED
@@ -0,0 +1,124 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration classes for SkySense++ MSL backbones."""
2
+
3
+ from transformers import PretrainedConfig
4
+
5
+
6
+ class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig):
7
+ """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical)."""
8
+
9
+ model_type = "skysensepp_swinv2_msl"
10
+
11
+ arch_zoo = {
12
+ "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
13
+ "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
14
+ "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0},
15
+ "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0},
16
+ "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6},
17
+ "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6},
18
+ }
19
+
20
+ def __init__(
21
+ self,
22
+ arch="huge",
23
+ img_size=224,
24
+ patch_size=4,
25
+ in_channels=3,
26
+ window_size=8,
27
+ drop_rate=0.0,
28
+ drop_path_rate=0.2,
29
+ out_indices=(0, 1, 2, 3),
30
+ use_abs_pos_embed=False,
31
+ with_cp=False,
32
+ pad_small_map=False,
33
+ pretrained_window_sizes=(0, 0, 0, 0),
34
+ is_post_norm_downsample=True,
35
+ vocabulary_size=64,
36
+ merge_stage=2,
37
+ use_attn=True,
38
+ **kwargs,
39
+ ):
40
+ super().__init__(**kwargs)
41
+
42
+ arch = arch.lower()
43
+ if arch not in self.arch_zoo:
44
+ raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
45
+ arch_settings = self.arch_zoo[arch]
46
+
47
+ self.arch = arch
48
+ self.embed_dims = arch_settings["embed_dims"]
49
+ self.depths = arch_settings["depths"]
50
+ self.num_heads = arch_settings["num_heads"]
51
+ self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"]
52
+
53
+ self.img_size = img_size
54
+ self.patch_size = patch_size
55
+ self.in_channels = in_channels
56
+ self.window_size = window_size
57
+ self.drop_rate = drop_rate
58
+ self.drop_path_rate = drop_path_rate
59
+ self.out_indices = list(out_indices)
60
+ self.use_abs_pos_embed = use_abs_pos_embed
61
+ self.with_cp = with_cp
62
+ self.pad_small_map = pad_small_map
63
+ self.pretrained_window_sizes = list(pretrained_window_sizes)
64
+ self.is_post_norm_downsample = is_post_norm_downsample
65
+
66
+ self.vocabulary_size = vocabulary_size
67
+ self.num_vocabulary_tokens = vocabulary_size + 1
68
+ self.merge_stage = merge_stage
69
+ self.use_attn = use_attn
70
+
71
+
72
+ class SkySensePlusPlusViTMSLConfig(PretrainedConfig):
73
+ """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1)."""
74
+
75
+ model_type = "skysensepp_vit_msl"
76
+
77
+ def __init__(
78
+ self,
79
+ img_size=16,
80
+ patch_size=4,
81
+ in_channels=10,
82
+ embed_dims=1024,
83
+ num_layers=24,
84
+ num_heads=16,
85
+ mlp_ratio=4,
86
+ out_indices=(5, 11, 17, 23),
87
+ qkv_bias=True,
88
+ drop_rate=0.0,
89
+ attn_drop_rate=0.0,
90
+ drop_path_rate=0.3,
91
+ with_cls_token=False,
92
+ output_cls_token=False,
93
+ patch_norm=False,
94
+ final_norm=False,
95
+ with_cp=False,
96
+ vocabulary_size=64,
97
+ merge_stage=4,
98
+ use_attn=False,
99
+ modality="s2",
100
+ **kwargs,
101
+ ):
102
+ super().__init__(**kwargs)
103
+ self.img_size = img_size
104
+ self.patch_size = patch_size
105
+ self.in_channels = in_channels
106
+ self.embed_dims = embed_dims
107
+ self.num_layers = num_layers
108
+ self.num_heads = num_heads
109
+ self.mlp_ratio = mlp_ratio
110
+ self.out_indices = list(out_indices)
111
+ self.qkv_bias = qkv_bias
112
+ self.drop_rate = drop_rate
113
+ self.attn_drop_rate = attn_drop_rate
114
+ self.drop_path_rate = drop_path_rate
115
+ self.with_cls_token = with_cls_token
116
+ self.output_cls_token = output_cls_token
117
+ self.patch_norm = patch_norm
118
+ self.final_norm = final_norm
119
+ self.with_cp = with_cp
120
+ self.vocabulary_size = vocabulary_size
121
+ self.num_vocabulary_tokens = vocabulary_size + 1
122
+ self.merge_stage = merge_stage
123
+ self.use_attn = use_attn
124
+ self.modality = modality
skysensepp-vit-msl-s1/conversion_manifest.json ADDED
@@ -0,0 +1,305 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release_s1.pth",
3
+ "modality": "s1",
4
+ "model_class": "SkySensePlusPlusViTMSLModel",
5
+ "num_tensors": 295,
6
+ "missing_keys": [],
7
+ "unexpected_keys": [],
8
+ "tensor_names": [
9
+ "cls_token",
10
+ "layers.0.attn.in_proj_bias",
11
+ "layers.0.attn.in_proj_weight",
12
+ "layers.0.attn.out_proj.bias",
13
+ "layers.0.attn.out_proj.weight",
14
+ "layers.0.ffn.layers.0.bias",
15
+ "layers.0.ffn.layers.0.weight",
16
+ "layers.0.ffn.layers.3.bias",
17
+ "layers.0.ffn.layers.3.weight",
18
+ "layers.0.norm1.bias",
19
+ "layers.0.norm1.weight",
20
+ "layers.0.norm2.bias",
21
+ "layers.0.norm2.weight",
22
+ "layers.1.attn.in_proj_bias",
23
+ "layers.1.attn.in_proj_weight",
24
+ "layers.1.attn.out_proj.bias",
25
+ "layers.1.attn.out_proj.weight",
26
+ "layers.1.ffn.layers.0.bias",
27
+ "layers.1.ffn.layers.0.weight",
28
+ "layers.1.ffn.layers.3.bias",
29
+ "layers.1.ffn.layers.3.weight",
30
+ "layers.1.norm1.bias",
31
+ "layers.1.norm1.weight",
32
+ "layers.1.norm2.bias",
33
+ "layers.1.norm2.weight",
34
+ "layers.10.attn.in_proj_bias",
35
+ "layers.10.attn.in_proj_weight",
36
+ "layers.10.attn.out_proj.bias",
37
+ "layers.10.attn.out_proj.weight",
38
+ "layers.10.ffn.layers.0.bias",
39
+ "layers.10.ffn.layers.0.weight",
40
+ "layers.10.ffn.layers.3.bias",
41
+ "layers.10.ffn.layers.3.weight",
42
+ "layers.10.norm1.bias",
43
+ "layers.10.norm1.weight",
44
+ "layers.10.norm2.bias",
45
+ "layers.10.norm2.weight",
46
+ "layers.11.attn.in_proj_bias",
47
+ "layers.11.attn.in_proj_weight",
48
+ "layers.11.attn.out_proj.bias",
49
+ "layers.11.attn.out_proj.weight",
50
+ "layers.11.ffn.layers.0.bias",
51
+ "layers.11.ffn.layers.0.weight",
52
+ "layers.11.ffn.layers.3.bias",
53
+ "layers.11.ffn.layers.3.weight",
54
+ "layers.11.norm1.bias",
55
+ "layers.11.norm1.weight",
56
+ "layers.11.norm2.bias",
57
+ "layers.11.norm2.weight",
58
+ "layers.12.attn.in_proj_bias",
59
+ "layers.12.attn.in_proj_weight",
60
+ "layers.12.attn.out_proj.bias",
61
+ "layers.12.attn.out_proj.weight",
62
+ "layers.12.ffn.layers.0.bias",
63
+ "layers.12.ffn.layers.0.weight",
64
+ "layers.12.ffn.layers.3.bias",
65
+ "layers.12.ffn.layers.3.weight",
66
+ "layers.12.norm1.bias",
67
+ "layers.12.norm1.weight",
68
+ "layers.12.norm2.bias",
69
+ "layers.12.norm2.weight",
70
+ "layers.13.attn.in_proj_bias",
71
+ "layers.13.attn.in_proj_weight",
72
+ "layers.13.attn.out_proj.bias",
73
+ "layers.13.attn.out_proj.weight",
74
+ "layers.13.ffn.layers.0.bias",
75
+ "layers.13.ffn.layers.0.weight",
76
+ "layers.13.ffn.layers.3.bias",
77
+ "layers.13.ffn.layers.3.weight",
78
+ "layers.13.norm1.bias",
79
+ "layers.13.norm1.weight",
80
+ "layers.13.norm2.bias",
81
+ "layers.13.norm2.weight",
82
+ "layers.14.attn.in_proj_bias",
83
+ "layers.14.attn.in_proj_weight",
84
+ "layers.14.attn.out_proj.bias",
85
+ "layers.14.attn.out_proj.weight",
86
+ "layers.14.ffn.layers.0.bias",
87
+ "layers.14.ffn.layers.0.weight",
88
+ "layers.14.ffn.layers.3.bias",
89
+ "layers.14.ffn.layers.3.weight",
90
+ "layers.14.norm1.bias",
91
+ "layers.14.norm1.weight",
92
+ "layers.14.norm2.bias",
93
+ "layers.14.norm2.weight",
94
+ "layers.15.attn.in_proj_bias",
95
+ "layers.15.attn.in_proj_weight",
96
+ "layers.15.attn.out_proj.bias",
97
+ "layers.15.attn.out_proj.weight",
98
+ "layers.15.ffn.layers.0.bias",
99
+ "layers.15.ffn.layers.0.weight",
100
+ "layers.15.ffn.layers.3.bias",
101
+ "layers.15.ffn.layers.3.weight",
102
+ "layers.15.norm1.bias",
103
+ "layers.15.norm1.weight",
104
+ "layers.15.norm2.bias",
105
+ "layers.15.norm2.weight",
106
+ "layers.16.attn.in_proj_bias",
107
+ "layers.16.attn.in_proj_weight",
108
+ "layers.16.attn.out_proj.bias",
109
+ "layers.16.attn.out_proj.weight",
110
+ "layers.16.ffn.layers.0.bias",
111
+ "layers.16.ffn.layers.0.weight",
112
+ "layers.16.ffn.layers.3.bias",
113
+ "layers.16.ffn.layers.3.weight",
114
+ "layers.16.norm1.bias",
115
+ "layers.16.norm1.weight",
116
+ "layers.16.norm2.bias",
117
+ "layers.16.norm2.weight",
118
+ "layers.17.attn.in_proj_bias",
119
+ "layers.17.attn.in_proj_weight",
120
+ "layers.17.attn.out_proj.bias",
121
+ "layers.17.attn.out_proj.weight",
122
+ "layers.17.ffn.layers.0.bias",
123
+ "layers.17.ffn.layers.0.weight",
124
+ "layers.17.ffn.layers.3.bias",
125
+ "layers.17.ffn.layers.3.weight",
126
+ "layers.17.norm1.bias",
127
+ "layers.17.norm1.weight",
128
+ "layers.17.norm2.bias",
129
+ "layers.17.norm2.weight",
130
+ "layers.18.attn.in_proj_bias",
131
+ "layers.18.attn.in_proj_weight",
132
+ "layers.18.attn.out_proj.bias",
133
+ "layers.18.attn.out_proj.weight",
134
+ "layers.18.ffn.layers.0.bias",
135
+ "layers.18.ffn.layers.0.weight",
136
+ "layers.18.ffn.layers.3.bias",
137
+ "layers.18.ffn.layers.3.weight",
138
+ "layers.18.norm1.bias",
139
+ "layers.18.norm1.weight",
140
+ "layers.18.norm2.bias",
141
+ "layers.18.norm2.weight",
142
+ "layers.19.attn.in_proj_bias",
143
+ "layers.19.attn.in_proj_weight",
144
+ "layers.19.attn.out_proj.bias",
145
+ "layers.19.attn.out_proj.weight",
146
+ "layers.19.ffn.layers.0.bias",
147
+ "layers.19.ffn.layers.0.weight",
148
+ "layers.19.ffn.layers.3.bias",
149
+ "layers.19.ffn.layers.3.weight",
150
+ "layers.19.norm1.bias",
151
+ "layers.19.norm1.weight",
152
+ "layers.19.norm2.bias",
153
+ "layers.19.norm2.weight",
154
+ "layers.2.attn.in_proj_bias",
155
+ "layers.2.attn.in_proj_weight",
156
+ "layers.2.attn.out_proj.bias",
157
+ "layers.2.attn.out_proj.weight",
158
+ "layers.2.ffn.layers.0.bias",
159
+ "layers.2.ffn.layers.0.weight",
160
+ "layers.2.ffn.layers.3.bias",
161
+ "layers.2.ffn.layers.3.weight",
162
+ "layers.2.norm1.bias",
163
+ "layers.2.norm1.weight",
164
+ "layers.2.norm2.bias",
165
+ "layers.2.norm2.weight",
166
+ "layers.20.attn.in_proj_bias",
167
+ "layers.20.attn.in_proj_weight",
168
+ "layers.20.attn.out_proj.bias",
169
+ "layers.20.attn.out_proj.weight",
170
+ "layers.20.ffn.layers.0.bias",
171
+ "layers.20.ffn.layers.0.weight",
172
+ "layers.20.ffn.layers.3.bias",
173
+ "layers.20.ffn.layers.3.weight",
174
+ "layers.20.norm1.bias",
175
+ "layers.20.norm1.weight",
176
+ "layers.20.norm2.bias",
177
+ "layers.20.norm2.weight",
178
+ "layers.21.attn.in_proj_bias",
179
+ "layers.21.attn.in_proj_weight",
180
+ "layers.21.attn.out_proj.bias",
181
+ "layers.21.attn.out_proj.weight",
182
+ "layers.21.ffn.layers.0.bias",
183
+ "layers.21.ffn.layers.0.weight",
184
+ "layers.21.ffn.layers.3.bias",
185
+ "layers.21.ffn.layers.3.weight",
186
+ "layers.21.norm1.bias",
187
+ "layers.21.norm1.weight",
188
+ "layers.21.norm2.bias",
189
+ "layers.21.norm2.weight",
190
+ "layers.22.attn.in_proj_bias",
191
+ "layers.22.attn.in_proj_weight",
192
+ "layers.22.attn.out_proj.bias",
193
+ "layers.22.attn.out_proj.weight",
194
+ "layers.22.ffn.layers.0.bias",
195
+ "layers.22.ffn.layers.0.weight",
196
+ "layers.22.ffn.layers.3.bias",
197
+ "layers.22.ffn.layers.3.weight",
198
+ "layers.22.norm1.bias",
199
+ "layers.22.norm1.weight",
200
+ "layers.22.norm2.bias",
201
+ "layers.22.norm2.weight",
202
+ "layers.23.attn.in_proj_bias",
203
+ "layers.23.attn.in_proj_weight",
204
+ "layers.23.attn.out_proj.bias",
205
+ "layers.23.attn.out_proj.weight",
206
+ "layers.23.ffn.layers.0.bias",
207
+ "layers.23.ffn.layers.0.weight",
208
+ "layers.23.ffn.layers.3.bias",
209
+ "layers.23.ffn.layers.3.weight",
210
+ "layers.23.norm1.bias",
211
+ "layers.23.norm1.weight",
212
+ "layers.23.norm2.bias",
213
+ "layers.23.norm2.weight",
214
+ "layers.3.attn.in_proj_bias",
215
+ "layers.3.attn.in_proj_weight",
216
+ "layers.3.attn.out_proj.bias",
217
+ "layers.3.attn.out_proj.weight",
218
+ "layers.3.ffn.layers.0.bias",
219
+ "layers.3.ffn.layers.0.weight",
220
+ "layers.3.ffn.layers.3.bias",
221
+ "layers.3.ffn.layers.3.weight",
222
+ "layers.3.norm1.bias",
223
+ "layers.3.norm1.weight",
224
+ "layers.3.norm2.bias",
225
+ "layers.3.norm2.weight",
226
+ "layers.4.attn.in_proj_bias",
227
+ "layers.4.attn.in_proj_weight",
228
+ "layers.4.attn.out_proj.bias",
229
+ "layers.4.attn.out_proj.weight",
230
+ "layers.4.ffn.layers.0.bias",
231
+ "layers.4.ffn.layers.0.weight",
232
+ "layers.4.ffn.layers.3.bias",
233
+ "layers.4.ffn.layers.3.weight",
234
+ "layers.4.norm1.bias",
235
+ "layers.4.norm1.weight",
236
+ "layers.4.norm2.bias",
237
+ "layers.4.norm2.weight",
238
+ "layers.5.attn.in_proj_bias",
239
+ "layers.5.attn.in_proj_weight",
240
+ "layers.5.attn.out_proj.bias",
241
+ "layers.5.attn.out_proj.weight",
242
+ "layers.5.ffn.layers.0.bias",
243
+ "layers.5.ffn.layers.0.weight",
244
+ "layers.5.ffn.layers.3.bias",
245
+ "layers.5.ffn.layers.3.weight",
246
+ "layers.5.norm1.bias",
247
+ "layers.5.norm1.weight",
248
+ "layers.5.norm2.bias",
249
+ "layers.5.norm2.weight",
250
+ "layers.6.attn.in_proj_bias",
251
+ "layers.6.attn.in_proj_weight",
252
+ "layers.6.attn.out_proj.bias",
253
+ "layers.6.attn.out_proj.weight",
254
+ "layers.6.ffn.layers.0.bias",
255
+ "layers.6.ffn.layers.0.weight",
256
+ "layers.6.ffn.layers.3.bias",
257
+ "layers.6.ffn.layers.3.weight",
258
+ "layers.6.norm1.bias",
259
+ "layers.6.norm1.weight",
260
+ "layers.6.norm2.bias",
261
+ "layers.6.norm2.weight",
262
+ "layers.7.attn.in_proj_bias",
263
+ "layers.7.attn.in_proj_weight",
264
+ "layers.7.attn.out_proj.bias",
265
+ "layers.7.attn.out_proj.weight",
266
+ "layers.7.ffn.layers.0.bias",
267
+ "layers.7.ffn.layers.0.weight",
268
+ "layers.7.ffn.layers.3.bias",
269
+ "layers.7.ffn.layers.3.weight",
270
+ "layers.7.norm1.bias",
271
+ "layers.7.norm1.weight",
272
+ "layers.7.norm2.bias",
273
+ "layers.7.norm2.weight",
274
+ "layers.8.attn.in_proj_bias",
275
+ "layers.8.attn.in_proj_weight",
276
+ "layers.8.attn.out_proj.bias",
277
+ "layers.8.attn.out_proj.weight",
278
+ "layers.8.ffn.layers.0.bias",
279
+ "layers.8.ffn.layers.0.weight",
280
+ "layers.8.ffn.layers.3.bias",
281
+ "layers.8.ffn.layers.3.weight",
282
+ "layers.8.norm1.bias",
283
+ "layers.8.norm1.weight",
284
+ "layers.8.norm2.bias",
285
+ "layers.8.norm2.weight",
286
+ "layers.9.attn.in_proj_bias",
287
+ "layers.9.attn.in_proj_weight",
288
+ "layers.9.attn.out_proj.bias",
289
+ "layers.9.attn.out_proj.weight",
290
+ "layers.9.ffn.layers.0.bias",
291
+ "layers.9.ffn.layers.0.weight",
292
+ "layers.9.ffn.layers.3.bias",
293
+ "layers.9.ffn.layers.3.weight",
294
+ "layers.9.norm1.bias",
295
+ "layers.9.norm1.weight",
296
+ "layers.9.norm2.bias",
297
+ "layers.9.norm2.weight",
298
+ "mask_token",
299
+ "patch_embed.projection.bias",
300
+ "patch_embed.projection.weight",
301
+ "pos_embed",
302
+ "vocabulary_token",
303
+ "vocabulary_weight"
304
+ ]
305
+ }
skysensepp-vit-msl-s1/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9a9a658f9efd8f0ebb65d92fe075bec23c1b9edf1fd653b2d7457f2065acaf1f
3
+ size 1209741688
skysensepp-vit-msl-s1/modeling_skysensepp_swinv2_msl.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ from copy import deepcopy
4
+ from typing import Optional, Sequence, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig
14
+ from .modeling_utils import (
15
+ DropPath,
16
+ FFN,
17
+ PatchEmbed,
18
+ PatchMerging,
19
+ ShiftWindowMSA,
20
+ to_2tuple,
21
+ )
22
+
23
+
24
+ class SwinBlockV2(nn.Module):
25
+ def __init__(
26
+ self,
27
+ embed_dims: int,
28
+ num_heads: int,
29
+ window_size: int = 8,
30
+ shift: bool = False,
31
+ extra_norm: bool = False,
32
+ ffn_ratio: float = 4.0,
33
+ drop_path: float = 0.0,
34
+ pad_small_map: bool = False,
35
+ with_cp: bool = False,
36
+ pretrained_window_size: int = 0,
37
+ ):
38
+ super().__init__()
39
+ self.with_cp = with_cp
40
+ self.extra_norm = extra_norm
41
+ self.attn = ShiftWindowMSA(
42
+ embed_dims=embed_dims,
43
+ num_heads=num_heads,
44
+ window_size=window_size,
45
+ shift_size=window_size // 2 if shift else 0,
46
+ drop_path=drop_path,
47
+ pad_small_map=pad_small_map,
48
+ pretrained_window_size=pretrained_window_size,
49
+ )
50
+ self.norm1 = nn.LayerNorm(embed_dims)
51
+ self.ffn = FFN(
52
+ embed_dims=embed_dims,
53
+ feedforward_channels=int(embed_dims * ffn_ratio),
54
+ num_fcs=2,
55
+ drop_path=drop_path,
56
+ act_layer=nn.GELU,
57
+ add_identity=False,
58
+ )
59
+ self.norm2 = nn.LayerNorm(embed_dims)
60
+ if self.extra_norm:
61
+ self.norm3 = nn.LayerNorm(embed_dims)
62
+
63
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
64
+ def _inner_forward(x):
65
+ identity = x
66
+ x = self.attn(x, hw_shape)
67
+ x = self.norm1(x)
68
+ x = x + identity
69
+
70
+ identity = x
71
+ x = self.ffn(x)
72
+ x = self.norm2(x)
73
+ x = x + identity
74
+
75
+ if self.extra_norm:
76
+ x = self.norm3(x)
77
+ return x
78
+
79
+ if self.with_cp and x.requires_grad:
80
+ x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
81
+ else:
82
+ x = _inner_forward(x)
83
+ return x
84
+
85
+
86
+ class SwinBlockV2Sequence(nn.Module):
87
+ def __init__(
88
+ self,
89
+ embed_dims: int,
90
+ depth: int,
91
+ num_heads: int,
92
+ window_size: int = 8,
93
+ downsample: bool = False,
94
+ drop_paths: Union[Sequence[float], float] = 0.0,
95
+ with_cp: bool = False,
96
+ pad_small_map: bool = False,
97
+ extra_norm_every_n_blocks: int = 0,
98
+ pretrained_window_size: int = 0,
99
+ is_post_norm_downsample: bool = True,
100
+ ):
101
+ super().__init__()
102
+ if not isinstance(drop_paths, Sequence):
103
+ drop_paths = [drop_paths] * depth
104
+
105
+ if downsample:
106
+ self.out_channels = 2 * embed_dims
107
+ self.downsample = PatchMerging(
108
+ in_channels=embed_dims,
109
+ out_channels=self.out_channels,
110
+ is_post_norm=is_post_norm_downsample,
111
+ )
112
+ else:
113
+ self.out_channels = embed_dims
114
+ self.downsample = None
115
+
116
+ self.blocks = nn.ModuleList()
117
+ for i in range(depth):
118
+ extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0
119
+ self.blocks.append(
120
+ SwinBlockV2(
121
+ embed_dims=self.out_channels,
122
+ num_heads=num_heads,
123
+ window_size=window_size,
124
+ shift=(i % 2 == 1),
125
+ extra_norm=extra_norm,
126
+ drop_path=drop_paths[i],
127
+ with_cp=with_cp,
128
+ pad_small_map=pad_small_map,
129
+ pretrained_window_size=pretrained_window_size,
130
+ )
131
+ )
132
+
133
+ def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
134
+ if self.downsample is not None:
135
+ x, out_shape = self.downsample(x, in_shape)
136
+ else:
137
+ out_shape = in_shape
138
+
139
+ for block in self.blocks:
140
+ x = block(x, out_shape)
141
+ return x, out_shape
142
+
143
+
144
+ class ProjMHSA(nn.Module):
145
+ """Projected multi-head self-attention used in SkySense++ HR backbone."""
146
+
147
+ def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True):
148
+ super().__init__()
149
+ self.proj_in = nn.Linear(embed_dims, proj_dims)
150
+ self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias)
151
+ self.proj_out = nn.Linear(proj_dims, embed_dims)
152
+
153
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
154
+ x = self.proj_in(x)
155
+ x, _ = self.attn(x, x, x)
156
+ return self.proj_out(x)
157
+
158
+
159
+ class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel):
160
+ config_class = SkySensePlusPlusSwinV2MSLConfig
161
+ base_model_prefix = "skysensepp_swinv2_msl"
162
+ supports_gradient_checkpointing = True
163
+
164
+ def _init_weights(self, module):
165
+ if isinstance(module, nn.Linear):
166
+ nn.init.trunc_normal_(module.weight, std=0.02)
167
+ if module.bias is not None:
168
+ nn.init.zeros_(module.bias)
169
+ elif isinstance(module, nn.LayerNorm):
170
+ nn.init.ones_(module.weight)
171
+ nn.init.zeros_(module.bias)
172
+ elif isinstance(module, nn.Conv2d):
173
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
174
+ if module.bias is not None:
175
+ nn.init.zeros_(module.bias)
176
+
177
+
178
+ class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel):
179
+ """SkySense++ HR backbone with semantic vocabulary and annotation conditioning."""
180
+
181
+ def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig):
182
+ super().__init__(config)
183
+
184
+ self.num_layers = len(config.depths)
185
+ self.out_indices = config.out_indices
186
+ self.merge_stage = config.merge_stage
187
+ self.use_attn = config.use_attn
188
+ self.patch_size = config.patch_size
189
+
190
+ if isinstance(config.window_size, int):
191
+ window_sizes = [config.window_size] * self.num_layers
192
+ else:
193
+ window_sizes = list(config.window_size)
194
+
195
+ self.patch_embed = PatchEmbed(
196
+ in_channels=config.in_channels,
197
+ embed_dims=config.embed_dims,
198
+ kernel_size=config.patch_size,
199
+ stride=config.patch_size,
200
+ norm_layer=nn.LayerNorm,
201
+ input_size=config.img_size,
202
+ )
203
+
204
+ self.use_abs_pos_embed = config.use_abs_pos_embed
205
+ if self.use_abs_pos_embed:
206
+ patch_resolution = self.patch_embed.init_out_size
207
+ num_patches = patch_resolution[0] * patch_resolution[1]
208
+ self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
209
+
210
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
211
+
212
+ total_depth = sum(config.depths)
213
+ if total_depth > 1:
214
+ dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
215
+ else:
216
+ dpr = [0.0]
217
+
218
+ self.stages = nn.ModuleList()
219
+ embed_dims_list = [config.embed_dims]
220
+ for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)):
221
+ stage = SwinBlockV2Sequence(
222
+ embed_dims=embed_dims_list[-1],
223
+ depth=depth,
224
+ num_heads=num_heads,
225
+ window_size=window_sizes[i],
226
+ downsample=(i > 0),
227
+ drop_paths=dpr[:depth],
228
+ with_cp=config.with_cp,
229
+ pad_small_map=config.pad_small_map,
230
+ extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
231
+ pretrained_window_size=config.pretrained_window_sizes[i],
232
+ is_post_norm_downsample=config.is_post_norm_downsample,
233
+ )
234
+ self.stages.append(stage)
235
+ dpr = dpr[depth:]
236
+ embed_dims_list.append(stage.out_channels)
237
+
238
+ for i in self.out_indices:
239
+ self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1]))
240
+
241
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
242
+ self.vocabulary_token = nn.Parameter(
243
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
244
+ )
245
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
246
+
247
+ if self.use_attn:
248
+ self.attn1 = ProjMHSA(352, 256, num_heads=16)
249
+ self.attn2 = ProjMHSA(704, 512, num_heads=16)
250
+ self.attn3 = ProjMHSA(1408, 1024, num_heads=16)
251
+ self.norm_attn = nn.LayerNorm(1408)
252
+
253
+ self.post_init()
254
+
255
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
256
+ batch_size, height, width = anno_img.shape
257
+ ann_token = torch.index_select(
258
+ self.vocabulary_token, 0, anno_img.reshape(-1)
259
+ ).reshape(batch_size, height, width, -1)
260
+
261
+ num_patch_h = height // self.patch_size
262
+ num_patch_w = width // self.patch_size
263
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
264
+ weight = (
265
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
266
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
267
+ .reshape(1, height, width, 1)
268
+ )
269
+ ann_token = ann_token * weight
270
+ ann_token = F.avg_pool2d(
271
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
272
+ )
273
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
274
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
275
+ )
276
+
277
+ def forward(
278
+ self,
279
+ pixel_values: torch.Tensor,
280
+ annotation: torch.Tensor,
281
+ mask: Optional[torch.Tensor] = None,
282
+ output_hidden_states: Optional[bool] = None,
283
+ return_dict: Optional[bool] = None,
284
+ ) -> Union[Tuple, BaseModelOutput]:
285
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
286
+
287
+ x, hw_shape = self.patch_embed(pixel_values)
288
+ y = self.create_ann_token(annotation)
289
+ batch_size, num_tokens, channels = y.shape
290
+
291
+ if mask is not None:
292
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
293
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
294
+ y = y * (1.0 - weight) + mask_tokens * weight
295
+
296
+ if self.merge_stage == 0:
297
+ x = (x + y) * 0.5
298
+ else:
299
+ x = x.reshape(batch_size, *hw_shape, channels)
300
+ y = y.reshape(batch_size, *hw_shape, channels)
301
+ x = torch.cat((x, y), dim=2)
302
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
303
+ x = x.reshape(batch_size, -1, channels)
304
+
305
+ if self.use_abs_pos_embed:
306
+ x = x + self.absolute_pos_embed
307
+ x = self.drop_after_pos(x)
308
+
309
+ all_hidden_states = () if output_hidden_states else None
310
+ feature_maps = []
311
+ merge_idx = self.merge_stage - 1
312
+
313
+ for i, stage in enumerate(self.stages):
314
+ x, hw_shape = stage(x, hw_shape)
315
+ if i == merge_idx:
316
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
317
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
318
+ x = x.reshape(batch_size, -1, x.shape[-1])
319
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
320
+
321
+ if self.use_attn:
322
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
323
+ if i <= len(attention_blocks) - 1:
324
+ x = x + attention_blocks[i](x)
325
+ if i == len(attention_blocks) - 1:
326
+ x = self.norm_attn(x)
327
+
328
+ if output_hidden_states:
329
+ all_hidden_states = all_hidden_states + (x,)
330
+
331
+ if i in self.out_indices:
332
+ norm_layer = getattr(self, f"norm{i}")
333
+ out = norm_layer(x)
334
+ out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous()
335
+ feature_maps.append(out)
336
+
337
+ if not return_dict:
338
+ return tuple(feature_maps)
339
+
340
+ return BaseModelOutput(
341
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
342
+ hidden_states=all_hidden_states,
343
+ )
skysensepp-vit-msl-s1/modeling_skysensepp_vit_msl.py ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ import math
4
+ from typing import Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig
14
+ from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
15
+
16
+
17
+ class TransformerEncoderLayer(nn.Module):
18
+ def __init__(
19
+ self,
20
+ embed_dims: int,
21
+ num_heads: int,
22
+ feedforward_channels: int,
23
+ drop_rate: float = 0.0,
24
+ attn_drop_rate: float = 0.0,
25
+ drop_path_rate: float = 0.0,
26
+ num_fcs: int = 2,
27
+ qkv_bias: bool = True,
28
+ with_cp: bool = False,
29
+ ):
30
+ super().__init__()
31
+ self.with_cp = with_cp
32
+ self.norm1 = nn.LayerNorm(embed_dims)
33
+ self.attn = nn.MultiheadAttention(
34
+ embed_dim=embed_dims,
35
+ num_heads=num_heads,
36
+ dropout=attn_drop_rate,
37
+ bias=qkv_bias,
38
+ batch_first=True,
39
+ )
40
+ self.proj_drop = nn.Dropout(drop_rate)
41
+ self.norm2 = nn.LayerNorm(embed_dims)
42
+ self.ffn = FFN(
43
+ embed_dims=embed_dims,
44
+ feedforward_channels=feedforward_channels,
45
+ num_fcs=num_fcs,
46
+ ffn_drop=drop_rate,
47
+ drop_path=drop_path_rate,
48
+ act_layer=nn.GELU,
49
+ add_identity=True,
50
+ )
51
+ self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
52
+
53
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
54
+ def _inner_forward(x):
55
+ residual = x
56
+ x_norm = self.norm1(x)
57
+ attn_out, _ = self.attn(x_norm, x_norm, x_norm)
58
+ attn_out = self.proj_drop(attn_out)
59
+ x = residual + self.drop_path(attn_out)
60
+ return self.ffn(self.norm2(x), identity=x)
61
+
62
+ if self.with_cp and x.requires_grad:
63
+ return cp.checkpoint(_inner_forward, x, use_reentrant=False)
64
+ return _inner_forward(x)
65
+
66
+
67
+ class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel):
68
+ config_class = SkySensePlusPlusViTMSLConfig
69
+ base_model_prefix = "skysensepp_vit_msl"
70
+ supports_gradient_checkpointing = True
71
+
72
+ def _init_weights(self, module):
73
+ if isinstance(module, nn.Linear):
74
+ nn.init.trunc_normal_(module.weight, std=0.02)
75
+ if module.bias is not None:
76
+ nn.init.zeros_(module.bias)
77
+ elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
78
+ nn.init.ones_(module.weight)
79
+ nn.init.zeros_(module.bias)
80
+ elif isinstance(module, nn.Conv2d):
81
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
82
+ if module.bias is not None:
83
+ nn.init.zeros_(module.bias)
84
+
85
+
86
+ class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel):
87
+ """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning."""
88
+
89
+ def __init__(self, config: SkySensePlusPlusViTMSLConfig):
90
+ super().__init__(config)
91
+
92
+ img_size = to_2tuple(config.img_size)
93
+ self.img_size = img_size
94
+ self.patch_size = config.patch_size
95
+ self.with_cls_token = config.with_cls_token
96
+ self.output_cls_token = config.output_cls_token
97
+ self.merge_stage = config.merge_stage
98
+ self.use_attn = config.use_attn
99
+ self.interpolate_mode = "bicubic"
100
+
101
+ self.patch_embed = PatchEmbed(
102
+ in_channels=config.in_channels,
103
+ embed_dims=config.embed_dims,
104
+ kernel_size=config.patch_size,
105
+ stride=config.patch_size,
106
+ norm_layer=nn.LayerNorm if config.patch_norm else None,
107
+ )
108
+
109
+ num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size)
110
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
111
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
112
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
113
+
114
+ out_indices = list(config.out_indices)
115
+ self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices]
116
+
117
+ num_layers = config.num_layers
118
+ if num_layers > 1:
119
+ dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
120
+ else:
121
+ dpr = [0.0]
122
+
123
+ self.layers = nn.ModuleList()
124
+ for i in range(config.num_layers):
125
+ self.layers.append(
126
+ TransformerEncoderLayer(
127
+ embed_dims=config.embed_dims,
128
+ num_heads=config.num_heads,
129
+ feedforward_channels=config.mlp_ratio * config.embed_dims,
130
+ attn_drop_rate=config.attn_drop_rate,
131
+ drop_rate=config.drop_rate,
132
+ drop_path_rate=dpr[i],
133
+ num_fcs=2,
134
+ qkv_bias=config.qkv_bias,
135
+ with_cp=config.with_cp,
136
+ )
137
+ )
138
+
139
+ self.final_norm = config.final_norm
140
+ if config.final_norm:
141
+ self.norm = nn.LayerNorm(config.embed_dims)
142
+
143
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
144
+ self.vocabulary_token = nn.Parameter(
145
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
146
+ )
147
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
148
+
149
+ if self.use_attn:
150
+ self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
151
+ self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
152
+ self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
153
+ self.norm_attn = nn.LayerNorm(config.embed_dims)
154
+
155
+ self.post_init()
156
+
157
+ @staticmethod
158
+ def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"):
159
+ pos_h, pos_w = pos_shape
160
+ pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :]
161
+ pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
162
+ pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode)
163
+ return torch.flatten(pos_embed_weight, 2).transpose(1, 2)
164
+
165
+ def _pos_embedding(self, patched_img, hw_shape, pos_embed):
166
+ x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
167
+ if x_len != pos_len:
168
+ pos_h = self.img_size[0] // self.patch_size
169
+ pos_w = self.img_size[1] // self.patch_size
170
+ pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode)
171
+ return self.drop_after_pos(patched_img + pos_embed)
172
+
173
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
174
+ batch_size, height, width = anno_img.shape
175
+ ann_token = torch.index_select(
176
+ self.vocabulary_token, 0, anno_img.reshape(-1)
177
+ ).reshape(batch_size, height, width, -1)
178
+
179
+ num_patch_h = height // self.patch_size
180
+ num_patch_w = width // self.patch_size
181
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
182
+ weight = (
183
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
184
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
185
+ .reshape(1, height, width, 1)
186
+ )
187
+ ann_token = ann_token * weight
188
+ ann_token = F.avg_pool2d(
189
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
190
+ )
191
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
192
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
193
+ )
194
+
195
+ def forward(
196
+ self,
197
+ pixel_values: torch.Tensor,
198
+ annotation: torch.Tensor,
199
+ mask: Optional[torch.Tensor] = None,
200
+ output_hidden_states: Optional[bool] = None,
201
+ return_dict: Optional[bool] = None,
202
+ ) -> Union[Tuple, BaseModelOutput]:
203
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
204
+
205
+ x, hw_shape = self.patch_embed(pixel_values)
206
+ y = self.create_ann_token(annotation)
207
+ batch_size, num_tokens, channels = y.shape
208
+
209
+ if mask is not None:
210
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
211
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
212
+ y = y * (1.0 - weight) + mask_tokens * weight
213
+
214
+ if self.merge_stage == 0:
215
+ x = (x + y) * 0.5
216
+ else:
217
+ x = x.reshape(batch_size, *hw_shape, channels)
218
+ y = y.reshape(batch_size, *hw_shape, channels)
219
+ x = torch.cat((x, y), dim=2)
220
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
221
+ x = x.reshape(batch_size, -1, channels)
222
+
223
+ x = self._pos_embedding(x, hw_shape, self.pos_embed)
224
+
225
+ all_hidden_states = () if output_hidden_states else None
226
+ feature_maps = []
227
+ merge_idx = self.merge_stage - 1
228
+
229
+ for i, layer in enumerate(self.layers):
230
+ x = layer(x)
231
+
232
+ if i == merge_idx:
233
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
234
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
235
+ x = x.reshape(batch_size, -1, x.shape[-1])
236
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
237
+
238
+ if self.use_attn:
239
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
240
+ if i <= len(attention_blocks) - 1:
241
+ attn_out, _ = attention_blocks[i](x, x, x)
242
+ x = x + attn_out
243
+ if i == len(attention_blocks) - 1:
244
+ x = self.norm_attn(x)
245
+
246
+ if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm:
247
+ x = self.norm(x)
248
+
249
+ if output_hidden_states:
250
+ all_hidden_states = all_hidden_states + (x,)
251
+
252
+ if i in self.out_indices:
253
+ out = x
254
+ out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous()
255
+ if self.output_cls_token:
256
+ out = [out, x[:, 0]]
257
+ feature_maps.append(out)
258
+
259
+ if not return_dict:
260
+ return tuple(feature_maps)
261
+
262
+ return BaseModelOutput(
263
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
264
+ hidden_states=all_hidden_states,
265
+ )
skysensepp-vit-msl-s1/modeling_utils.py ADDED
@@ -0,0 +1,557 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense: Pure PyTorch + HuggingFace Transformers implementation.
2
+
3
+ Shared utility modules used across SkySense model implementations.
4
+ """
5
+
6
+ import math
7
+ from typing import Optional, Tuple
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+
14
+ def to_2tuple(x):
15
+ """Convert to a 2-tuple."""
16
+ if isinstance(x, (list, tuple)):
17
+ return tuple(x)
18
+ return (x, x)
19
+
20
+
21
+ class DropPath(nn.Module):
22
+ """Drop paths (stochastic depth) per sample.
23
+
24
+ Args:
25
+ drop_prob (float): Probability of dropping a path. Default: 0.0.
26
+ """
27
+
28
+ def __init__(self, drop_prob: float = 0.0):
29
+ super().__init__()
30
+ self.drop_prob = drop_prob
31
+
32
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
33
+ if self.drop_prob == 0.0 or not self.training:
34
+ return x
35
+ keep_prob = 1 - self.drop_prob
36
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1)
37
+ random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
38
+ random_tensor = torch.floor(random_tensor + keep_prob)
39
+ output = x / keep_prob * random_tensor
40
+ return output
41
+
42
+
43
+ class PatchEmbed(nn.Module):
44
+ """Image to Patch Embedding using Conv2d.
45
+
46
+ Args:
47
+ in_channels (int): Number of input channels. Default: 3.
48
+ embed_dims (int): Embedding dimension. Default: 96.
49
+ kernel_size (int): Kernel size of the projection. Default: 4.
50
+ stride (int): Stride of the projection. Default: 4.
51
+ padding (int): Padding of the projection. Default: 0.
52
+ norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
53
+ input_size (int or tuple or None): Input resolution for calculating output size.
54
+ """
55
+
56
+ def __init__(
57
+ self,
58
+ in_channels: int = 3,
59
+ embed_dims: int = 96,
60
+ kernel_size: int = 4,
61
+ stride: int = 4,
62
+ padding: int = 0,
63
+ norm_layer: Optional[type] = nn.LayerNorm,
64
+ input_size: Optional[int] = None,
65
+ ):
66
+ super().__init__()
67
+ self.projection = nn.Conv2d(
68
+ in_channels, embed_dims,
69
+ kernel_size=kernel_size, stride=stride, padding=padding,
70
+ )
71
+ self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
72
+
73
+ # Compute init output size if input_size is given
74
+ if input_size is not None:
75
+ input_size = to_2tuple(input_size)
76
+ self.init_out_size = (
77
+ (input_size[0] - kernel_size + 2 * padding) // stride + 1,
78
+ (input_size[1] - kernel_size + 2 * padding) // stride + 1,
79
+ )
80
+ else:
81
+ self.init_out_size = None
82
+
83
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
84
+ x = self.projection(x) # (B, C, H, W)
85
+ out_size = (x.shape[2], x.shape[3])
86
+ x = x.flatten(2).transpose(1, 2) # (B, H*W, C)
87
+ x = self.norm(x)
88
+ return x, out_size
89
+
90
+
91
+ class FFN(nn.Module):
92
+ """Feed-Forward Network.
93
+
94
+ Args:
95
+ embed_dims (int): Input dimension.
96
+ feedforward_channels (int): Hidden dimension.
97
+ num_fcs (int): Number of FC layers. Default: 2.
98
+ ffn_drop (float): Dropout rate. Default: 0.0.
99
+ drop_path (float): Drop path rate. Default: 0.0.
100
+ act_layer (nn.Module): Activation layer class. Default: nn.GELU.
101
+ add_identity (bool): Whether to add identity connection. Default: True.
102
+ """
103
+
104
+ def __init__(
105
+ self,
106
+ embed_dims: int,
107
+ feedforward_channels: int,
108
+ num_fcs: int = 2,
109
+ ffn_drop: float = 0.0,
110
+ drop_path: float = 0.0,
111
+ act_layer: type = nn.GELU,
112
+ add_identity: bool = True,
113
+ ):
114
+ super().__init__()
115
+ assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
116
+ self.embed_dims = embed_dims
117
+ self.feedforward_channels = feedforward_channels
118
+ self.add_identity = add_identity
119
+
120
+ layers = []
121
+ in_channels = embed_dims
122
+ for i in range(num_fcs - 1):
123
+ layers.append(nn.Linear(in_channels, feedforward_channels))
124
+ layers.append(act_layer())
125
+ layers.append(nn.Dropout(ffn_drop))
126
+ in_channels = feedforward_channels
127
+ layers.append(nn.Linear(feedforward_channels, embed_dims))
128
+ layers.append(nn.Dropout(ffn_drop))
129
+ self.layers = nn.Sequential(*layers)
130
+
131
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
132
+
133
+ def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
134
+ out = self.layers(x)
135
+ out = self.drop_path(out)
136
+ if self.add_identity:
137
+ if identity is None:
138
+ identity = x
139
+ out = out + identity
140
+ return out
141
+
142
+
143
+ class WindowMSAV2(nn.Module):
144
+ """Window-based Multi-head Self-Attention for Swin Transformer V2.
145
+
146
+ Uses cosine attention and log-spaced continuous position bias (log-CPB).
147
+
148
+ Args:
149
+ embed_dims (int): Number of input channels.
150
+ num_heads (int): Number of attention heads.
151
+ window_size (tuple[int]): Window size (Wh, Ww).
152
+ pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
153
+ qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
154
+ attn_drop (float): Attention dropout rate. Default: 0.0.
155
+ proj_drop (float): Output projection dropout rate. Default: 0.0.
156
+ """
157
+
158
+ def __init__(
159
+ self,
160
+ embed_dims: int,
161
+ num_heads: int,
162
+ window_size: Tuple[int, int],
163
+ pretrained_window_size: Tuple[int, int] = (0, 0),
164
+ qkv_bias: bool = True,
165
+ attn_drop: float = 0.0,
166
+ proj_drop: float = 0.0,
167
+ ):
168
+ super().__init__()
169
+ self.embed_dims = embed_dims
170
+ self.num_heads = num_heads
171
+ self.window_size = window_size
172
+ self.pretrained_window_size = pretrained_window_size
173
+
174
+ self.logit_scale = nn.Parameter(
175
+ torch.log(10 * torch.ones((num_heads, 1, 1))))
176
+
177
+ # MLP for continuous relative position bias (log-CPB)
178
+ self.cpb_mlp = nn.Sequential(
179
+ nn.Linear(2, 512, bias=True),
180
+ nn.ReLU(inplace=True),
181
+ nn.Linear(512, num_heads, bias=False),
182
+ )
183
+
184
+ # Build relative coords table
185
+ self._build_relative_coords_table()
186
+ # Build relative position index
187
+ self._build_relative_position_index()
188
+
189
+ self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
190
+ if qkv_bias:
191
+ self.q_bias = nn.Parameter(torch.zeros(embed_dims))
192
+ self.v_bias = nn.Parameter(torch.zeros(embed_dims))
193
+ else:
194
+ self.q_bias = None
195
+ self.v_bias = None
196
+
197
+ self.attn_drop = nn.Dropout(attn_drop)
198
+ self.proj = nn.Linear(embed_dims, embed_dims)
199
+ self.proj_drop = nn.Dropout(proj_drop)
200
+ self.softmax = nn.Softmax(dim=-1)
201
+
202
+ def _build_relative_coords_table(self):
203
+ """Build the relative coordinates table for log-CPB."""
204
+ Wh, Ww = self.window_size
205
+ # Table of relative coordinates
206
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
207
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
208
+ coords_table = torch.stack(
209
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
210
+ ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2)
211
+
212
+ # Normalize to [-1, 1] and apply log-scale
213
+ if self.pretrained_window_size[0] > 0:
214
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
215
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
216
+ else:
217
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
218
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
219
+ coords_table *= 8 # normalize to -8, 8
220
+ coords_table = (
221
+ torch.sign(coords_table)
222
+ * torch.log2(torch.abs(coords_table) + 1.0)
223
+ / math.log2(8)
224
+ )
225
+ self.register_buffer("relative_coords_table", coords_table)
226
+
227
+ def _build_relative_position_index(self):
228
+ """Build the pairwise relative position index for each window token."""
229
+ Wh, Ww = self.window_size
230
+ coords_h = torch.arange(Wh)
231
+ coords_w = torch.arange(Ww)
232
+ coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
233
+ coords_flatten = coords.view(2, -1)
234
+
235
+ relative_coords = (
236
+ coords_flatten[:, :, None] - coords_flatten[:, None, :]
237
+ ) # (2, Wh*Ww, Wh*Ww)
238
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
239
+ relative_coords[:, :, 0] += Wh - 1
240
+ relative_coords[:, :, 1] += Ww - 1
241
+ relative_coords[:, :, 0] *= 2 * Ww - 1
242
+ relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww)
243
+ self.register_buffer("relative_position_index", relative_position_index)
244
+
245
+ def _compute_position_bias(self, N):
246
+ """Compute relative position bias, supporting dynamic window sizes.
247
+
248
+ The log-CPB (Continuous Position Bias) MLP can generalize to any window
249
+ size by computing bias from normalized relative coordinates.
250
+ """
251
+ init_N = self.window_size[0] * self.window_size[1]
252
+ if N == init_N:
253
+ # Use pre-built tables
254
+ relative_position_bias_table = self.cpb_mlp(
255
+ self.relative_coords_table
256
+ ).view(-1, self.num_heads)
257
+ relative_position_bias = relative_position_bias_table[
258
+ self.relative_position_index.view(-1)
259
+ ].view(N, N, -1)
260
+ else:
261
+ # Dynamic: compute for actual window size on-the-fly
262
+ Wh = Ww = int(math.sqrt(N))
263
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
264
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
265
+ coords_table = torch.stack(
266
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
267
+ ).flatten(1).transpose(0, 1).unsqueeze(0)
268
+ if self.pretrained_window_size[0] > 0:
269
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
270
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
271
+ else:
272
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
273
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
274
+ coords_table *= 8
275
+ coords_table = (
276
+ torch.sign(coords_table)
277
+ * torch.log2(torch.abs(coords_table) + 1.0)
278
+ / math.log2(8)
279
+ )
280
+ # Build position index for actual window size
281
+ ch = torch.arange(Wh, device=self.logit_scale.device)
282
+ cw = torch.arange(Ww, device=self.logit_scale.device)
283
+ coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
284
+ coords_flat = coords.view(2, -1)
285
+ rel = coords_flat[:, :, None] - coords_flat[:, None, :]
286
+ rel = rel.permute(1, 2, 0).contiguous()
287
+ rel[:, :, 0] += Wh - 1
288
+ rel[:, :, 1] += Ww - 1
289
+ rel[:, :, 0] *= 2 * Ww - 1
290
+ pos_index = rel.sum(-1)
291
+
292
+ bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
293
+ relative_position_bias = bias_table[
294
+ pos_index.view(-1)
295
+ ].view(N, N, -1)
296
+
297
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
298
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
299
+ return relative_position_bias
300
+
301
+ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
302
+ """
303
+ Args:
304
+ x: (num_windows*B, N, C) where N = Wh*Ww
305
+ mask: (num_windows, N, N) or None
306
+ """
307
+ B_, N, C = x.shape
308
+
309
+ # Compute QKV with bias
310
+ if self.q_bias is not None:
311
+ qkv_bias = torch.cat(
312
+ (self.q_bias,
313
+ torch.zeros_like(self.v_bias, requires_grad=False),
314
+ self.v_bias))
315
+ qkv = F.linear(x, self.qkv.weight, qkv_bias)
316
+ else:
317
+ qkv = self.qkv(x)
318
+
319
+ qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
320
+ qkv = qkv.permute(2, 0, 3, 1, 4)
321
+ q, k, v = qkv.unbind(0)
322
+
323
+ # Cosine attention
324
+ attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
325
+ logit_scale = torch.clamp(
326
+ self.logit_scale, max=math.log(1.0 / 0.01)
327
+ ).exp()
328
+ attn = attn * logit_scale
329
+
330
+ # Log-CPB relative position bias (supports dynamic window sizes)
331
+ relative_position_bias = self._compute_position_bias(N)
332
+ attn = attn + relative_position_bias.unsqueeze(0)
333
+
334
+ if mask is not None:
335
+ nW = mask.shape[0]
336
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
337
+ attn = attn + mask.unsqueeze(1).unsqueeze(0)
338
+ attn = attn.view(-1, self.num_heads, N, N)
339
+
340
+ attn = self.softmax(attn)
341
+ attn = self.attn_drop(attn)
342
+
343
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
344
+ x = self.proj(x)
345
+ x = self.proj_drop(x)
346
+ return x
347
+
348
+
349
+ class ShiftWindowMSA(nn.Module):
350
+ """Shifted Window Multi-head Self-Attention.
351
+
352
+ Args:
353
+ embed_dims (int): Number of input channels.
354
+ num_heads (int): Number of attention heads.
355
+ window_size (int): Window size.
356
+ shift_size (int): Shift size for SW-MSA. Default: 0.
357
+ attn_drop (float): Attention dropout rate. Default: 0.0.
358
+ proj_drop (float): Projection dropout rate. Default: 0.0.
359
+ drop_path (float): Drop path rate. Default: 0.0.
360
+ pad_small_map (bool): Pad small feature maps to window size. Default: False.
361
+ pretrained_window_size (int): Pretrained window size. Default: 0.
362
+ """
363
+
364
+ def __init__(
365
+ self,
366
+ embed_dims: int,
367
+ num_heads: int,
368
+ window_size: int,
369
+ shift_size: int = 0,
370
+ attn_drop: float = 0.0,
371
+ proj_drop: float = 0.0,
372
+ drop_path: float = 0.0,
373
+ pad_small_map: bool = False,
374
+ pretrained_window_size: int = 0,
375
+ ):
376
+ super().__init__()
377
+ self.window_size = window_size
378
+ self.shift_size = shift_size
379
+ self.pad_small_map = pad_small_map
380
+
381
+ self.w_msa = WindowMSAV2(
382
+ embed_dims=embed_dims,
383
+ num_heads=num_heads,
384
+ window_size=to_2tuple(window_size),
385
+ pretrained_window_size=to_2tuple(pretrained_window_size),
386
+ attn_drop=attn_drop,
387
+ proj_drop=proj_drop,
388
+ )
389
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
390
+
391
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
392
+ B, L, C = x.shape
393
+ H, W = hw_shape
394
+ assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
395
+
396
+ x = x.view(B, H, W, C)
397
+
398
+ window_size = self.window_size
399
+ shift_size = self.shift_size
400
+
401
+ # Pad or shrink window
402
+ if self.pad_small_map:
403
+ pad_r = (window_size - W % window_size) % window_size
404
+ pad_b = (window_size - H % window_size) % window_size
405
+ x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
406
+ _, Hp, Wp, _ = x.shape
407
+ else:
408
+ Hp, Wp = H, W
409
+ if window_size > Hp:
410
+ window_size = Hp
411
+ shift_size = 0
412
+ if window_size > Wp:
413
+ window_size = Wp
414
+ shift_size = 0
415
+
416
+ # Compute attention mask for SW-MSA
417
+ attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
418
+
419
+ # Cyclic shift
420
+ if shift_size > 0:
421
+ x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
422
+
423
+ # Partition windows
424
+ x_windows = self._window_partition(x, window_size)
425
+ # (num_windows*B, window_size*window_size, C)
426
+
427
+ # W-MSA/SW-MSA
428
+ attn_windows = self.w_msa(x_windows, mask=attn_mask)
429
+
430
+ # Merge windows
431
+ x = self._window_reverse(attn_windows, window_size, Hp, Wp)
432
+
433
+ # Reverse cyclic shift
434
+ if shift_size > 0:
435
+ x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
436
+
437
+ if self.pad_small_map and (pad_r > 0 or pad_b > 0):
438
+ x = x[:, :H, :W, :].contiguous()
439
+
440
+ x = x.view(B, H * W, C)
441
+ x = self.drop_path(x)
442
+ return x
443
+
444
+ @staticmethod
445
+ def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
446
+ """Partition into non-overlapping windows."""
447
+ B, H, W, C = x.shape
448
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
449
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
450
+ windows = windows.view(-1, window_size * window_size, C)
451
+ return windows
452
+
453
+ @staticmethod
454
+ def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
455
+ """Reverse window partition."""
456
+ B_nW = windows.shape[0]
457
+ nH = H // window_size
458
+ nW = W // window_size
459
+ B = B_nW // (nH * nW)
460
+ x = windows.view(B, nH, nW, window_size, window_size, -1)
461
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
462
+ x = x.view(B, H, W, -1)
463
+ return x
464
+
465
+ @staticmethod
466
+ def _compute_attn_mask(H, W, window_size, shift_size, device):
467
+ """Compute attention mask for shifted window attention."""
468
+ if shift_size <= 0:
469
+ return None
470
+ img_mask = torch.zeros((1, H, W, 1), device=device)
471
+ h_slices = (
472
+ slice(0, -window_size),
473
+ slice(-window_size, -shift_size),
474
+ slice(-shift_size, None),
475
+ )
476
+ w_slices = (
477
+ slice(0, -window_size),
478
+ slice(-window_size, -shift_size),
479
+ slice(-shift_size, None),
480
+ )
481
+ cnt = 0
482
+ for h in h_slices:
483
+ for w in w_slices:
484
+ img_mask[:, h, w, :] = cnt
485
+ cnt += 1
486
+
487
+ # Partition mask
488
+ mask_windows = img_mask.view(
489
+ 1, H // window_size, window_size, W // window_size, window_size, 1
490
+ )
491
+ mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
492
+ mask_windows = mask_windows.view(-1, window_size * window_size)
493
+
494
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
495
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
496
+ attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
497
+ return attn_mask
498
+
499
+
500
+ class PatchMerging(nn.Module):
501
+ """Patch Merging Layer for downsampling (2x).
502
+
503
+ Args:
504
+ in_channels (int): Input channels.
505
+ out_channels (int): Output channels.
506
+ norm_layer (type): Normalization layer. Default: nn.LayerNorm.
507
+ is_post_norm (bool): Apply norm after linear. Default: True.
508
+ """
509
+
510
+ def __init__(
511
+ self,
512
+ in_channels: int,
513
+ out_channels: int,
514
+ norm_layer: type = nn.LayerNorm,
515
+ is_post_norm: bool = True,
516
+ ):
517
+ super().__init__()
518
+ self.in_channels = in_channels
519
+ self.out_channels = out_channels
520
+ self.is_post_norm = is_post_norm
521
+ self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
522
+ if is_post_norm:
523
+ self.norm = norm_layer(out_channels)
524
+ else:
525
+ self.norm = norm_layer(4 * in_channels)
526
+
527
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
528
+ B, L, C = x.shape
529
+ H, W = hw_shape
530
+ assert L == H * W
531
+
532
+ x = x.view(B, H, W, C)
533
+
534
+ # Pad if needed
535
+ pad_h = H % 2
536
+ pad_w = W % 2
537
+ if pad_h or pad_w:
538
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
539
+
540
+ x0 = x[:, 0::2, 0::2, :]
541
+ x1 = x[:, 1::2, 0::2, :]
542
+ x2 = x[:, 0::2, 1::2, :]
543
+ x3 = x[:, 1::2, 1::2, :]
544
+ x = torch.cat([x0, x1, x2, x3], dim=-1)
545
+
546
+ out_h = (H + pad_h) // 2
547
+ out_w = (W + pad_w) // 2
548
+ x = x.view(B, out_h * out_w, 4 * C)
549
+
550
+ if self.is_post_norm:
551
+ x = self.reduction(x)
552
+ x = self.norm(x)
553
+ else:
554
+ x = self.norm(x)
555
+ x = self.reduction(x)
556
+
557
+ return x, (out_h, out_w)
skysensepp-vit-msl-s1/pipeline_skysensepp.py ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Custom HuggingFace pipeline for SkySense++ MSL feature extraction."""
2
+
3
+ from typing import Any, Dict, Optional, Union
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline):
11
+ """Pipeline for SkySense++ MSL backbones.
12
+
13
+ Expects image tensors plus semantic annotation maps (class indices).
14
+ """
15
+
16
+ def _sanitize_parameters(
17
+ self,
18
+ annotation=None,
19
+ mask=None,
20
+ output_hidden_states=None,
21
+ **kwargs,
22
+ ):
23
+ preprocess_params = {}
24
+ forward_params = {}
25
+ postprocess_params = {}
26
+
27
+ if annotation is not None:
28
+ preprocess_params["annotation"] = annotation
29
+ if mask is not None:
30
+ forward_params["mask"] = mask
31
+ if output_hidden_states is not None:
32
+ forward_params["output_hidden_states"] = output_hidden_states
33
+
34
+ return preprocess_params, forward_params, postprocess_params
35
+
36
+ def preprocess(
37
+ self,
38
+ pixel_values: Any,
39
+ annotation: Optional[Any] = None,
40
+ **kwargs,
41
+ ) -> Dict[str, torch.Tensor]:
42
+ if isinstance(pixel_values, dict):
43
+ annotation = pixel_values.get("annotation", annotation)
44
+ pixel_values = pixel_values.get("pixel_values", pixel_values)
45
+
46
+ if isinstance(pixel_values, np.ndarray):
47
+ pixel_values = torch.from_numpy(pixel_values).float()
48
+ elif not isinstance(pixel_values, torch.Tensor):
49
+ raise TypeError(
50
+ f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}"
51
+ )
52
+
53
+ if annotation is None:
54
+ raise ValueError("SkySense++ MSL models require an `annotation` semantic map.")
55
+
56
+ if isinstance(annotation, np.ndarray):
57
+ annotation = torch.from_numpy(annotation).long()
58
+ elif not isinstance(annotation, torch.Tensor):
59
+ raise TypeError(
60
+ f"Expected tensor or ndarray for annotation, got {type(annotation)}"
61
+ )
62
+
63
+ if pixel_values.ndim == 3:
64
+ pixel_values = pixel_values.unsqueeze(0)
65
+ if annotation.ndim == 2:
66
+ annotation = annotation.unsqueeze(0)
67
+
68
+ return {"pixel_values": pixel_values, "annotation": annotation}
69
+
70
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
71
+ with torch.no_grad():
72
+ outputs = self.model(
73
+ pixel_values=model_inputs["pixel_values"],
74
+ annotation=model_inputs["annotation"],
75
+ mask=kwargs.get("mask"),
76
+ output_hidden_states=kwargs.get("output_hidden_states", False),
77
+ return_dict=True,
78
+ )
79
+ return {"outputs": outputs}
80
+
81
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
82
+ outputs = model_outputs["outputs"]
83
+ result = {"last_hidden_state": outputs.last_hidden_state}
84
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
85
+ result["hidden_states"] = outputs.hidden_states
86
+ return result
skysensepp-vit-msl-s1/pipeline_skysensepp_fusion.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Optional pipeline for SkySense++ fusion neck."""
2
+
3
+ from typing import Any, Dict
4
+
5
+ import numpy as np
6
+ import torch
7
+ from transformers import Pipeline
8
+
9
+
10
+ class SkySensePlusPlusFusionNeckPipeline(Pipeline):
11
+ """Pipeline for the optional SkySense++ fusion neck module.
12
+
13
+ Expects concatenated multi-modal tokens per spatial location:
14
+ ``(batch, num_modalities, input_dims)``.
15
+ """
16
+
17
+ def _sanitize_parameters(self, output_hidden_states=None, **kwargs):
18
+ preprocess_params = {}
19
+ forward_params = {}
20
+ postprocess_params = {}
21
+ if output_hidden_states is not None:
22
+ forward_params["output_hidden_states"] = output_hidden_states
23
+ return preprocess_params, forward_params, postprocess_params
24
+
25
+ def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]:
26
+ if isinstance(hidden_states, dict):
27
+ hidden_states = hidden_states["hidden_states"]
28
+
29
+ if isinstance(hidden_states, np.ndarray):
30
+ hidden_states = torch.from_numpy(hidden_states).float()
31
+ elif not isinstance(hidden_states, torch.Tensor):
32
+ raise TypeError(
33
+ f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}"
34
+ )
35
+ if hidden_states.ndim == 2:
36
+ hidden_states = hidden_states.unsqueeze(0)
37
+ return {"hidden_states": hidden_states}
38
+
39
+ def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]:
40
+ with torch.no_grad():
41
+ outputs = self.model(
42
+ hidden_states=model_inputs["hidden_states"],
43
+ output_hidden_states=kwargs.get("output_hidden_states", False),
44
+ return_dict=True,
45
+ )
46
+ return {"outputs": outputs}
47
+
48
+ def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
49
+ outputs = model_outputs["outputs"]
50
+ result = {"pooler_output": outputs.pooler_output}
51
+ if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
52
+ result["hidden_states"] = outputs.hidden_states
53
+ return result
skysensepp-vit-msl-s2/__init__.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++: Multi-Modal Remote Sensing Foundation Model (HuggingFace)."""
2
+
3
+ from .configuration_skysensepp import (
4
+ SkySensePlusPlusSwinV2MSLConfig,
5
+ SkySensePlusPlusViTMSLConfig,
6
+ )
7
+ from .modeling_skysensepp_swinv2_msl import (
8
+ SkySensePlusPlusSwinV2MSLModel,
9
+ SkySensePlusPlusSwinV2MSLPreTrainedModel,
10
+ )
11
+ from .modeling_skysensepp_vit_msl import (
12
+ SkySensePlusPlusViTMSLModel,
13
+ SkySensePlusPlusViTMSLPreTrainedModel,
14
+ )
15
+ from .pipeline_skysensepp import SkySensePlusPlusMSLFeatureExtractionPipeline
16
+
17
+ __all__ = [
18
+ "SkySensePlusPlusSwinV2MSLConfig",
19
+ "SkySensePlusPlusViTMSLConfig",
20
+ "SkySensePlusPlusSwinV2MSLModel",
21
+ "SkySensePlusPlusSwinV2MSLPreTrainedModel",
22
+ "SkySensePlusPlusViTMSLModel",
23
+ "SkySensePlusPlusViTMSLPreTrainedModel",
24
+ "SkySensePlusPlusMSLFeatureExtractionPipeline",
25
+ ]
skysensepp-vit-msl-s2/config.json ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "return_dict": true,
3
+ "output_hidden_states": false,
4
+ "dtype": "float32",
5
+ "chunk_size_feed_forward": 0,
6
+ "is_encoder_decoder": false,
7
+ "architectures": [
8
+ "SkySensePlusPlusViTMSLModel"
9
+ ],
10
+ "id2label": {
11
+ "0": "LABEL_0",
12
+ "1": "LABEL_1"
13
+ },
14
+ "label2id": {
15
+ "LABEL_0": 0,
16
+ "LABEL_1": 1
17
+ },
18
+ "problem_type": null,
19
+ "_name_or_path": "",
20
+ "transformers_version": "5.0.0",
21
+ "img_size": 16,
22
+ "patch_size": 4,
23
+ "in_channels": 10,
24
+ "embed_dims": 1024,
25
+ "num_layers": 24,
26
+ "num_heads": 16,
27
+ "mlp_ratio": 4,
28
+ "out_indices": [
29
+ 5,
30
+ 11,
31
+ 17,
32
+ 23
33
+ ],
34
+ "qkv_bias": true,
35
+ "drop_rate": 0.0,
36
+ "attn_drop_rate": 0.0,
37
+ "drop_path_rate": 0.3,
38
+ "with_cls_token": false,
39
+ "output_cls_token": false,
40
+ "patch_norm": false,
41
+ "final_norm": false,
42
+ "with_cp": false,
43
+ "vocabulary_size": 64,
44
+ "num_vocabulary_tokens": 65,
45
+ "merge_stage": 4,
46
+ "use_attn": false,
47
+ "modality": "s2",
48
+ "model_type": "skysensepp_vit_msl",
49
+ "output_attentions": false,
50
+ "auto_map": {
51
+ "AutoConfig": "configuration_skysensepp.SkySensePlusPlusViTMSLConfig",
52
+ "AutoModel": "modeling_skysensepp_vit_msl.SkySensePlusPlusViTMSLModel"
53
+ },
54
+ "custom_pipelines": {
55
+ "skysensepp-feature-extraction": {
56
+ "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline",
57
+ "pt": [
58
+ "AutoModel"
59
+ ]
60
+ },
61
+ "image-feature-extraction": {
62
+ "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline",
63
+ "pt": [
64
+ "AutoModel"
65
+ ]
66
+ }
67
+ }
68
+ }
skysensepp-vit-msl-s2/configuration_skysensepp.py ADDED
@@ -0,0 +1,124 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Configuration classes for SkySense++ MSL backbones."""
2
+
3
+ from transformers import PretrainedConfig
4
+
5
+
6
+ class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig):
7
+ """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical)."""
8
+
9
+ model_type = "skysensepp_swinv2_msl"
10
+
11
+ arch_zoo = {
12
+ "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
13
+ "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0},
14
+ "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0},
15
+ "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0},
16
+ "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6},
17
+ "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6},
18
+ }
19
+
20
+ def __init__(
21
+ self,
22
+ arch="huge",
23
+ img_size=224,
24
+ patch_size=4,
25
+ in_channels=3,
26
+ window_size=8,
27
+ drop_rate=0.0,
28
+ drop_path_rate=0.2,
29
+ out_indices=(0, 1, 2, 3),
30
+ use_abs_pos_embed=False,
31
+ with_cp=False,
32
+ pad_small_map=False,
33
+ pretrained_window_sizes=(0, 0, 0, 0),
34
+ is_post_norm_downsample=True,
35
+ vocabulary_size=64,
36
+ merge_stage=2,
37
+ use_attn=True,
38
+ **kwargs,
39
+ ):
40
+ super().__init__(**kwargs)
41
+
42
+ arch = arch.lower()
43
+ if arch not in self.arch_zoo:
44
+ raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}")
45
+ arch_settings = self.arch_zoo[arch]
46
+
47
+ self.arch = arch
48
+ self.embed_dims = arch_settings["embed_dims"]
49
+ self.depths = arch_settings["depths"]
50
+ self.num_heads = arch_settings["num_heads"]
51
+ self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"]
52
+
53
+ self.img_size = img_size
54
+ self.patch_size = patch_size
55
+ self.in_channels = in_channels
56
+ self.window_size = window_size
57
+ self.drop_rate = drop_rate
58
+ self.drop_path_rate = drop_path_rate
59
+ self.out_indices = list(out_indices)
60
+ self.use_abs_pos_embed = use_abs_pos_embed
61
+ self.with_cp = with_cp
62
+ self.pad_small_map = pad_small_map
63
+ self.pretrained_window_sizes = list(pretrained_window_sizes)
64
+ self.is_post_norm_downsample = is_post_norm_downsample
65
+
66
+ self.vocabulary_size = vocabulary_size
67
+ self.num_vocabulary_tokens = vocabulary_size + 1
68
+ self.merge_stage = merge_stage
69
+ self.use_attn = use_attn
70
+
71
+
72
+ class SkySensePlusPlusViTMSLConfig(PretrainedConfig):
73
+ """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1)."""
74
+
75
+ model_type = "skysensepp_vit_msl"
76
+
77
+ def __init__(
78
+ self,
79
+ img_size=16,
80
+ patch_size=4,
81
+ in_channels=10,
82
+ embed_dims=1024,
83
+ num_layers=24,
84
+ num_heads=16,
85
+ mlp_ratio=4,
86
+ out_indices=(5, 11, 17, 23),
87
+ qkv_bias=True,
88
+ drop_rate=0.0,
89
+ attn_drop_rate=0.0,
90
+ drop_path_rate=0.3,
91
+ with_cls_token=False,
92
+ output_cls_token=False,
93
+ patch_norm=False,
94
+ final_norm=False,
95
+ with_cp=False,
96
+ vocabulary_size=64,
97
+ merge_stage=4,
98
+ use_attn=False,
99
+ modality="s2",
100
+ **kwargs,
101
+ ):
102
+ super().__init__(**kwargs)
103
+ self.img_size = img_size
104
+ self.patch_size = patch_size
105
+ self.in_channels = in_channels
106
+ self.embed_dims = embed_dims
107
+ self.num_layers = num_layers
108
+ self.num_heads = num_heads
109
+ self.mlp_ratio = mlp_ratio
110
+ self.out_indices = list(out_indices)
111
+ self.qkv_bias = qkv_bias
112
+ self.drop_rate = drop_rate
113
+ self.attn_drop_rate = attn_drop_rate
114
+ self.drop_path_rate = drop_path_rate
115
+ self.with_cls_token = with_cls_token
116
+ self.output_cls_token = output_cls_token
117
+ self.patch_norm = patch_norm
118
+ self.final_norm = final_norm
119
+ self.with_cp = with_cp
120
+ self.vocabulary_size = vocabulary_size
121
+ self.num_vocabulary_tokens = vocabulary_size + 1
122
+ self.merge_stage = merge_stage
123
+ self.use_attn = use_attn
124
+ self.modality = modality
skysensepp-vit-msl-s2/conversion_manifest.json ADDED
@@ -0,0 +1,305 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release_s2.pth",
3
+ "modality": "s2",
4
+ "model_class": "SkySensePlusPlusViTMSLModel",
5
+ "num_tensors": 295,
6
+ "missing_keys": [],
7
+ "unexpected_keys": [],
8
+ "tensor_names": [
9
+ "cls_token",
10
+ "layers.0.attn.in_proj_bias",
11
+ "layers.0.attn.in_proj_weight",
12
+ "layers.0.attn.out_proj.bias",
13
+ "layers.0.attn.out_proj.weight",
14
+ "layers.0.ffn.layers.0.bias",
15
+ "layers.0.ffn.layers.0.weight",
16
+ "layers.0.ffn.layers.3.bias",
17
+ "layers.0.ffn.layers.3.weight",
18
+ "layers.0.norm1.bias",
19
+ "layers.0.norm1.weight",
20
+ "layers.0.norm2.bias",
21
+ "layers.0.norm2.weight",
22
+ "layers.1.attn.in_proj_bias",
23
+ "layers.1.attn.in_proj_weight",
24
+ "layers.1.attn.out_proj.bias",
25
+ "layers.1.attn.out_proj.weight",
26
+ "layers.1.ffn.layers.0.bias",
27
+ "layers.1.ffn.layers.0.weight",
28
+ "layers.1.ffn.layers.3.bias",
29
+ "layers.1.ffn.layers.3.weight",
30
+ "layers.1.norm1.bias",
31
+ "layers.1.norm1.weight",
32
+ "layers.1.norm2.bias",
33
+ "layers.1.norm2.weight",
34
+ "layers.10.attn.in_proj_bias",
35
+ "layers.10.attn.in_proj_weight",
36
+ "layers.10.attn.out_proj.bias",
37
+ "layers.10.attn.out_proj.weight",
38
+ "layers.10.ffn.layers.0.bias",
39
+ "layers.10.ffn.layers.0.weight",
40
+ "layers.10.ffn.layers.3.bias",
41
+ "layers.10.ffn.layers.3.weight",
42
+ "layers.10.norm1.bias",
43
+ "layers.10.norm1.weight",
44
+ "layers.10.norm2.bias",
45
+ "layers.10.norm2.weight",
46
+ "layers.11.attn.in_proj_bias",
47
+ "layers.11.attn.in_proj_weight",
48
+ "layers.11.attn.out_proj.bias",
49
+ "layers.11.attn.out_proj.weight",
50
+ "layers.11.ffn.layers.0.bias",
51
+ "layers.11.ffn.layers.0.weight",
52
+ "layers.11.ffn.layers.3.bias",
53
+ "layers.11.ffn.layers.3.weight",
54
+ "layers.11.norm1.bias",
55
+ "layers.11.norm1.weight",
56
+ "layers.11.norm2.bias",
57
+ "layers.11.norm2.weight",
58
+ "layers.12.attn.in_proj_bias",
59
+ "layers.12.attn.in_proj_weight",
60
+ "layers.12.attn.out_proj.bias",
61
+ "layers.12.attn.out_proj.weight",
62
+ "layers.12.ffn.layers.0.bias",
63
+ "layers.12.ffn.layers.0.weight",
64
+ "layers.12.ffn.layers.3.bias",
65
+ "layers.12.ffn.layers.3.weight",
66
+ "layers.12.norm1.bias",
67
+ "layers.12.norm1.weight",
68
+ "layers.12.norm2.bias",
69
+ "layers.12.norm2.weight",
70
+ "layers.13.attn.in_proj_bias",
71
+ "layers.13.attn.in_proj_weight",
72
+ "layers.13.attn.out_proj.bias",
73
+ "layers.13.attn.out_proj.weight",
74
+ "layers.13.ffn.layers.0.bias",
75
+ "layers.13.ffn.layers.0.weight",
76
+ "layers.13.ffn.layers.3.bias",
77
+ "layers.13.ffn.layers.3.weight",
78
+ "layers.13.norm1.bias",
79
+ "layers.13.norm1.weight",
80
+ "layers.13.norm2.bias",
81
+ "layers.13.norm2.weight",
82
+ "layers.14.attn.in_proj_bias",
83
+ "layers.14.attn.in_proj_weight",
84
+ "layers.14.attn.out_proj.bias",
85
+ "layers.14.attn.out_proj.weight",
86
+ "layers.14.ffn.layers.0.bias",
87
+ "layers.14.ffn.layers.0.weight",
88
+ "layers.14.ffn.layers.3.bias",
89
+ "layers.14.ffn.layers.3.weight",
90
+ "layers.14.norm1.bias",
91
+ "layers.14.norm1.weight",
92
+ "layers.14.norm2.bias",
93
+ "layers.14.norm2.weight",
94
+ "layers.15.attn.in_proj_bias",
95
+ "layers.15.attn.in_proj_weight",
96
+ "layers.15.attn.out_proj.bias",
97
+ "layers.15.attn.out_proj.weight",
98
+ "layers.15.ffn.layers.0.bias",
99
+ "layers.15.ffn.layers.0.weight",
100
+ "layers.15.ffn.layers.3.bias",
101
+ "layers.15.ffn.layers.3.weight",
102
+ "layers.15.norm1.bias",
103
+ "layers.15.norm1.weight",
104
+ "layers.15.norm2.bias",
105
+ "layers.15.norm2.weight",
106
+ "layers.16.attn.in_proj_bias",
107
+ "layers.16.attn.in_proj_weight",
108
+ "layers.16.attn.out_proj.bias",
109
+ "layers.16.attn.out_proj.weight",
110
+ "layers.16.ffn.layers.0.bias",
111
+ "layers.16.ffn.layers.0.weight",
112
+ "layers.16.ffn.layers.3.bias",
113
+ "layers.16.ffn.layers.3.weight",
114
+ "layers.16.norm1.bias",
115
+ "layers.16.norm1.weight",
116
+ "layers.16.norm2.bias",
117
+ "layers.16.norm2.weight",
118
+ "layers.17.attn.in_proj_bias",
119
+ "layers.17.attn.in_proj_weight",
120
+ "layers.17.attn.out_proj.bias",
121
+ "layers.17.attn.out_proj.weight",
122
+ "layers.17.ffn.layers.0.bias",
123
+ "layers.17.ffn.layers.0.weight",
124
+ "layers.17.ffn.layers.3.bias",
125
+ "layers.17.ffn.layers.3.weight",
126
+ "layers.17.norm1.bias",
127
+ "layers.17.norm1.weight",
128
+ "layers.17.norm2.bias",
129
+ "layers.17.norm2.weight",
130
+ "layers.18.attn.in_proj_bias",
131
+ "layers.18.attn.in_proj_weight",
132
+ "layers.18.attn.out_proj.bias",
133
+ "layers.18.attn.out_proj.weight",
134
+ "layers.18.ffn.layers.0.bias",
135
+ "layers.18.ffn.layers.0.weight",
136
+ "layers.18.ffn.layers.3.bias",
137
+ "layers.18.ffn.layers.3.weight",
138
+ "layers.18.norm1.bias",
139
+ "layers.18.norm1.weight",
140
+ "layers.18.norm2.bias",
141
+ "layers.18.norm2.weight",
142
+ "layers.19.attn.in_proj_bias",
143
+ "layers.19.attn.in_proj_weight",
144
+ "layers.19.attn.out_proj.bias",
145
+ "layers.19.attn.out_proj.weight",
146
+ "layers.19.ffn.layers.0.bias",
147
+ "layers.19.ffn.layers.0.weight",
148
+ "layers.19.ffn.layers.3.bias",
149
+ "layers.19.ffn.layers.3.weight",
150
+ "layers.19.norm1.bias",
151
+ "layers.19.norm1.weight",
152
+ "layers.19.norm2.bias",
153
+ "layers.19.norm2.weight",
154
+ "layers.2.attn.in_proj_bias",
155
+ "layers.2.attn.in_proj_weight",
156
+ "layers.2.attn.out_proj.bias",
157
+ "layers.2.attn.out_proj.weight",
158
+ "layers.2.ffn.layers.0.bias",
159
+ "layers.2.ffn.layers.0.weight",
160
+ "layers.2.ffn.layers.3.bias",
161
+ "layers.2.ffn.layers.3.weight",
162
+ "layers.2.norm1.bias",
163
+ "layers.2.norm1.weight",
164
+ "layers.2.norm2.bias",
165
+ "layers.2.norm2.weight",
166
+ "layers.20.attn.in_proj_bias",
167
+ "layers.20.attn.in_proj_weight",
168
+ "layers.20.attn.out_proj.bias",
169
+ "layers.20.attn.out_proj.weight",
170
+ "layers.20.ffn.layers.0.bias",
171
+ "layers.20.ffn.layers.0.weight",
172
+ "layers.20.ffn.layers.3.bias",
173
+ "layers.20.ffn.layers.3.weight",
174
+ "layers.20.norm1.bias",
175
+ "layers.20.norm1.weight",
176
+ "layers.20.norm2.bias",
177
+ "layers.20.norm2.weight",
178
+ "layers.21.attn.in_proj_bias",
179
+ "layers.21.attn.in_proj_weight",
180
+ "layers.21.attn.out_proj.bias",
181
+ "layers.21.attn.out_proj.weight",
182
+ "layers.21.ffn.layers.0.bias",
183
+ "layers.21.ffn.layers.0.weight",
184
+ "layers.21.ffn.layers.3.bias",
185
+ "layers.21.ffn.layers.3.weight",
186
+ "layers.21.norm1.bias",
187
+ "layers.21.norm1.weight",
188
+ "layers.21.norm2.bias",
189
+ "layers.21.norm2.weight",
190
+ "layers.22.attn.in_proj_bias",
191
+ "layers.22.attn.in_proj_weight",
192
+ "layers.22.attn.out_proj.bias",
193
+ "layers.22.attn.out_proj.weight",
194
+ "layers.22.ffn.layers.0.bias",
195
+ "layers.22.ffn.layers.0.weight",
196
+ "layers.22.ffn.layers.3.bias",
197
+ "layers.22.ffn.layers.3.weight",
198
+ "layers.22.norm1.bias",
199
+ "layers.22.norm1.weight",
200
+ "layers.22.norm2.bias",
201
+ "layers.22.norm2.weight",
202
+ "layers.23.attn.in_proj_bias",
203
+ "layers.23.attn.in_proj_weight",
204
+ "layers.23.attn.out_proj.bias",
205
+ "layers.23.attn.out_proj.weight",
206
+ "layers.23.ffn.layers.0.bias",
207
+ "layers.23.ffn.layers.0.weight",
208
+ "layers.23.ffn.layers.3.bias",
209
+ "layers.23.ffn.layers.3.weight",
210
+ "layers.23.norm1.bias",
211
+ "layers.23.norm1.weight",
212
+ "layers.23.norm2.bias",
213
+ "layers.23.norm2.weight",
214
+ "layers.3.attn.in_proj_bias",
215
+ "layers.3.attn.in_proj_weight",
216
+ "layers.3.attn.out_proj.bias",
217
+ "layers.3.attn.out_proj.weight",
218
+ "layers.3.ffn.layers.0.bias",
219
+ "layers.3.ffn.layers.0.weight",
220
+ "layers.3.ffn.layers.3.bias",
221
+ "layers.3.ffn.layers.3.weight",
222
+ "layers.3.norm1.bias",
223
+ "layers.3.norm1.weight",
224
+ "layers.3.norm2.bias",
225
+ "layers.3.norm2.weight",
226
+ "layers.4.attn.in_proj_bias",
227
+ "layers.4.attn.in_proj_weight",
228
+ "layers.4.attn.out_proj.bias",
229
+ "layers.4.attn.out_proj.weight",
230
+ "layers.4.ffn.layers.0.bias",
231
+ "layers.4.ffn.layers.0.weight",
232
+ "layers.4.ffn.layers.3.bias",
233
+ "layers.4.ffn.layers.3.weight",
234
+ "layers.4.norm1.bias",
235
+ "layers.4.norm1.weight",
236
+ "layers.4.norm2.bias",
237
+ "layers.4.norm2.weight",
238
+ "layers.5.attn.in_proj_bias",
239
+ "layers.5.attn.in_proj_weight",
240
+ "layers.5.attn.out_proj.bias",
241
+ "layers.5.attn.out_proj.weight",
242
+ "layers.5.ffn.layers.0.bias",
243
+ "layers.5.ffn.layers.0.weight",
244
+ "layers.5.ffn.layers.3.bias",
245
+ "layers.5.ffn.layers.3.weight",
246
+ "layers.5.norm1.bias",
247
+ "layers.5.norm1.weight",
248
+ "layers.5.norm2.bias",
249
+ "layers.5.norm2.weight",
250
+ "layers.6.attn.in_proj_bias",
251
+ "layers.6.attn.in_proj_weight",
252
+ "layers.6.attn.out_proj.bias",
253
+ "layers.6.attn.out_proj.weight",
254
+ "layers.6.ffn.layers.0.bias",
255
+ "layers.6.ffn.layers.0.weight",
256
+ "layers.6.ffn.layers.3.bias",
257
+ "layers.6.ffn.layers.3.weight",
258
+ "layers.6.norm1.bias",
259
+ "layers.6.norm1.weight",
260
+ "layers.6.norm2.bias",
261
+ "layers.6.norm2.weight",
262
+ "layers.7.attn.in_proj_bias",
263
+ "layers.7.attn.in_proj_weight",
264
+ "layers.7.attn.out_proj.bias",
265
+ "layers.7.attn.out_proj.weight",
266
+ "layers.7.ffn.layers.0.bias",
267
+ "layers.7.ffn.layers.0.weight",
268
+ "layers.7.ffn.layers.3.bias",
269
+ "layers.7.ffn.layers.3.weight",
270
+ "layers.7.norm1.bias",
271
+ "layers.7.norm1.weight",
272
+ "layers.7.norm2.bias",
273
+ "layers.7.norm2.weight",
274
+ "layers.8.attn.in_proj_bias",
275
+ "layers.8.attn.in_proj_weight",
276
+ "layers.8.attn.out_proj.bias",
277
+ "layers.8.attn.out_proj.weight",
278
+ "layers.8.ffn.layers.0.bias",
279
+ "layers.8.ffn.layers.0.weight",
280
+ "layers.8.ffn.layers.3.bias",
281
+ "layers.8.ffn.layers.3.weight",
282
+ "layers.8.norm1.bias",
283
+ "layers.8.norm1.weight",
284
+ "layers.8.norm2.bias",
285
+ "layers.8.norm2.weight",
286
+ "layers.9.attn.in_proj_bias",
287
+ "layers.9.attn.in_proj_weight",
288
+ "layers.9.attn.out_proj.bias",
289
+ "layers.9.attn.out_proj.weight",
290
+ "layers.9.ffn.layers.0.bias",
291
+ "layers.9.ffn.layers.0.weight",
292
+ "layers.9.ffn.layers.3.bias",
293
+ "layers.9.ffn.layers.3.weight",
294
+ "layers.9.norm1.bias",
295
+ "layers.9.norm1.weight",
296
+ "layers.9.norm2.bias",
297
+ "layers.9.norm2.weight",
298
+ "mask_token",
299
+ "patch_embed.projection.bias",
300
+ "patch_embed.projection.weight",
301
+ "pos_embed",
302
+ "vocabulary_token",
303
+ "vocabulary_weight"
304
+ ]
305
+ }
skysensepp-vit-msl-s2/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:27f2a40bdad5ffc10598d808ec7b5d481f9848c6d6a63155fb4f4f113480c3e5
3
+ size 1210265976
skysensepp-vit-msl-s2/modeling_skysensepp_swinv2_msl.py ADDED
@@ -0,0 +1,343 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ from copy import deepcopy
4
+ from typing import Optional, Sequence, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig
14
+ from .modeling_utils import (
15
+ DropPath,
16
+ FFN,
17
+ PatchEmbed,
18
+ PatchMerging,
19
+ ShiftWindowMSA,
20
+ to_2tuple,
21
+ )
22
+
23
+
24
+ class SwinBlockV2(nn.Module):
25
+ def __init__(
26
+ self,
27
+ embed_dims: int,
28
+ num_heads: int,
29
+ window_size: int = 8,
30
+ shift: bool = False,
31
+ extra_norm: bool = False,
32
+ ffn_ratio: float = 4.0,
33
+ drop_path: float = 0.0,
34
+ pad_small_map: bool = False,
35
+ with_cp: bool = False,
36
+ pretrained_window_size: int = 0,
37
+ ):
38
+ super().__init__()
39
+ self.with_cp = with_cp
40
+ self.extra_norm = extra_norm
41
+ self.attn = ShiftWindowMSA(
42
+ embed_dims=embed_dims,
43
+ num_heads=num_heads,
44
+ window_size=window_size,
45
+ shift_size=window_size // 2 if shift else 0,
46
+ drop_path=drop_path,
47
+ pad_small_map=pad_small_map,
48
+ pretrained_window_size=pretrained_window_size,
49
+ )
50
+ self.norm1 = nn.LayerNorm(embed_dims)
51
+ self.ffn = FFN(
52
+ embed_dims=embed_dims,
53
+ feedforward_channels=int(embed_dims * ffn_ratio),
54
+ num_fcs=2,
55
+ drop_path=drop_path,
56
+ act_layer=nn.GELU,
57
+ add_identity=False,
58
+ )
59
+ self.norm2 = nn.LayerNorm(embed_dims)
60
+ if self.extra_norm:
61
+ self.norm3 = nn.LayerNorm(embed_dims)
62
+
63
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
64
+ def _inner_forward(x):
65
+ identity = x
66
+ x = self.attn(x, hw_shape)
67
+ x = self.norm1(x)
68
+ x = x + identity
69
+
70
+ identity = x
71
+ x = self.ffn(x)
72
+ x = self.norm2(x)
73
+ x = x + identity
74
+
75
+ if self.extra_norm:
76
+ x = self.norm3(x)
77
+ return x
78
+
79
+ if self.with_cp and x.requires_grad:
80
+ x = cp.checkpoint(_inner_forward, x, use_reentrant=False)
81
+ else:
82
+ x = _inner_forward(x)
83
+ return x
84
+
85
+
86
+ class SwinBlockV2Sequence(nn.Module):
87
+ def __init__(
88
+ self,
89
+ embed_dims: int,
90
+ depth: int,
91
+ num_heads: int,
92
+ window_size: int = 8,
93
+ downsample: bool = False,
94
+ drop_paths: Union[Sequence[float], float] = 0.0,
95
+ with_cp: bool = False,
96
+ pad_small_map: bool = False,
97
+ extra_norm_every_n_blocks: int = 0,
98
+ pretrained_window_size: int = 0,
99
+ is_post_norm_downsample: bool = True,
100
+ ):
101
+ super().__init__()
102
+ if not isinstance(drop_paths, Sequence):
103
+ drop_paths = [drop_paths] * depth
104
+
105
+ if downsample:
106
+ self.out_channels = 2 * embed_dims
107
+ self.downsample = PatchMerging(
108
+ in_channels=embed_dims,
109
+ out_channels=self.out_channels,
110
+ is_post_norm=is_post_norm_downsample,
111
+ )
112
+ else:
113
+ self.out_channels = embed_dims
114
+ self.downsample = None
115
+
116
+ self.blocks = nn.ModuleList()
117
+ for i in range(depth):
118
+ extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0
119
+ self.blocks.append(
120
+ SwinBlockV2(
121
+ embed_dims=self.out_channels,
122
+ num_heads=num_heads,
123
+ window_size=window_size,
124
+ shift=(i % 2 == 1),
125
+ extra_norm=extra_norm,
126
+ drop_path=drop_paths[i],
127
+ with_cp=with_cp,
128
+ pad_small_map=pad_small_map,
129
+ pretrained_window_size=pretrained_window_size,
130
+ )
131
+ )
132
+
133
+ def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
134
+ if self.downsample is not None:
135
+ x, out_shape = self.downsample(x, in_shape)
136
+ else:
137
+ out_shape = in_shape
138
+
139
+ for block in self.blocks:
140
+ x = block(x, out_shape)
141
+ return x, out_shape
142
+
143
+
144
+ class ProjMHSA(nn.Module):
145
+ """Projected multi-head self-attention used in SkySense++ HR backbone."""
146
+
147
+ def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True):
148
+ super().__init__()
149
+ self.proj_in = nn.Linear(embed_dims, proj_dims)
150
+ self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias)
151
+ self.proj_out = nn.Linear(proj_dims, embed_dims)
152
+
153
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
154
+ x = self.proj_in(x)
155
+ x, _ = self.attn(x, x, x)
156
+ return self.proj_out(x)
157
+
158
+
159
+ class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel):
160
+ config_class = SkySensePlusPlusSwinV2MSLConfig
161
+ base_model_prefix = "skysensepp_swinv2_msl"
162
+ supports_gradient_checkpointing = True
163
+
164
+ def _init_weights(self, module):
165
+ if isinstance(module, nn.Linear):
166
+ nn.init.trunc_normal_(module.weight, std=0.02)
167
+ if module.bias is not None:
168
+ nn.init.zeros_(module.bias)
169
+ elif isinstance(module, nn.LayerNorm):
170
+ nn.init.ones_(module.weight)
171
+ nn.init.zeros_(module.bias)
172
+ elif isinstance(module, nn.Conv2d):
173
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
174
+ if module.bias is not None:
175
+ nn.init.zeros_(module.bias)
176
+
177
+
178
+ class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel):
179
+ """SkySense++ HR backbone with semantic vocabulary and annotation conditioning."""
180
+
181
+ def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig):
182
+ super().__init__(config)
183
+
184
+ self.num_layers = len(config.depths)
185
+ self.out_indices = config.out_indices
186
+ self.merge_stage = config.merge_stage
187
+ self.use_attn = config.use_attn
188
+ self.patch_size = config.patch_size
189
+
190
+ if isinstance(config.window_size, int):
191
+ window_sizes = [config.window_size] * self.num_layers
192
+ else:
193
+ window_sizes = list(config.window_size)
194
+
195
+ self.patch_embed = PatchEmbed(
196
+ in_channels=config.in_channels,
197
+ embed_dims=config.embed_dims,
198
+ kernel_size=config.patch_size,
199
+ stride=config.patch_size,
200
+ norm_layer=nn.LayerNorm,
201
+ input_size=config.img_size,
202
+ )
203
+
204
+ self.use_abs_pos_embed = config.use_abs_pos_embed
205
+ if self.use_abs_pos_embed:
206
+ patch_resolution = self.patch_embed.init_out_size
207
+ num_patches = patch_resolution[0] * patch_resolution[1]
208
+ self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
209
+
210
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
211
+
212
+ total_depth = sum(config.depths)
213
+ if total_depth > 1:
214
+ dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)]
215
+ else:
216
+ dpr = [0.0]
217
+
218
+ self.stages = nn.ModuleList()
219
+ embed_dims_list = [config.embed_dims]
220
+ for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)):
221
+ stage = SwinBlockV2Sequence(
222
+ embed_dims=embed_dims_list[-1],
223
+ depth=depth,
224
+ num_heads=num_heads,
225
+ window_size=window_sizes[i],
226
+ downsample=(i > 0),
227
+ drop_paths=dpr[:depth],
228
+ with_cp=config.with_cp,
229
+ pad_small_map=config.pad_small_map,
230
+ extra_norm_every_n_blocks=config.extra_norm_every_n_blocks,
231
+ pretrained_window_size=config.pretrained_window_sizes[i],
232
+ is_post_norm_downsample=config.is_post_norm_downsample,
233
+ )
234
+ self.stages.append(stage)
235
+ dpr = dpr[depth:]
236
+ embed_dims_list.append(stage.out_channels)
237
+
238
+ for i in self.out_indices:
239
+ self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1]))
240
+
241
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
242
+ self.vocabulary_token = nn.Parameter(
243
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
244
+ )
245
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
246
+
247
+ if self.use_attn:
248
+ self.attn1 = ProjMHSA(352, 256, num_heads=16)
249
+ self.attn2 = ProjMHSA(704, 512, num_heads=16)
250
+ self.attn3 = ProjMHSA(1408, 1024, num_heads=16)
251
+ self.norm_attn = nn.LayerNorm(1408)
252
+
253
+ self.post_init()
254
+
255
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
256
+ batch_size, height, width = anno_img.shape
257
+ ann_token = torch.index_select(
258
+ self.vocabulary_token, 0, anno_img.reshape(-1)
259
+ ).reshape(batch_size, height, width, -1)
260
+
261
+ num_patch_h = height // self.patch_size
262
+ num_patch_w = width // self.patch_size
263
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
264
+ weight = (
265
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
266
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
267
+ .reshape(1, height, width, 1)
268
+ )
269
+ ann_token = ann_token * weight
270
+ ann_token = F.avg_pool2d(
271
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
272
+ )
273
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
274
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
275
+ )
276
+
277
+ def forward(
278
+ self,
279
+ pixel_values: torch.Tensor,
280
+ annotation: torch.Tensor,
281
+ mask: Optional[torch.Tensor] = None,
282
+ output_hidden_states: Optional[bool] = None,
283
+ return_dict: Optional[bool] = None,
284
+ ) -> Union[Tuple, BaseModelOutput]:
285
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
286
+
287
+ x, hw_shape = self.patch_embed(pixel_values)
288
+ y = self.create_ann_token(annotation)
289
+ batch_size, num_tokens, channels = y.shape
290
+
291
+ if mask is not None:
292
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
293
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
294
+ y = y * (1.0 - weight) + mask_tokens * weight
295
+
296
+ if self.merge_stage == 0:
297
+ x = (x + y) * 0.5
298
+ else:
299
+ x = x.reshape(batch_size, *hw_shape, channels)
300
+ y = y.reshape(batch_size, *hw_shape, channels)
301
+ x = torch.cat((x, y), dim=2)
302
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
303
+ x = x.reshape(batch_size, -1, channels)
304
+
305
+ if self.use_abs_pos_embed:
306
+ x = x + self.absolute_pos_embed
307
+ x = self.drop_after_pos(x)
308
+
309
+ all_hidden_states = () if output_hidden_states else None
310
+ feature_maps = []
311
+ merge_idx = self.merge_stage - 1
312
+
313
+ for i, stage in enumerate(self.stages):
314
+ x, hw_shape = stage(x, hw_shape)
315
+ if i == merge_idx:
316
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
317
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
318
+ x = x.reshape(batch_size, -1, x.shape[-1])
319
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
320
+
321
+ if self.use_attn:
322
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
323
+ if i <= len(attention_blocks) - 1:
324
+ x = x + attention_blocks[i](x)
325
+ if i == len(attention_blocks) - 1:
326
+ x = self.norm_attn(x)
327
+
328
+ if output_hidden_states:
329
+ all_hidden_states = all_hidden_states + (x,)
330
+
331
+ if i in self.out_indices:
332
+ norm_layer = getattr(self, f"norm{i}")
333
+ out = norm_layer(x)
334
+ out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous()
335
+ feature_maps.append(out)
336
+
337
+ if not return_dict:
338
+ return tuple(feature_maps)
339
+
340
+ return BaseModelOutput(
341
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
342
+ hidden_states=all_hidden_states,
343
+ )
skysensepp-vit-msl-s2/modeling_skysensepp_vit_msl.py ADDED
@@ -0,0 +1,265 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace)."""
2
+
3
+ import math
4
+ from typing import Optional, Tuple, Union
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ import torch.nn.functional as F
9
+ import torch.utils.checkpoint as cp
10
+ from transformers import PreTrainedModel
11
+ from transformers.modeling_outputs import BaseModelOutput
12
+
13
+ from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig
14
+ from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple
15
+
16
+
17
+ class TransformerEncoderLayer(nn.Module):
18
+ def __init__(
19
+ self,
20
+ embed_dims: int,
21
+ num_heads: int,
22
+ feedforward_channels: int,
23
+ drop_rate: float = 0.0,
24
+ attn_drop_rate: float = 0.0,
25
+ drop_path_rate: float = 0.0,
26
+ num_fcs: int = 2,
27
+ qkv_bias: bool = True,
28
+ with_cp: bool = False,
29
+ ):
30
+ super().__init__()
31
+ self.with_cp = with_cp
32
+ self.norm1 = nn.LayerNorm(embed_dims)
33
+ self.attn = nn.MultiheadAttention(
34
+ embed_dim=embed_dims,
35
+ num_heads=num_heads,
36
+ dropout=attn_drop_rate,
37
+ bias=qkv_bias,
38
+ batch_first=True,
39
+ )
40
+ self.proj_drop = nn.Dropout(drop_rate)
41
+ self.norm2 = nn.LayerNorm(embed_dims)
42
+ self.ffn = FFN(
43
+ embed_dims=embed_dims,
44
+ feedforward_channels=feedforward_channels,
45
+ num_fcs=num_fcs,
46
+ ffn_drop=drop_rate,
47
+ drop_path=drop_path_rate,
48
+ act_layer=nn.GELU,
49
+ add_identity=True,
50
+ )
51
+ self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
52
+
53
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
54
+ def _inner_forward(x):
55
+ residual = x
56
+ x_norm = self.norm1(x)
57
+ attn_out, _ = self.attn(x_norm, x_norm, x_norm)
58
+ attn_out = self.proj_drop(attn_out)
59
+ x = residual + self.drop_path(attn_out)
60
+ return self.ffn(self.norm2(x), identity=x)
61
+
62
+ if self.with_cp and x.requires_grad:
63
+ return cp.checkpoint(_inner_forward, x, use_reentrant=False)
64
+ return _inner_forward(x)
65
+
66
+
67
+ class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel):
68
+ config_class = SkySensePlusPlusViTMSLConfig
69
+ base_model_prefix = "skysensepp_vit_msl"
70
+ supports_gradient_checkpointing = True
71
+
72
+ def _init_weights(self, module):
73
+ if isinstance(module, nn.Linear):
74
+ nn.init.trunc_normal_(module.weight, std=0.02)
75
+ if module.bias is not None:
76
+ nn.init.zeros_(module.bias)
77
+ elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)):
78
+ nn.init.ones_(module.weight)
79
+ nn.init.zeros_(module.bias)
80
+ elif isinstance(module, nn.Conv2d):
81
+ nn.init.kaiming_normal_(module.weight, mode="fan_in")
82
+ if module.bias is not None:
83
+ nn.init.zeros_(module.bias)
84
+
85
+
86
+ class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel):
87
+ """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning."""
88
+
89
+ def __init__(self, config: SkySensePlusPlusViTMSLConfig):
90
+ super().__init__(config)
91
+
92
+ img_size = to_2tuple(config.img_size)
93
+ self.img_size = img_size
94
+ self.patch_size = config.patch_size
95
+ self.with_cls_token = config.with_cls_token
96
+ self.output_cls_token = config.output_cls_token
97
+ self.merge_stage = config.merge_stage
98
+ self.use_attn = config.use_attn
99
+ self.interpolate_mode = "bicubic"
100
+
101
+ self.patch_embed = PatchEmbed(
102
+ in_channels=config.in_channels,
103
+ embed_dims=config.embed_dims,
104
+ kernel_size=config.patch_size,
105
+ stride=config.patch_size,
106
+ norm_layer=nn.LayerNorm if config.patch_norm else None,
107
+ )
108
+
109
+ num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size)
110
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
111
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims))
112
+ self.drop_after_pos = nn.Dropout(p=config.drop_rate)
113
+
114
+ out_indices = list(config.out_indices)
115
+ self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices]
116
+
117
+ num_layers = config.num_layers
118
+ if num_layers > 1:
119
+ dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)]
120
+ else:
121
+ dpr = [0.0]
122
+
123
+ self.layers = nn.ModuleList()
124
+ for i in range(config.num_layers):
125
+ self.layers.append(
126
+ TransformerEncoderLayer(
127
+ embed_dims=config.embed_dims,
128
+ num_heads=config.num_heads,
129
+ feedforward_channels=config.mlp_ratio * config.embed_dims,
130
+ attn_drop_rate=config.attn_drop_rate,
131
+ drop_rate=config.drop_rate,
132
+ drop_path_rate=dpr[i],
133
+ num_fcs=2,
134
+ qkv_bias=config.qkv_bias,
135
+ with_cp=config.with_cp,
136
+ )
137
+ )
138
+
139
+ self.final_norm = config.final_norm
140
+ if config.final_norm:
141
+ self.norm = nn.LayerNorm(config.embed_dims)
142
+
143
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims))
144
+ self.vocabulary_token = nn.Parameter(
145
+ torch.zeros(config.num_vocabulary_tokens, config.embed_dims)
146
+ )
147
+ self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size))
148
+
149
+ if self.use_attn:
150
+ self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
151
+ self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
152
+ self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True)
153
+ self.norm_attn = nn.LayerNorm(config.embed_dims)
154
+
155
+ self.post_init()
156
+
157
+ @staticmethod
158
+ def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"):
159
+ pos_h, pos_w = pos_shape
160
+ pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :]
161
+ pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
162
+ pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode)
163
+ return torch.flatten(pos_embed_weight, 2).transpose(1, 2)
164
+
165
+ def _pos_embedding(self, patched_img, hw_shape, pos_embed):
166
+ x_len, pos_len = patched_img.shape[1], pos_embed.shape[1]
167
+ if x_len != pos_len:
168
+ pos_h = self.img_size[0] // self.patch_size
169
+ pos_w = self.img_size[1] // self.patch_size
170
+ pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode)
171
+ return self.drop_after_pos(patched_img + pos_embed)
172
+
173
+ def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor:
174
+ batch_size, height, width = anno_img.shape
175
+ ann_token = torch.index_select(
176
+ self.vocabulary_token, 0, anno_img.reshape(-1)
177
+ ).reshape(batch_size, height, width, -1)
178
+
179
+ num_patch_h = height // self.patch_size
180
+ num_patch_w = width // self.patch_size
181
+ weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size
182
+ weight = (
183
+ weight.reshape(1, 1, self.patch_size, 1, self.patch_size)
184
+ .repeat(1, num_patch_h, 1, num_patch_w, 1)
185
+ .reshape(1, height, width, 1)
186
+ )
187
+ ann_token = ann_token * weight
188
+ ann_token = F.avg_pool2d(
189
+ torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size
190
+ )
191
+ return torch.einsum("bchw->bhwc", ann_token).reshape(
192
+ batch_size, num_patch_h * num_patch_w, self.config.embed_dims
193
+ )
194
+
195
+ def forward(
196
+ self,
197
+ pixel_values: torch.Tensor,
198
+ annotation: torch.Tensor,
199
+ mask: Optional[torch.Tensor] = None,
200
+ output_hidden_states: Optional[bool] = None,
201
+ return_dict: Optional[bool] = None,
202
+ ) -> Union[Tuple, BaseModelOutput]:
203
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
204
+
205
+ x, hw_shape = self.patch_embed(pixel_values)
206
+ y = self.create_ann_token(annotation)
207
+ batch_size, num_tokens, channels = y.shape
208
+
209
+ if mask is not None:
210
+ mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1)
211
+ weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens)
212
+ y = y * (1.0 - weight) + mask_tokens * weight
213
+
214
+ if self.merge_stage == 0:
215
+ x = (x + y) * 0.5
216
+ else:
217
+ x = x.reshape(batch_size, *hw_shape, channels)
218
+ y = y.reshape(batch_size, *hw_shape, channels)
219
+ x = torch.cat((x, y), dim=2)
220
+ hw_shape = (hw_shape[0], hw_shape[1] * 2)
221
+ x = x.reshape(batch_size, -1, channels)
222
+
223
+ x = self._pos_embedding(x, hw_shape, self.pos_embed)
224
+
225
+ all_hidden_states = () if output_hidden_states else None
226
+ feature_maps = []
227
+ merge_idx = self.merge_stage - 1
228
+
229
+ for i, layer in enumerate(self.layers):
230
+ x = layer(x)
231
+
232
+ if i == merge_idx:
233
+ x = x.reshape(batch_size, *hw_shape, x.shape[-1])
234
+ x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5
235
+ x = x.reshape(batch_size, -1, x.shape[-1])
236
+ hw_shape = (hw_shape[0], hw_shape[1] // 2)
237
+
238
+ if self.use_attn:
239
+ attention_blocks = [self.attn1, self.attn2, self.attn3]
240
+ if i <= len(attention_blocks) - 1:
241
+ attn_out, _ = attention_blocks[i](x, x, x)
242
+ x = x + attn_out
243
+ if i == len(attention_blocks) - 1:
244
+ x = self.norm_attn(x)
245
+
246
+ if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm:
247
+ x = self.norm(x)
248
+
249
+ if output_hidden_states:
250
+ all_hidden_states = all_hidden_states + (x,)
251
+
252
+ if i in self.out_indices:
253
+ out = x
254
+ out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous()
255
+ if self.output_cls_token:
256
+ out = [out, x[:, 0]]
257
+ feature_maps.append(out)
258
+
259
+ if not return_dict:
260
+ return tuple(feature_maps)
261
+
262
+ return BaseModelOutput(
263
+ last_hidden_state=feature_maps[-1] if feature_maps else x,
264
+ hidden_states=all_hidden_states,
265
+ )
skysensepp-vit-msl-s2/modeling_utils.py ADDED
@@ -0,0 +1,557 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SkySense: Pure PyTorch + HuggingFace Transformers implementation.
2
+
3
+ Shared utility modules used across SkySense model implementations.
4
+ """
5
+
6
+ import math
7
+ from typing import Optional, Tuple
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+
14
+ def to_2tuple(x):
15
+ """Convert to a 2-tuple."""
16
+ if isinstance(x, (list, tuple)):
17
+ return tuple(x)
18
+ return (x, x)
19
+
20
+
21
+ class DropPath(nn.Module):
22
+ """Drop paths (stochastic depth) per sample.
23
+
24
+ Args:
25
+ drop_prob (float): Probability of dropping a path. Default: 0.0.
26
+ """
27
+
28
+ def __init__(self, drop_prob: float = 0.0):
29
+ super().__init__()
30
+ self.drop_prob = drop_prob
31
+
32
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
33
+ if self.drop_prob == 0.0 or not self.training:
34
+ return x
35
+ keep_prob = 1 - self.drop_prob
36
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1)
37
+ random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device)
38
+ random_tensor = torch.floor(random_tensor + keep_prob)
39
+ output = x / keep_prob * random_tensor
40
+ return output
41
+
42
+
43
+ class PatchEmbed(nn.Module):
44
+ """Image to Patch Embedding using Conv2d.
45
+
46
+ Args:
47
+ in_channels (int): Number of input channels. Default: 3.
48
+ embed_dims (int): Embedding dimension. Default: 96.
49
+ kernel_size (int): Kernel size of the projection. Default: 4.
50
+ stride (int): Stride of the projection. Default: 4.
51
+ padding (int): Padding of the projection. Default: 0.
52
+ norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm.
53
+ input_size (int or tuple or None): Input resolution for calculating output size.
54
+ """
55
+
56
+ def __init__(
57
+ self,
58
+ in_channels: int = 3,
59
+ embed_dims: int = 96,
60
+ kernel_size: int = 4,
61
+ stride: int = 4,
62
+ padding: int = 0,
63
+ norm_layer: Optional[type] = nn.LayerNorm,
64
+ input_size: Optional[int] = None,
65
+ ):
66
+ super().__init__()
67
+ self.projection = nn.Conv2d(
68
+ in_channels, embed_dims,
69
+ kernel_size=kernel_size, stride=stride, padding=padding,
70
+ )
71
+ self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity()
72
+
73
+ # Compute init output size if input_size is given
74
+ if input_size is not None:
75
+ input_size = to_2tuple(input_size)
76
+ self.init_out_size = (
77
+ (input_size[0] - kernel_size + 2 * padding) // stride + 1,
78
+ (input_size[1] - kernel_size + 2 * padding) // stride + 1,
79
+ )
80
+ else:
81
+ self.init_out_size = None
82
+
83
+ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
84
+ x = self.projection(x) # (B, C, H, W)
85
+ out_size = (x.shape[2], x.shape[3])
86
+ x = x.flatten(2).transpose(1, 2) # (B, H*W, C)
87
+ x = self.norm(x)
88
+ return x, out_size
89
+
90
+
91
+ class FFN(nn.Module):
92
+ """Feed-Forward Network.
93
+
94
+ Args:
95
+ embed_dims (int): Input dimension.
96
+ feedforward_channels (int): Hidden dimension.
97
+ num_fcs (int): Number of FC layers. Default: 2.
98
+ ffn_drop (float): Dropout rate. Default: 0.0.
99
+ drop_path (float): Drop path rate. Default: 0.0.
100
+ act_layer (nn.Module): Activation layer class. Default: nn.GELU.
101
+ add_identity (bool): Whether to add identity connection. Default: True.
102
+ """
103
+
104
+ def __init__(
105
+ self,
106
+ embed_dims: int,
107
+ feedforward_channels: int,
108
+ num_fcs: int = 2,
109
+ ffn_drop: float = 0.0,
110
+ drop_path: float = 0.0,
111
+ act_layer: type = nn.GELU,
112
+ add_identity: bool = True,
113
+ ):
114
+ super().__init__()
115
+ assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}"
116
+ self.embed_dims = embed_dims
117
+ self.feedforward_channels = feedforward_channels
118
+ self.add_identity = add_identity
119
+
120
+ layers = []
121
+ in_channels = embed_dims
122
+ for i in range(num_fcs - 1):
123
+ layers.append(nn.Linear(in_channels, feedforward_channels))
124
+ layers.append(act_layer())
125
+ layers.append(nn.Dropout(ffn_drop))
126
+ in_channels = feedforward_channels
127
+ layers.append(nn.Linear(feedforward_channels, embed_dims))
128
+ layers.append(nn.Dropout(ffn_drop))
129
+ self.layers = nn.Sequential(*layers)
130
+
131
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
132
+
133
+ def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor:
134
+ out = self.layers(x)
135
+ out = self.drop_path(out)
136
+ if self.add_identity:
137
+ if identity is None:
138
+ identity = x
139
+ out = out + identity
140
+ return out
141
+
142
+
143
+ class WindowMSAV2(nn.Module):
144
+ """Window-based Multi-head Self-Attention for Swin Transformer V2.
145
+
146
+ Uses cosine attention and log-spaced continuous position bias (log-CPB).
147
+
148
+ Args:
149
+ embed_dims (int): Number of input channels.
150
+ num_heads (int): Number of attention heads.
151
+ window_size (tuple[int]): Window size (Wh, Ww).
152
+ pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0).
153
+ qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True.
154
+ attn_drop (float): Attention dropout rate. Default: 0.0.
155
+ proj_drop (float): Output projection dropout rate. Default: 0.0.
156
+ """
157
+
158
+ def __init__(
159
+ self,
160
+ embed_dims: int,
161
+ num_heads: int,
162
+ window_size: Tuple[int, int],
163
+ pretrained_window_size: Tuple[int, int] = (0, 0),
164
+ qkv_bias: bool = True,
165
+ attn_drop: float = 0.0,
166
+ proj_drop: float = 0.0,
167
+ ):
168
+ super().__init__()
169
+ self.embed_dims = embed_dims
170
+ self.num_heads = num_heads
171
+ self.window_size = window_size
172
+ self.pretrained_window_size = pretrained_window_size
173
+
174
+ self.logit_scale = nn.Parameter(
175
+ torch.log(10 * torch.ones((num_heads, 1, 1))))
176
+
177
+ # MLP for continuous relative position bias (log-CPB)
178
+ self.cpb_mlp = nn.Sequential(
179
+ nn.Linear(2, 512, bias=True),
180
+ nn.ReLU(inplace=True),
181
+ nn.Linear(512, num_heads, bias=False),
182
+ )
183
+
184
+ # Build relative coords table
185
+ self._build_relative_coords_table()
186
+ # Build relative position index
187
+ self._build_relative_position_index()
188
+
189
+ self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False)
190
+ if qkv_bias:
191
+ self.q_bias = nn.Parameter(torch.zeros(embed_dims))
192
+ self.v_bias = nn.Parameter(torch.zeros(embed_dims))
193
+ else:
194
+ self.q_bias = None
195
+ self.v_bias = None
196
+
197
+ self.attn_drop = nn.Dropout(attn_drop)
198
+ self.proj = nn.Linear(embed_dims, embed_dims)
199
+ self.proj_drop = nn.Dropout(proj_drop)
200
+ self.softmax = nn.Softmax(dim=-1)
201
+
202
+ def _build_relative_coords_table(self):
203
+ """Build the relative coordinates table for log-CPB."""
204
+ Wh, Ww = self.window_size
205
+ # Table of relative coordinates
206
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32)
207
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32)
208
+ coords_table = torch.stack(
209
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
210
+ ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2)
211
+
212
+ # Normalize to [-1, 1] and apply log-scale
213
+ if self.pretrained_window_size[0] > 0:
214
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
215
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
216
+ else:
217
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
218
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
219
+ coords_table *= 8 # normalize to -8, 8
220
+ coords_table = (
221
+ torch.sign(coords_table)
222
+ * torch.log2(torch.abs(coords_table) + 1.0)
223
+ / math.log2(8)
224
+ )
225
+ self.register_buffer("relative_coords_table", coords_table)
226
+
227
+ def _build_relative_position_index(self):
228
+ """Build the pairwise relative position index for each window token."""
229
+ Wh, Ww = self.window_size
230
+ coords_h = torch.arange(Wh)
231
+ coords_w = torch.arange(Ww)
232
+ coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij'))
233
+ coords_flatten = coords.view(2, -1)
234
+
235
+ relative_coords = (
236
+ coords_flatten[:, :, None] - coords_flatten[:, None, :]
237
+ ) # (2, Wh*Ww, Wh*Ww)
238
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
239
+ relative_coords[:, :, 0] += Wh - 1
240
+ relative_coords[:, :, 1] += Ww - 1
241
+ relative_coords[:, :, 0] *= 2 * Ww - 1
242
+ relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww)
243
+ self.register_buffer("relative_position_index", relative_position_index)
244
+
245
+ def _compute_position_bias(self, N):
246
+ """Compute relative position bias, supporting dynamic window sizes.
247
+
248
+ The log-CPB (Continuous Position Bias) MLP can generalize to any window
249
+ size by computing bias from normalized relative coordinates.
250
+ """
251
+ init_N = self.window_size[0] * self.window_size[1]
252
+ if N == init_N:
253
+ # Use pre-built tables
254
+ relative_position_bias_table = self.cpb_mlp(
255
+ self.relative_coords_table
256
+ ).view(-1, self.num_heads)
257
+ relative_position_bias = relative_position_bias_table[
258
+ self.relative_position_index.view(-1)
259
+ ].view(N, N, -1)
260
+ else:
261
+ # Dynamic: compute for actual window size on-the-fly
262
+ Wh = Ww = int(math.sqrt(N))
263
+ coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device)
264
+ coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device)
265
+ coords_table = torch.stack(
266
+ torch.meshgrid(coords_h, coords_w, indexing='ij')
267
+ ).flatten(1).transpose(0, 1).unsqueeze(0)
268
+ if self.pretrained_window_size[0] > 0:
269
+ coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1)
270
+ coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1)
271
+ else:
272
+ coords_table[:, :, 0] /= max(Wh - 1, 1)
273
+ coords_table[:, :, 1] /= max(Ww - 1, 1)
274
+ coords_table *= 8
275
+ coords_table = (
276
+ torch.sign(coords_table)
277
+ * torch.log2(torch.abs(coords_table) + 1.0)
278
+ / math.log2(8)
279
+ )
280
+ # Build position index for actual window size
281
+ ch = torch.arange(Wh, device=self.logit_scale.device)
282
+ cw = torch.arange(Ww, device=self.logit_scale.device)
283
+ coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij'))
284
+ coords_flat = coords.view(2, -1)
285
+ rel = coords_flat[:, :, None] - coords_flat[:, None, :]
286
+ rel = rel.permute(1, 2, 0).contiguous()
287
+ rel[:, :, 0] += Wh - 1
288
+ rel[:, :, 1] += Ww - 1
289
+ rel[:, :, 0] *= 2 * Ww - 1
290
+ pos_index = rel.sum(-1)
291
+
292
+ bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads)
293
+ relative_position_bias = bias_table[
294
+ pos_index.view(-1)
295
+ ].view(N, N, -1)
296
+
297
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
298
+ relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
299
+ return relative_position_bias
300
+
301
+ def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
302
+ """
303
+ Args:
304
+ x: (num_windows*B, N, C) where N = Wh*Ww
305
+ mask: (num_windows, N, N) or None
306
+ """
307
+ B_, N, C = x.shape
308
+
309
+ # Compute QKV with bias
310
+ if self.q_bias is not None:
311
+ qkv_bias = torch.cat(
312
+ (self.q_bias,
313
+ torch.zeros_like(self.v_bias, requires_grad=False),
314
+ self.v_bias))
315
+ qkv = F.linear(x, self.qkv.weight, qkv_bias)
316
+ else:
317
+ qkv = self.qkv(x)
318
+
319
+ qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads)
320
+ qkv = qkv.permute(2, 0, 3, 1, 4)
321
+ q, k, v = qkv.unbind(0)
322
+
323
+ # Cosine attention
324
+ attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)
325
+ logit_scale = torch.clamp(
326
+ self.logit_scale, max=math.log(1.0 / 0.01)
327
+ ).exp()
328
+ attn = attn * logit_scale
329
+
330
+ # Log-CPB relative position bias (supports dynamic window sizes)
331
+ relative_position_bias = self._compute_position_bias(N)
332
+ attn = attn + relative_position_bias.unsqueeze(0)
333
+
334
+ if mask is not None:
335
+ nW = mask.shape[0]
336
+ attn = attn.view(B_ // nW, nW, self.num_heads, N, N)
337
+ attn = attn + mask.unsqueeze(1).unsqueeze(0)
338
+ attn = attn.view(-1, self.num_heads, N, N)
339
+
340
+ attn = self.softmax(attn)
341
+ attn = self.attn_drop(attn)
342
+
343
+ x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
344
+ x = self.proj(x)
345
+ x = self.proj_drop(x)
346
+ return x
347
+
348
+
349
+ class ShiftWindowMSA(nn.Module):
350
+ """Shifted Window Multi-head Self-Attention.
351
+
352
+ Args:
353
+ embed_dims (int): Number of input channels.
354
+ num_heads (int): Number of attention heads.
355
+ window_size (int): Window size.
356
+ shift_size (int): Shift size for SW-MSA. Default: 0.
357
+ attn_drop (float): Attention dropout rate. Default: 0.0.
358
+ proj_drop (float): Projection dropout rate. Default: 0.0.
359
+ drop_path (float): Drop path rate. Default: 0.0.
360
+ pad_small_map (bool): Pad small feature maps to window size. Default: False.
361
+ pretrained_window_size (int): Pretrained window size. Default: 0.
362
+ """
363
+
364
+ def __init__(
365
+ self,
366
+ embed_dims: int,
367
+ num_heads: int,
368
+ window_size: int,
369
+ shift_size: int = 0,
370
+ attn_drop: float = 0.0,
371
+ proj_drop: float = 0.0,
372
+ drop_path: float = 0.0,
373
+ pad_small_map: bool = False,
374
+ pretrained_window_size: int = 0,
375
+ ):
376
+ super().__init__()
377
+ self.window_size = window_size
378
+ self.shift_size = shift_size
379
+ self.pad_small_map = pad_small_map
380
+
381
+ self.w_msa = WindowMSAV2(
382
+ embed_dims=embed_dims,
383
+ num_heads=num_heads,
384
+ window_size=to_2tuple(window_size),
385
+ pretrained_window_size=to_2tuple(pretrained_window_size),
386
+ attn_drop=attn_drop,
387
+ proj_drop=proj_drop,
388
+ )
389
+ self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
390
+
391
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor:
392
+ B, L, C = x.shape
393
+ H, W = hw_shape
394
+ assert L == H * W, f"Input length {L} != H*W ({H}*{W})"
395
+
396
+ x = x.view(B, H, W, C)
397
+
398
+ window_size = self.window_size
399
+ shift_size = self.shift_size
400
+
401
+ # Pad or shrink window
402
+ if self.pad_small_map:
403
+ pad_r = (window_size - W % window_size) % window_size
404
+ pad_b = (window_size - H % window_size) % window_size
405
+ x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
406
+ _, Hp, Wp, _ = x.shape
407
+ else:
408
+ Hp, Wp = H, W
409
+ if window_size > Hp:
410
+ window_size = Hp
411
+ shift_size = 0
412
+ if window_size > Wp:
413
+ window_size = Wp
414
+ shift_size = 0
415
+
416
+ # Compute attention mask for SW-MSA
417
+ attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device)
418
+
419
+ # Cyclic shift
420
+ if shift_size > 0:
421
+ x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2))
422
+
423
+ # Partition windows
424
+ x_windows = self._window_partition(x, window_size)
425
+ # (num_windows*B, window_size*window_size, C)
426
+
427
+ # W-MSA/SW-MSA
428
+ attn_windows = self.w_msa(x_windows, mask=attn_mask)
429
+
430
+ # Merge windows
431
+ x = self._window_reverse(attn_windows, window_size, Hp, Wp)
432
+
433
+ # Reverse cyclic shift
434
+ if shift_size > 0:
435
+ x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2))
436
+
437
+ if self.pad_small_map and (pad_r > 0 or pad_b > 0):
438
+ x = x[:, :H, :W, :].contiguous()
439
+
440
+ x = x.view(B, H * W, C)
441
+ x = self.drop_path(x)
442
+ return x
443
+
444
+ @staticmethod
445
+ def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor:
446
+ """Partition into non-overlapping windows."""
447
+ B, H, W, C = x.shape
448
+ x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
449
+ windows = x.permute(0, 1, 3, 2, 4, 5).contiguous()
450
+ windows = windows.view(-1, window_size * window_size, C)
451
+ return windows
452
+
453
+ @staticmethod
454
+ def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor:
455
+ """Reverse window partition."""
456
+ B_nW = windows.shape[0]
457
+ nH = H // window_size
458
+ nW = W // window_size
459
+ B = B_nW // (nH * nW)
460
+ x = windows.view(B, nH, nW, window_size, window_size, -1)
461
+ x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
462
+ x = x.view(B, H, W, -1)
463
+ return x
464
+
465
+ @staticmethod
466
+ def _compute_attn_mask(H, W, window_size, shift_size, device):
467
+ """Compute attention mask for shifted window attention."""
468
+ if shift_size <= 0:
469
+ return None
470
+ img_mask = torch.zeros((1, H, W, 1), device=device)
471
+ h_slices = (
472
+ slice(0, -window_size),
473
+ slice(-window_size, -shift_size),
474
+ slice(-shift_size, None),
475
+ )
476
+ w_slices = (
477
+ slice(0, -window_size),
478
+ slice(-window_size, -shift_size),
479
+ slice(-shift_size, None),
480
+ )
481
+ cnt = 0
482
+ for h in h_slices:
483
+ for w in w_slices:
484
+ img_mask[:, h, w, :] = cnt
485
+ cnt += 1
486
+
487
+ # Partition mask
488
+ mask_windows = img_mask.view(
489
+ 1, H // window_size, window_size, W // window_size, window_size, 1
490
+ )
491
+ mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous()
492
+ mask_windows = mask_windows.view(-1, window_size * window_size)
493
+
494
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
495
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0)
496
+ attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0)
497
+ return attn_mask
498
+
499
+
500
+ class PatchMerging(nn.Module):
501
+ """Patch Merging Layer for downsampling (2x).
502
+
503
+ Args:
504
+ in_channels (int): Input channels.
505
+ out_channels (int): Output channels.
506
+ norm_layer (type): Normalization layer. Default: nn.LayerNorm.
507
+ is_post_norm (bool): Apply norm after linear. Default: True.
508
+ """
509
+
510
+ def __init__(
511
+ self,
512
+ in_channels: int,
513
+ out_channels: int,
514
+ norm_layer: type = nn.LayerNorm,
515
+ is_post_norm: bool = True,
516
+ ):
517
+ super().__init__()
518
+ self.in_channels = in_channels
519
+ self.out_channels = out_channels
520
+ self.is_post_norm = is_post_norm
521
+ self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False)
522
+ if is_post_norm:
523
+ self.norm = norm_layer(out_channels)
524
+ else:
525
+ self.norm = norm_layer(4 * in_channels)
526
+
527
+ def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]:
528
+ B, L, C = x.shape
529
+ H, W = hw_shape
530
+ assert L == H * W
531
+
532
+ x = x.view(B, H, W, C)
533
+
534
+ # Pad if needed
535
+ pad_h = H % 2
536
+ pad_w = W % 2
537
+ if pad_h or pad_w:
538
+ x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
539
+
540
+ x0 = x[:, 0::2, 0::2, :]
541
+ x1 = x[:, 1::2, 0::2, :]
542
+ x2 = x[:, 0::2, 1::2, :]
543
+ x3 = x[:, 1::2, 1::2, :]
544
+ x = torch.cat([x0, x1, x2, x3], dim=-1)
545
+
546
+ out_h = (H + pad_h) // 2
547
+ out_w = (W + pad_w) // 2
548
+ x = x.view(B, out_h * out_w, 4 * C)
549
+
550
+ if self.is_post_norm:
551
+ x = self.reduction(x)
552
+ x = self.norm(x)
553
+ else:
554
+ x = self.norm(x)
555
+ x = self.reduction(x)
556
+
557
+ return x, (out_h, out_w)