diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..79d236351ff7b50dd5db782f4fc36c253ca49454 --- /dev/null +++ b/README.md @@ -0,0 +1,140 @@ +--- +license: apache-2.0 +tags: + - remote-sensing + - earth-observation + - skysensepp + - feature-extraction +pipeline_tag: feature-extraction +--- + +# SkySense++ Transformers + +HuggingFace-compatible checkpoints for SkySense++ zero-shot MSL backbones, converted from the official release weights. + +## Checkpoints + +| Directory | Modality | Architecture | Source | +|-----------|----------|--------------|--------| +| `skysensepp-swinv2-msl-hr` | High-res optical | SwinV2 Huge + MSL | `skysensepp_release_hr.pth` | +| `skysensepp-vit-msl-s2` | Sentinel-2 | ViT-Large + MSL | `skysensepp_release_s2.pth` | +| `skysensepp-vit-msl-s1` | Sentinel-1 | ViT-Large + MSL | `skysensepp_release_s1.pth` | +| `skysensepp-fusion-neck` | Multi-modal fusion (optional) | TransformerEncoder | `fusion.*` from `skysensepp_release.ckpt` | +| `skysensepp-fewshot-release` | Full 1-shot segmentation | HR + S2 + S1 + fusion + VAE + UPerHead | `skysensepp_release.ckpt` | + +Each subdirectory is a self-contained HuggingFace model repo with remote code (`trust_remote_code=True`). + +The fusion neck is an **optional** component — backbone checkpoints do not include or require it by default. + +The few-shot release bundles all submodules into one end-to-end model (~6.8 GB). + +## Usage + +```python +from transformers import pipeline +import torch + +MODEL = "/path/to/SkySensePlusPlus-transformers/skysensepp-swinv2-msl-hr" + +pipe = pipeline( + task="image-feature-extraction", + model=MODEL, + trust_remote_code=True, + device="cpu", +) + +hr_img = torch.randn(1, 3, 512, 512) +annotation = torch.zeros(1, 512, 512, dtype=torch.long) # semantic class indices + +features = pipe(hr_img, annotation=annotation) +print(features["last_hidden_state"].shape) # (1, 2816, 16, 16) +``` + +Sentinel-2 / Sentinel-1 backbones use the same pipeline pattern: + +```python +s2_pipe = pipeline( + task="image-feature-extraction", + model="/path/to/skysensepp-vit-msl-s2", + trust_remote_code=True, + device="cpu", +) + +s2_img = torch.randn(1, 10, 16, 16) +s2_anno = torch.zeros(1, 16, 16, dtype=torch.long) +features = s2_pipe(s2_img, annotation=s2_anno) +print(features["last_hidden_state"].shape) +``` + +SkySense++ MSL models require both imagery and a semantic annotation map. Use class index `0` for background/unlabeled regions during zero-shot feature extraction. + +### Optional fusion neck + +```python +fusion_pipe = pipeline( + task="skysensepp-fusion", + model="/path/to/skysensepp-fusion-neck", + trust_remote_code=True, + device="cpu", +) + +# Concatenated HR + S2 + S1 stage-3 tokens per spatial location +hidden_states = torch.randn(256, 3, 2816) +fused = fusion_pipe(hidden_states) + +print(fused["pooler_output"].shape) # (256, 1024) +``` + +### Few-shot / 1-shot segmentation + +The full release model expects vertically stacked prompt+query inputs (prompt on top, query on bottom): + +```python +from transformers import pipeline +import torch + +MODEL = "/path/to/SkySensePlusPlus-transformers/skysensepp-fewshot-release" + +pipe = pipeline( + task="skysensepp-fewshot", + model=MODEL, + trust_remote_code=True, + device=0, # GPU recommended (~24 GB); CPU OOMs at 1024×512 HR +) + +# Stacked HR (3, 1024, 512), S2/S1 with seq=2, RGB targets (ImageNet-normalized) +hr = torch.randn(1, 3, 1024, 512) +s2 = torch.randn(1, 10, 2, 32, 32) +s1 = torch.randn(1, 2, 2, 32, 32) +targets = torch.randn(1, 3, 1024, 512) # use real RGB annotation maps in practice +anno_mask = torch.zeros(1, 8, 4, dtype=torch.long) +anno_mask[:, 4:, :] = 1 # mask query (bottom) half + +result = pipe(hr, s2_img=s2, s1_img=s1, targets=targets, anno_mask=anno_mask) +print(result["logits"].shape) # (1, 65, 512, 512) — query region only +``` + +## Conversion + +Source project: `/home/czy/local/projects/SkySensePlusPlus-transformers` + +```bash +conda activate rsgen +python scripts/convert_checkpoint_to_hf.py \ + --input-path /path/to/skysensepp_release_hr.pth \ + --modality hr \ + --output-dir /path/to/skysensepp-swinv2-msl-hr \ + --clean-output + +# Full few-shot release (~6.8 GB) +python scripts/convert_checkpoint_to_hf.py \ + --input-path /path/to/skysensepp_release.ckpt \ + --modality fewshot \ + --output-dir /path/to/skysensepp-fewshot-release \ + --clean-output +``` + +## Notes + +- 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. +- The few-shot model uses the same 62 skipped HR buffers; all 1522 learned tensors load with 0 unexpected keys. diff --git a/skysensepp-fewshot-release/config.json b/skysensepp-fewshot-release/config.json new file mode 100644 index 0000000000000000000000000000000000000000..3ea93b475106b2ea319623a54a87bc08129f6b7a --- /dev/null +++ b/skysensepp-fewshot-release/config.json @@ -0,0 +1,357 @@ +{ + "return_dict": true, + "output_hidden_states": false, + "dtype": "float32", + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": [ + "SkySensePlusPlusModel" + ], + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "transformers_version": "5.0.0", + "sources": [ + "hr", + "s2", + "s1" + ], + "vocabulary_size": 64, + "use_modal_vae": true, + "upsample_results": true, + "backbone_hr": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "arch": "huge", + "embed_dims": 352, + "depths": [ + 2, + 2, + 18, + 2 + ], + "num_heads": [ + 8, + 16, + 32, + 64 + ], + "extra_norm_every_n_blocks": 6, + "img_size": 512, + "patch_size": 4, + "in_channels": 3, + "window_size": 8, + "drop_rate": 0.0, + "drop_path_rate": 0.2, + "out_indices": [ + 0, + 1, + 2, + 3 + ], + "use_abs_pos_embed": false, + "with_cp": false, + "pad_small_map": false, + "pretrained_window_sizes": [ + 0, + 0, + 0, + 0 + ], + "is_post_norm_downsample": true, + "vocabulary_size": 64, + "num_vocabulary_tokens": 65, + "merge_stage": 2, + "use_attn": true, + "model_type": "skysensepp_swinv2_msl", + "output_attentions": false + }, + "backbone_s2": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "img_size": 16, + "patch_size": 4, + "in_channels": 10, + "embed_dims": 1024, + "num_layers": 24, + "num_heads": 16, + "mlp_ratio": 4, + "out_indices": [ + 5, + 11, + 17, + 23 + ], + "qkv_bias": true, + "drop_rate": 0.0, + "attn_drop_rate": 0.0, + "drop_path_rate": 0.3, + "with_cls_token": false, + "output_cls_token": false, + "patch_norm": false, + "final_norm": false, + "with_cp": false, + "vocabulary_size": 64, + "num_vocabulary_tokens": 65, + "merge_stage": 4, + "use_attn": false, + "modality": "s2", + "model_type": "skysensepp_vit_msl", + "output_attentions": false + }, + "backbone_s1": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "img_size": 16, + "patch_size": 4, + "in_channels": 2, + "embed_dims": 1024, + "num_layers": 24, + "num_heads": 16, + "mlp_ratio": 4, + "out_indices": [ + 5, + 11, + 17, + 23 + ], + "qkv_bias": true, + "drop_rate": 0.0, + "attn_drop_rate": 0.0, + "drop_path_rate": 0.3, + "with_cls_token": false, + "output_cls_token": false, + "patch_norm": false, + "final_norm": false, + "with_cp": false, + "vocabulary_size": 64, + "num_vocabulary_tokens": 65, + "merge_stage": 4, + "use_attn": false, + "modality": "s1", + "model_type": "skysensepp_vit_msl", + "output_attentions": false + }, + "head_s2": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "in_dim": 1024, + "out_dim": 2816, + "up_scale": 4, + "model_type": "skysensepp_up_head", + "output_attentions": false + }, + "head_s1": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "in_dim": 1024, + "out_dim": 2816, + "up_scale": 4, + "model_type": "skysensepp_up_head", + "output_attentions": false + }, + "fusion": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "input_dims": 2816, + "embed_dims": 1024, + "num_layers": 24, + "num_heads": 16, + "mlp_ratio": 4, + "qkv_bias": true, + "drop_rate": 0.0, + "attn_drop_rate": 0.0, + "drop_path_rate": 0.3, + "with_cls_token": true, + "output_cls_token": true, + "with_cp": false, + "model_type": "skysensepp_fusion_neck", + "output_attentions": false + }, + "modality_vae": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "input_shape_hr": [ + 2816, + 32, + 16 + ], + "input_shape_s2": [ + 2816, + 32, + 16 + ], + "input_shape_s1": [ + 2816, + 32, + 16 + ], + "conv_dim": 256, + "z_dim": 256, + "n_codebook": 8192, + "model_type": "skysensepp_modality_vae", + "output_attentions": false + }, + "head_rec_hr": { + "return_dict": true, + "output_hidden_states": false, + "dtype": null, + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": null, + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "in_channels": [ + 704, + 704, + 1408, + 2816, + 1024 + ], + "channels": 512, + "num_classes": 65, + "pool_scales": [ + 1, + 2, + 3, + 6 + ], + "dropout_ratio": 0.1, + "align_corners": false, + "model_type": "skysensepp_uper_head", + "output_attentions": false + }, + "model_type": "skysensepp", + "output_attentions": false, + "auto_map": { + "AutoConfig": "configuration_skysensepp.SkySensePlusPlusConfig", + "AutoModel": "modeling_skysensepp.SkySensePlusPlusModel" + }, + "custom_pipelines": { + "skysensepp-fusion": { + "impl": "pipeline_skysensepp_fewshot.SkySensePlusPlusFewShotPipeline", + "pt": [ + "AutoModel" + ] + }, + "skysensepp-fewshot": { + "impl": "pipeline_skysensepp_fewshot.SkySensePlusPlusFewShotPipeline", + "pt": [ + "AutoModel" + ] + } + } +} diff --git a/skysensepp-fewshot-release/configuration_skysensepp.py b/skysensepp-fewshot-release/configuration_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..dcb3becfcce63f900903d5a90e3a237cd2fdccaf --- /dev/null +++ b/skysensepp-fewshot-release/configuration_skysensepp.py @@ -0,0 +1,281 @@ +"""Configuration classes for SkySense++ MSL backbones.""" + +from transformers import PretrainedConfig + + +class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig): + """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical).""" + + model_type = "skysensepp_swinv2_msl" + + arch_zoo = { + "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0}, + "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0}, + "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6}, + "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6}, + } + + def __init__( + self, + arch="huge", + img_size=512, + patch_size=4, + in_channels=3, + window_size=8, + drop_rate=0.0, + drop_path_rate=0.2, + out_indices=(0, 1, 2, 3), + use_abs_pos_embed=False, + with_cp=False, + pad_small_map=False, + pretrained_window_sizes=(0, 0, 0, 0), + is_post_norm_downsample=True, + vocabulary_size=64, + merge_stage=2, + use_attn=True, + **kwargs, + ): + super().__init__(**kwargs) + + arch = arch.lower() + if arch not in self.arch_zoo: + raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}") + arch_settings = self.arch_zoo[arch] + + self.arch = arch + self.embed_dims = arch_settings["embed_dims"] + self.depths = arch_settings["depths"] + self.num_heads = arch_settings["num_heads"] + self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"] + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.window_size = window_size + self.drop_rate = drop_rate + self.drop_path_rate = drop_path_rate + self.out_indices = list(out_indices) + self.use_abs_pos_embed = use_abs_pos_embed + self.with_cp = with_cp + self.pad_small_map = pad_small_map + self.pretrained_window_sizes = list(pretrained_window_sizes) + self.is_post_norm_downsample = is_post_norm_downsample + + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + + +class SkySensePlusPlusViTMSLConfig(PretrainedConfig): + """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1).""" + + model_type = "skysensepp_vit_msl" + + def __init__( + self, + img_size=16, + patch_size=4, + in_channels=10, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + out_indices=(5, 11, 17, 23), + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=False, + output_cls_token=False, + patch_norm=False, + final_norm=False, + with_cp=False, + vocabulary_size=64, + merge_stage=4, + use_attn=False, + modality="s2", + **kwargs, + ): + super().__init__(**kwargs) + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.out_indices = list(out_indices) + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.patch_norm = patch_norm + self.final_norm = final_norm + self.with_cp = with_cp + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + self.modality = modality + + +class SkySensePlusPlusFusionNeckConfig(PretrainedConfig): + """Configuration for SkySense++ multi-modal fusion neck (TransformerEncoder). + + Optional component — not used by default backbone checkpoints. + Fuses concatenated HR/S2/S1 stage-3 features (2816-dim) via a ViT encoder + with cls token output (1024-dim). + """ + + model_type = "skysensepp_fusion_neck" + + def __init__( + self, + input_dims=2816, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=True, + output_cls_token=True, + with_cp=False, + **kwargs, + ): + super().__init__(**kwargs) + self.input_dims = input_dims + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.with_cp = with_cp + + +class UPHeadConfig(PretrainedConfig): + model_type = "skysensepp_up_head" + + def __init__(self, in_dim=1024, out_dim=2816, up_scale=4, **kwargs): + super().__init__(**kwargs) + self.in_dim = in_dim + self.out_dim = out_dim + self.up_scale = up_scale + + +class UPerHeadConfig(PretrainedConfig): + model_type = "skysensepp_uper_head" + + def __init__( + self, + in_channels=(704, 704, 1408, 2816, 1024), + channels=512, + num_classes=65, + pool_scales=(1, 2, 3, 6), + dropout_ratio=0.1, + align_corners=False, + **kwargs, + ): + super().__init__(**kwargs) + self.in_channels = list(in_channels) + self.channels = channels + self.num_classes = num_classes + self.pool_scales = list(pool_scales) + self.dropout_ratio = dropout_ratio + self.align_corners = align_corners + + +class ModalityVAEConfig(PretrainedConfig): + model_type = "skysensepp_modality_vae" + + def __init__( + self, + input_shape_hr=(2816, 32, 16), + input_shape_s2=(2816, 32, 16), + input_shape_s1=(2816, 32, 16), + conv_dim=256, + z_dim=256, + n_codebook=8192, + **kwargs, + ): + super().__init__(**kwargs) + self.input_shape_hr = list(input_shape_hr) + self.input_shape_s2 = list(input_shape_s2) + self.input_shape_s1 = list(input_shape_s1) + self.conv_dim = conv_dim + self.z_dim = z_dim + self.n_codebook = n_codebook + + +class SkySensePlusPlusConfig(PretrainedConfig): + """Full SkySense++ config for few-shot / 1-shot release checkpoint.""" + + model_type = "skysensepp" + + def __init__( + self, + sources=("hr", "s2", "s1"), + vocabulary_size=64, + use_modal_vae=True, + upsample_results=True, + backbone_hr=None, + backbone_s2=None, + backbone_s1=None, + head_s2=None, + head_s1=None, + fusion=None, + modality_vae=None, + head_rec_hr=None, + **kwargs, + ): + super().__init__(**kwargs) + self.sources = list(sources) + self.vocabulary_size = vocabulary_size + self.use_modal_vae = use_modal_vae + self.upsample_results = upsample_results + self.backbone_hr = ( + backbone_hr + if isinstance(backbone_hr, SkySensePlusPlusSwinV2MSLConfig) + else SkySensePlusPlusSwinV2MSLConfig(**(backbone_hr or {})) + ) + self.backbone_s2 = ( + backbone_s2 + if isinstance(backbone_s2, SkySensePlusPlusViTMSLConfig) + else SkySensePlusPlusViTMSLConfig(**(backbone_s2 or {"modality": "s2"})) + ) + self.backbone_s1 = ( + backbone_s1 + if isinstance(backbone_s1, SkySensePlusPlusViTMSLConfig) + else SkySensePlusPlusViTMSLConfig( + **(backbone_s1 or {"modality": "s1", "in_channels": 2}) + ) + ) + self.head_s2 = head_s2 if isinstance(head_s2, UPHeadConfig) else UPHeadConfig(**(head_s2 or {})) + self.head_s1 = head_s1 if isinstance(head_s1, UPHeadConfig) else UPHeadConfig(**(head_s1 or {})) + self.fusion = ( + fusion + if isinstance(fusion, SkySensePlusPlusFusionNeckConfig) + else SkySensePlusPlusFusionNeckConfig(**(fusion or {})) + ) + self.modality_vae = ( + modality_vae + if isinstance(modality_vae, ModalityVAEConfig) + else ModalityVAEConfig(**(modality_vae or {})) + ) + self.head_rec_hr = ( + head_rec_hr + if isinstance(head_rec_hr, UPerHeadConfig) + else UPerHeadConfig(**(head_rec_hr or {})) + ) diff --git a/skysensepp-fewshot-release/conversion_manifest.json b/skysensepp-fewshot-release/conversion_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..de7f2dfad69700673242c21c7898a623ee9debf8 --- /dev/null +++ b/skysensepp-fewshot-release/conversion_manifest.json @@ -0,0 +1,1595 @@ +{ + "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release.ckpt", + "modality": "fewshot", + "model_class": "SkySensePlusPlusModel", + "num_tensors": 1522, + "missing_keys": [ + "backbone_hr.stages.0.blocks.0.attn.w_msa.relative_coords_table", + "backbone_hr.stages.0.blocks.0.attn.w_msa.relative_position_index", + "backbone_hr.stages.0.blocks.1.attn.w_msa.relative_coords_table", + "backbone_hr.stages.0.blocks.1.attn.w_msa.relative_position_index", + "backbone_hr.stages.1.blocks.0.attn.w_msa.relative_coords_table", + "backbone_hr.stages.1.blocks.0.attn.w_msa.relative_position_index", + "backbone_hr.stages.1.blocks.1.attn.w_msa.relative_coords_table", + "backbone_hr.stages.1.blocks.1.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.0.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.0.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.1.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.1.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.2.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.2.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.3.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.3.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.4.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.4.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.5.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.5.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.6.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.6.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.7.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.7.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.8.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.8.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.9.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.9.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.10.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.10.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.11.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.11.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.12.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.12.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.13.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.13.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.14.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.14.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.15.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.15.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.16.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.16.attn.w_msa.relative_position_index", + "backbone_hr.stages.2.blocks.17.attn.w_msa.relative_coords_table", + "backbone_hr.stages.2.blocks.17.attn.w_msa.relative_position_index", + "backbone_hr.stages.3.blocks.0.attn.w_msa.relative_coords_table", + "backbone_hr.stages.3.blocks.0.attn.w_msa.relative_position_index", + "backbone_hr.stages.3.blocks.1.attn.w_msa.relative_coords_table", + "backbone_hr.stages.3.blocks.1.attn.w_msa.relative_position_index", + "head_rec_hr.psp_modules.0.1.conv.bias", + "head_rec_hr.psp_modules.1.1.conv.bias", + "head_rec_hr.psp_modules.2.1.conv.bias", + "head_rec_hr.psp_modules.3.1.conv.bias", + "head_rec_hr.bottleneck.conv.bias", + "head_rec_hr.lateral_convs.0.conv.bias", + "head_rec_hr.lateral_convs.1.conv.bias", + "head_rec_hr.lateral_convs.2.conv.bias", + "head_rec_hr.lateral_convs.3.conv.bias", + "head_rec_hr.fpn_convs.0.conv.bias", + "head_rec_hr.fpn_convs.1.conv.bias", + "head_rec_hr.fpn_convs.2.conv.bias", + "head_rec_hr.fpn_convs.3.conv.bias", + "head_rec_hr.fpn_bottleneck.conv.bias" + ], + "unexpected_keys": [], + "tensor_names": [ + "backbone_hr.attn1.attn.in_proj_bias", + "backbone_hr.attn1.attn.in_proj_weight", + "backbone_hr.attn1.attn.out_proj.bias", + "backbone_hr.attn1.attn.out_proj.weight", + "backbone_hr.attn1.proj_in.bias", + "backbone_hr.attn1.proj_in.weight", + "backbone_hr.attn1.proj_out.bias", + "backbone_hr.attn1.proj_out.weight", + "backbone_hr.attn2.attn.in_proj_bias", + "backbone_hr.attn2.attn.in_proj_weight", + "backbone_hr.attn2.attn.out_proj.bias", + "backbone_hr.attn2.attn.out_proj.weight", + "backbone_hr.attn2.proj_in.bias", + "backbone_hr.attn2.proj_in.weight", + "backbone_hr.attn2.proj_out.bias", + "backbone_hr.attn2.proj_out.weight", + "backbone_hr.attn3.attn.in_proj_bias", + "backbone_hr.attn3.attn.in_proj_weight", + "backbone_hr.attn3.attn.out_proj.bias", + "backbone_hr.attn3.attn.out_proj.weight", + "backbone_hr.attn3.proj_in.bias", + "backbone_hr.attn3.proj_in.weight", + "backbone_hr.attn3.proj_out.bias", + "backbone_hr.attn3.proj_out.weight", + "backbone_hr.mask_token", + "backbone_hr.norm0.bias", + "backbone_hr.norm0.weight", + "backbone_hr.norm1.bias", + "backbone_hr.norm1.weight", + "backbone_hr.norm2.bias", + "backbone_hr.norm2.weight", + "backbone_hr.norm3.bias", + "backbone_hr.norm3.weight", + "backbone_hr.norm_attn.bias", + "backbone_hr.norm_attn.weight", + "backbone_hr.patch_embed.norm.bias", + "backbone_hr.patch_embed.norm.weight", + "backbone_hr.patch_embed.projection.bias", + "backbone_hr.patch_embed.projection.weight", + "backbone_hr.stages.0.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.0.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.0.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.0.blocks.0.attn.w_msa.logit_scale", + "backbone_hr.stages.0.blocks.0.attn.w_msa.proj.bias", + "backbone_hr.stages.0.blocks.0.attn.w_msa.proj.weight", + "backbone_hr.stages.0.blocks.0.attn.w_msa.q_bias", + "backbone_hr.stages.0.blocks.0.attn.w_msa.qkv.weight", + "backbone_hr.stages.0.blocks.0.attn.w_msa.v_bias", + "backbone_hr.stages.0.blocks.0.ffn.layers.0.bias", + "backbone_hr.stages.0.blocks.0.ffn.layers.0.weight", + "backbone_hr.stages.0.blocks.0.ffn.layers.3.bias", + "backbone_hr.stages.0.blocks.0.ffn.layers.3.weight", + "backbone_hr.stages.0.blocks.0.norm1.bias", + "backbone_hr.stages.0.blocks.0.norm1.weight", + "backbone_hr.stages.0.blocks.0.norm2.bias", + "backbone_hr.stages.0.blocks.0.norm2.weight", + "backbone_hr.stages.0.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.0.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.0.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.0.blocks.1.attn.w_msa.logit_scale", + "backbone_hr.stages.0.blocks.1.attn.w_msa.proj.bias", + "backbone_hr.stages.0.blocks.1.attn.w_msa.proj.weight", + "backbone_hr.stages.0.blocks.1.attn.w_msa.q_bias", + "backbone_hr.stages.0.blocks.1.attn.w_msa.qkv.weight", + "backbone_hr.stages.0.blocks.1.attn.w_msa.v_bias", + "backbone_hr.stages.0.blocks.1.ffn.layers.0.bias", + "backbone_hr.stages.0.blocks.1.ffn.layers.0.weight", + "backbone_hr.stages.0.blocks.1.ffn.layers.3.bias", + "backbone_hr.stages.0.blocks.1.ffn.layers.3.weight", + "backbone_hr.stages.0.blocks.1.norm1.bias", + "backbone_hr.stages.0.blocks.1.norm1.weight", + "backbone_hr.stages.0.blocks.1.norm2.bias", + "backbone_hr.stages.0.blocks.1.norm2.weight", + "backbone_hr.stages.1.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.1.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.1.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.1.blocks.0.attn.w_msa.logit_scale", + "backbone_hr.stages.1.blocks.0.attn.w_msa.proj.bias", + "backbone_hr.stages.1.blocks.0.attn.w_msa.proj.weight", + "backbone_hr.stages.1.blocks.0.attn.w_msa.q_bias", + "backbone_hr.stages.1.blocks.0.attn.w_msa.qkv.weight", + "backbone_hr.stages.1.blocks.0.attn.w_msa.v_bias", + "backbone_hr.stages.1.blocks.0.ffn.layers.0.bias", + "backbone_hr.stages.1.blocks.0.ffn.layers.0.weight", + "backbone_hr.stages.1.blocks.0.ffn.layers.3.bias", + "backbone_hr.stages.1.blocks.0.ffn.layers.3.weight", + "backbone_hr.stages.1.blocks.0.norm1.bias", + "backbone_hr.stages.1.blocks.0.norm1.weight", + "backbone_hr.stages.1.blocks.0.norm2.bias", + "backbone_hr.stages.1.blocks.0.norm2.weight", + "backbone_hr.stages.1.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.1.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.1.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.1.blocks.1.attn.w_msa.logit_scale", + "backbone_hr.stages.1.blocks.1.attn.w_msa.proj.bias", + "backbone_hr.stages.1.blocks.1.attn.w_msa.proj.weight", + "backbone_hr.stages.1.blocks.1.attn.w_msa.q_bias", + "backbone_hr.stages.1.blocks.1.attn.w_msa.qkv.weight", + "backbone_hr.stages.1.blocks.1.attn.w_msa.v_bias", + "backbone_hr.stages.1.blocks.1.ffn.layers.0.bias", + "backbone_hr.stages.1.blocks.1.ffn.layers.0.weight", + "backbone_hr.stages.1.blocks.1.ffn.layers.3.bias", + "backbone_hr.stages.1.blocks.1.ffn.layers.3.weight", + "backbone_hr.stages.1.blocks.1.norm1.bias", + "backbone_hr.stages.1.blocks.1.norm1.weight", + "backbone_hr.stages.1.blocks.1.norm2.bias", + "backbone_hr.stages.1.blocks.1.norm2.weight", + "backbone_hr.stages.1.downsample.norm.bias", + "backbone_hr.stages.1.downsample.norm.weight", + "backbone_hr.stages.1.downsample.reduction.weight", + "backbone_hr.stages.2.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.0.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.0.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.0.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.0.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.0.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.0.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.0.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.0.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.0.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.0.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.0.norm1.bias", + "backbone_hr.stages.2.blocks.0.norm1.weight", + "backbone_hr.stages.2.blocks.0.norm2.bias", + "backbone_hr.stages.2.blocks.0.norm2.weight", + "backbone_hr.stages.2.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.1.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.1.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.1.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.1.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.1.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.1.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.1.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.1.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.1.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.1.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.1.norm1.bias", + "backbone_hr.stages.2.blocks.1.norm1.weight", + "backbone_hr.stages.2.blocks.1.norm2.bias", + "backbone_hr.stages.2.blocks.1.norm2.weight", + "backbone_hr.stages.2.blocks.10.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.10.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.10.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.10.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.10.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.10.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.10.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.10.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.10.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.10.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.10.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.10.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.10.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.10.norm1.bias", + "backbone_hr.stages.2.blocks.10.norm1.weight", + "backbone_hr.stages.2.blocks.10.norm2.bias", + "backbone_hr.stages.2.blocks.10.norm2.weight", + "backbone_hr.stages.2.blocks.11.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.11.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.11.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.11.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.11.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.11.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.11.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.11.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.11.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.11.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.11.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.11.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.11.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.11.norm1.bias", + "backbone_hr.stages.2.blocks.11.norm1.weight", + "backbone_hr.stages.2.blocks.11.norm2.bias", + "backbone_hr.stages.2.blocks.11.norm2.weight", + "backbone_hr.stages.2.blocks.11.norm3.bias", + "backbone_hr.stages.2.blocks.11.norm3.weight", + "backbone_hr.stages.2.blocks.12.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.12.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.12.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.12.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.12.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.12.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.12.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.12.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.12.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.12.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.12.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.12.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.12.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.12.norm1.bias", + "backbone_hr.stages.2.blocks.12.norm1.weight", + "backbone_hr.stages.2.blocks.12.norm2.bias", + "backbone_hr.stages.2.blocks.12.norm2.weight", + "backbone_hr.stages.2.blocks.13.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.13.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.13.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.13.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.13.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.13.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.13.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.13.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.13.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.13.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.13.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.13.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.13.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.13.norm1.bias", + "backbone_hr.stages.2.blocks.13.norm1.weight", + "backbone_hr.stages.2.blocks.13.norm2.bias", + "backbone_hr.stages.2.blocks.13.norm2.weight", + "backbone_hr.stages.2.blocks.14.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.14.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.14.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.14.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.14.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.14.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.14.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.14.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.14.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.14.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.14.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.14.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.14.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.14.norm1.bias", + "backbone_hr.stages.2.blocks.14.norm1.weight", + "backbone_hr.stages.2.blocks.14.norm2.bias", + "backbone_hr.stages.2.blocks.14.norm2.weight", + "backbone_hr.stages.2.blocks.15.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.15.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.15.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.15.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.15.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.15.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.15.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.15.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.15.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.15.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.15.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.15.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.15.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.15.norm1.bias", + "backbone_hr.stages.2.blocks.15.norm1.weight", + "backbone_hr.stages.2.blocks.15.norm2.bias", + "backbone_hr.stages.2.blocks.15.norm2.weight", + "backbone_hr.stages.2.blocks.16.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.16.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.16.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.16.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.16.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.16.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.16.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.16.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.16.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.16.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.16.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.16.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.16.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.16.norm1.bias", + "backbone_hr.stages.2.blocks.16.norm1.weight", + "backbone_hr.stages.2.blocks.16.norm2.bias", + "backbone_hr.stages.2.blocks.16.norm2.weight", + "backbone_hr.stages.2.blocks.17.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.17.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.17.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.17.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.17.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.17.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.17.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.17.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.17.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.17.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.17.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.17.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.17.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.17.norm1.bias", + "backbone_hr.stages.2.blocks.17.norm1.weight", + "backbone_hr.stages.2.blocks.17.norm2.bias", + "backbone_hr.stages.2.blocks.17.norm2.weight", + "backbone_hr.stages.2.blocks.17.norm3.bias", + "backbone_hr.stages.2.blocks.17.norm3.weight", + "backbone_hr.stages.2.blocks.2.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.2.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.2.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.2.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.2.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.2.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.2.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.2.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.2.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.2.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.2.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.2.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.2.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.2.norm1.bias", + "backbone_hr.stages.2.blocks.2.norm1.weight", + "backbone_hr.stages.2.blocks.2.norm2.bias", + "backbone_hr.stages.2.blocks.2.norm2.weight", + "backbone_hr.stages.2.blocks.3.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.3.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.3.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.3.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.3.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.3.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.3.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.3.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.3.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.3.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.3.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.3.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.3.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.3.norm1.bias", + "backbone_hr.stages.2.blocks.3.norm1.weight", + "backbone_hr.stages.2.blocks.3.norm2.bias", + "backbone_hr.stages.2.blocks.3.norm2.weight", + "backbone_hr.stages.2.blocks.4.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.4.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.4.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.4.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.4.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.4.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.4.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.4.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.4.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.4.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.4.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.4.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.4.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.4.norm1.bias", + "backbone_hr.stages.2.blocks.4.norm1.weight", + "backbone_hr.stages.2.blocks.4.norm2.bias", + "backbone_hr.stages.2.blocks.4.norm2.weight", + "backbone_hr.stages.2.blocks.5.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.5.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.5.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.5.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.5.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.5.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.5.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.5.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.5.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.5.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.5.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.5.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.5.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.5.norm1.bias", + "backbone_hr.stages.2.blocks.5.norm1.weight", + "backbone_hr.stages.2.blocks.5.norm2.bias", + "backbone_hr.stages.2.blocks.5.norm2.weight", + "backbone_hr.stages.2.blocks.5.norm3.bias", + "backbone_hr.stages.2.blocks.5.norm3.weight", + "backbone_hr.stages.2.blocks.6.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.6.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.6.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.6.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.6.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.6.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.6.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.6.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.6.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.6.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.6.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.6.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.6.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.6.norm1.bias", + "backbone_hr.stages.2.blocks.6.norm1.weight", + "backbone_hr.stages.2.blocks.6.norm2.bias", + "backbone_hr.stages.2.blocks.6.norm2.weight", + "backbone_hr.stages.2.blocks.7.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.7.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.7.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.7.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.7.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.7.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.7.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.7.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.7.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.7.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.7.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.7.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.7.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.7.norm1.bias", + "backbone_hr.stages.2.blocks.7.norm1.weight", + "backbone_hr.stages.2.blocks.7.norm2.bias", + "backbone_hr.stages.2.blocks.7.norm2.weight", + "backbone_hr.stages.2.blocks.8.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.8.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.8.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.8.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.8.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.8.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.8.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.8.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.8.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.8.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.8.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.8.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.8.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.8.norm1.bias", + "backbone_hr.stages.2.blocks.8.norm1.weight", + "backbone_hr.stages.2.blocks.8.norm2.bias", + "backbone_hr.stages.2.blocks.8.norm2.weight", + "backbone_hr.stages.2.blocks.9.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.2.blocks.9.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.2.blocks.9.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.2.blocks.9.attn.w_msa.logit_scale", + "backbone_hr.stages.2.blocks.9.attn.w_msa.proj.bias", + "backbone_hr.stages.2.blocks.9.attn.w_msa.proj.weight", + "backbone_hr.stages.2.blocks.9.attn.w_msa.q_bias", + "backbone_hr.stages.2.blocks.9.attn.w_msa.qkv.weight", + "backbone_hr.stages.2.blocks.9.attn.w_msa.v_bias", + "backbone_hr.stages.2.blocks.9.ffn.layers.0.bias", + "backbone_hr.stages.2.blocks.9.ffn.layers.0.weight", + "backbone_hr.stages.2.blocks.9.ffn.layers.3.bias", + "backbone_hr.stages.2.blocks.9.ffn.layers.3.weight", + "backbone_hr.stages.2.blocks.9.norm1.bias", + "backbone_hr.stages.2.blocks.9.norm1.weight", + "backbone_hr.stages.2.blocks.9.norm2.bias", + "backbone_hr.stages.2.blocks.9.norm2.weight", + "backbone_hr.stages.2.downsample.norm.bias", + "backbone_hr.stages.2.downsample.norm.weight", + "backbone_hr.stages.2.downsample.reduction.weight", + "backbone_hr.stages.3.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.3.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.3.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.3.blocks.0.attn.w_msa.logit_scale", + "backbone_hr.stages.3.blocks.0.attn.w_msa.proj.bias", + "backbone_hr.stages.3.blocks.0.attn.w_msa.proj.weight", + "backbone_hr.stages.3.blocks.0.attn.w_msa.q_bias", + "backbone_hr.stages.3.blocks.0.attn.w_msa.qkv.weight", + "backbone_hr.stages.3.blocks.0.attn.w_msa.v_bias", + "backbone_hr.stages.3.blocks.0.ffn.layers.0.bias", + "backbone_hr.stages.3.blocks.0.ffn.layers.0.weight", + "backbone_hr.stages.3.blocks.0.ffn.layers.3.bias", + "backbone_hr.stages.3.blocks.0.ffn.layers.3.weight", + "backbone_hr.stages.3.blocks.0.norm1.bias", + "backbone_hr.stages.3.blocks.0.norm1.weight", + "backbone_hr.stages.3.blocks.0.norm2.bias", + "backbone_hr.stages.3.blocks.0.norm2.weight", + "backbone_hr.stages.3.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "backbone_hr.stages.3.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "backbone_hr.stages.3.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "backbone_hr.stages.3.blocks.1.attn.w_msa.logit_scale", + "backbone_hr.stages.3.blocks.1.attn.w_msa.proj.bias", + "backbone_hr.stages.3.blocks.1.attn.w_msa.proj.weight", + "backbone_hr.stages.3.blocks.1.attn.w_msa.q_bias", + "backbone_hr.stages.3.blocks.1.attn.w_msa.qkv.weight", + "backbone_hr.stages.3.blocks.1.attn.w_msa.v_bias", + "backbone_hr.stages.3.blocks.1.ffn.layers.0.bias", + "backbone_hr.stages.3.blocks.1.ffn.layers.0.weight", + "backbone_hr.stages.3.blocks.1.ffn.layers.3.bias", + "backbone_hr.stages.3.blocks.1.ffn.layers.3.weight", + "backbone_hr.stages.3.blocks.1.norm1.bias", + "backbone_hr.stages.3.blocks.1.norm1.weight", + "backbone_hr.stages.3.blocks.1.norm2.bias", + "backbone_hr.stages.3.blocks.1.norm2.weight", + "backbone_hr.stages.3.downsample.norm.bias", + "backbone_hr.stages.3.downsample.norm.weight", + "backbone_hr.stages.3.downsample.reduction.weight", + "backbone_hr.vocabulary_token", + "backbone_hr.vocabulary_weight", + "backbone_s1.cls_token", + "backbone_s1.layers.0.attn.in_proj_bias", + "backbone_s1.layers.0.attn.in_proj_weight", + "backbone_s1.layers.0.attn.out_proj.bias", + "backbone_s1.layers.0.attn.out_proj.weight", + "backbone_s1.layers.0.ffn.layers.0.bias", + "backbone_s1.layers.0.ffn.layers.0.weight", + "backbone_s1.layers.0.ffn.layers.3.bias", + "backbone_s1.layers.0.ffn.layers.3.weight", + "backbone_s1.layers.0.norm1.bias", + "backbone_s1.layers.0.norm1.weight", + "backbone_s1.layers.0.norm2.bias", + "backbone_s1.layers.0.norm2.weight", + "backbone_s1.layers.1.attn.in_proj_bias", + "backbone_s1.layers.1.attn.in_proj_weight", + "backbone_s1.layers.1.attn.out_proj.bias", + "backbone_s1.layers.1.attn.out_proj.weight", + "backbone_s1.layers.1.ffn.layers.0.bias", + "backbone_s1.layers.1.ffn.layers.0.weight", + "backbone_s1.layers.1.ffn.layers.3.bias", + "backbone_s1.layers.1.ffn.layers.3.weight", + "backbone_s1.layers.1.norm1.bias", + "backbone_s1.layers.1.norm1.weight", + "backbone_s1.layers.1.norm2.bias", + "backbone_s1.layers.1.norm2.weight", + "backbone_s1.layers.10.attn.in_proj_bias", + "backbone_s1.layers.10.attn.in_proj_weight", + "backbone_s1.layers.10.attn.out_proj.bias", + "backbone_s1.layers.10.attn.out_proj.weight", + "backbone_s1.layers.10.ffn.layers.0.bias", + "backbone_s1.layers.10.ffn.layers.0.weight", + "backbone_s1.layers.10.ffn.layers.3.bias", + "backbone_s1.layers.10.ffn.layers.3.weight", + "backbone_s1.layers.10.norm1.bias", + "backbone_s1.layers.10.norm1.weight", + "backbone_s1.layers.10.norm2.bias", + "backbone_s1.layers.10.norm2.weight", + "backbone_s1.layers.11.attn.in_proj_bias", + "backbone_s1.layers.11.attn.in_proj_weight", + "backbone_s1.layers.11.attn.out_proj.bias", + "backbone_s1.layers.11.attn.out_proj.weight", + "backbone_s1.layers.11.ffn.layers.0.bias", + "backbone_s1.layers.11.ffn.layers.0.weight", + "backbone_s1.layers.11.ffn.layers.3.bias", + "backbone_s1.layers.11.ffn.layers.3.weight", + "backbone_s1.layers.11.norm1.bias", + "backbone_s1.layers.11.norm1.weight", + "backbone_s1.layers.11.norm2.bias", + "backbone_s1.layers.11.norm2.weight", + "backbone_s1.layers.12.attn.in_proj_bias", + "backbone_s1.layers.12.attn.in_proj_weight", + "backbone_s1.layers.12.attn.out_proj.bias", + "backbone_s1.layers.12.attn.out_proj.weight", + "backbone_s1.layers.12.ffn.layers.0.bias", + "backbone_s1.layers.12.ffn.layers.0.weight", + "backbone_s1.layers.12.ffn.layers.3.bias", + "backbone_s1.layers.12.ffn.layers.3.weight", + "backbone_s1.layers.12.norm1.bias", + "backbone_s1.layers.12.norm1.weight", + "backbone_s1.layers.12.norm2.bias", + "backbone_s1.layers.12.norm2.weight", + "backbone_s1.layers.13.attn.in_proj_bias", + "backbone_s1.layers.13.attn.in_proj_weight", + "backbone_s1.layers.13.attn.out_proj.bias", + "backbone_s1.layers.13.attn.out_proj.weight", + "backbone_s1.layers.13.ffn.layers.0.bias", + "backbone_s1.layers.13.ffn.layers.0.weight", + "backbone_s1.layers.13.ffn.layers.3.bias", + "backbone_s1.layers.13.ffn.layers.3.weight", + "backbone_s1.layers.13.norm1.bias", + "backbone_s1.layers.13.norm1.weight", + "backbone_s1.layers.13.norm2.bias", + "backbone_s1.layers.13.norm2.weight", + "backbone_s1.layers.14.attn.in_proj_bias", + "backbone_s1.layers.14.attn.in_proj_weight", + "backbone_s1.layers.14.attn.out_proj.bias", + "backbone_s1.layers.14.attn.out_proj.weight", + "backbone_s1.layers.14.ffn.layers.0.bias", + "backbone_s1.layers.14.ffn.layers.0.weight", + "backbone_s1.layers.14.ffn.layers.3.bias", + "backbone_s1.layers.14.ffn.layers.3.weight", + "backbone_s1.layers.14.norm1.bias", + "backbone_s1.layers.14.norm1.weight", + "backbone_s1.layers.14.norm2.bias", + "backbone_s1.layers.14.norm2.weight", + "backbone_s1.layers.15.attn.in_proj_bias", + "backbone_s1.layers.15.attn.in_proj_weight", + "backbone_s1.layers.15.attn.out_proj.bias", + "backbone_s1.layers.15.attn.out_proj.weight", + "backbone_s1.layers.15.ffn.layers.0.bias", + "backbone_s1.layers.15.ffn.layers.0.weight", + "backbone_s1.layers.15.ffn.layers.3.bias", + "backbone_s1.layers.15.ffn.layers.3.weight", + "backbone_s1.layers.15.norm1.bias", + "backbone_s1.layers.15.norm1.weight", + "backbone_s1.layers.15.norm2.bias", + "backbone_s1.layers.15.norm2.weight", + "backbone_s1.layers.16.attn.in_proj_bias", + "backbone_s1.layers.16.attn.in_proj_weight", + "backbone_s1.layers.16.attn.out_proj.bias", + "backbone_s1.layers.16.attn.out_proj.weight", + "backbone_s1.layers.16.ffn.layers.0.bias", + "backbone_s1.layers.16.ffn.layers.0.weight", + "backbone_s1.layers.16.ffn.layers.3.bias", + "backbone_s1.layers.16.ffn.layers.3.weight", + "backbone_s1.layers.16.norm1.bias", + "backbone_s1.layers.16.norm1.weight", + "backbone_s1.layers.16.norm2.bias", + "backbone_s1.layers.16.norm2.weight", + "backbone_s1.layers.17.attn.in_proj_bias", + "backbone_s1.layers.17.attn.in_proj_weight", + "backbone_s1.layers.17.attn.out_proj.bias", + "backbone_s1.layers.17.attn.out_proj.weight", + "backbone_s1.layers.17.ffn.layers.0.bias", + "backbone_s1.layers.17.ffn.layers.0.weight", + "backbone_s1.layers.17.ffn.layers.3.bias", + "backbone_s1.layers.17.ffn.layers.3.weight", + "backbone_s1.layers.17.norm1.bias", + "backbone_s1.layers.17.norm1.weight", + "backbone_s1.layers.17.norm2.bias", + "backbone_s1.layers.17.norm2.weight", + "backbone_s1.layers.18.attn.in_proj_bias", + "backbone_s1.layers.18.attn.in_proj_weight", + "backbone_s1.layers.18.attn.out_proj.bias", + "backbone_s1.layers.18.attn.out_proj.weight", + "backbone_s1.layers.18.ffn.layers.0.bias", + "backbone_s1.layers.18.ffn.layers.0.weight", + "backbone_s1.layers.18.ffn.layers.3.bias", + "backbone_s1.layers.18.ffn.layers.3.weight", + "backbone_s1.layers.18.norm1.bias", + "backbone_s1.layers.18.norm1.weight", + "backbone_s1.layers.18.norm2.bias", + "backbone_s1.layers.18.norm2.weight", + "backbone_s1.layers.19.attn.in_proj_bias", + "backbone_s1.layers.19.attn.in_proj_weight", + "backbone_s1.layers.19.attn.out_proj.bias", + "backbone_s1.layers.19.attn.out_proj.weight", + "backbone_s1.layers.19.ffn.layers.0.bias", + "backbone_s1.layers.19.ffn.layers.0.weight", + "backbone_s1.layers.19.ffn.layers.3.bias", + "backbone_s1.layers.19.ffn.layers.3.weight", + "backbone_s1.layers.19.norm1.bias", + "backbone_s1.layers.19.norm1.weight", + "backbone_s1.layers.19.norm2.bias", + "backbone_s1.layers.19.norm2.weight", + "backbone_s1.layers.2.attn.in_proj_bias", + "backbone_s1.layers.2.attn.in_proj_weight", + "backbone_s1.layers.2.attn.out_proj.bias", + "backbone_s1.layers.2.attn.out_proj.weight", + "backbone_s1.layers.2.ffn.layers.0.bias", + "backbone_s1.layers.2.ffn.layers.0.weight", + "backbone_s1.layers.2.ffn.layers.3.bias", + "backbone_s1.layers.2.ffn.layers.3.weight", + "backbone_s1.layers.2.norm1.bias", + "backbone_s1.layers.2.norm1.weight", + "backbone_s1.layers.2.norm2.bias", + "backbone_s1.layers.2.norm2.weight", + "backbone_s1.layers.20.attn.in_proj_bias", + "backbone_s1.layers.20.attn.in_proj_weight", + "backbone_s1.layers.20.attn.out_proj.bias", + "backbone_s1.layers.20.attn.out_proj.weight", + "backbone_s1.layers.20.ffn.layers.0.bias", + "backbone_s1.layers.20.ffn.layers.0.weight", + "backbone_s1.layers.20.ffn.layers.3.bias", + "backbone_s1.layers.20.ffn.layers.3.weight", + "backbone_s1.layers.20.norm1.bias", + "backbone_s1.layers.20.norm1.weight", + "backbone_s1.layers.20.norm2.bias", + "backbone_s1.layers.20.norm2.weight", + "backbone_s1.layers.21.attn.in_proj_bias", + "backbone_s1.layers.21.attn.in_proj_weight", + "backbone_s1.layers.21.attn.out_proj.bias", + "backbone_s1.layers.21.attn.out_proj.weight", + "backbone_s1.layers.21.ffn.layers.0.bias", + "backbone_s1.layers.21.ffn.layers.0.weight", + "backbone_s1.layers.21.ffn.layers.3.bias", + "backbone_s1.layers.21.ffn.layers.3.weight", + "backbone_s1.layers.21.norm1.bias", + "backbone_s1.layers.21.norm1.weight", + "backbone_s1.layers.21.norm2.bias", + "backbone_s1.layers.21.norm2.weight", + "backbone_s1.layers.22.attn.in_proj_bias", + "backbone_s1.layers.22.attn.in_proj_weight", + "backbone_s1.layers.22.attn.out_proj.bias", + "backbone_s1.layers.22.attn.out_proj.weight", + "backbone_s1.layers.22.ffn.layers.0.bias", + "backbone_s1.layers.22.ffn.layers.0.weight", + "backbone_s1.layers.22.ffn.layers.3.bias", + "backbone_s1.layers.22.ffn.layers.3.weight", + "backbone_s1.layers.22.norm1.bias", + "backbone_s1.layers.22.norm1.weight", + "backbone_s1.layers.22.norm2.bias", + "backbone_s1.layers.22.norm2.weight", + "backbone_s1.layers.23.attn.in_proj_bias", + "backbone_s1.layers.23.attn.in_proj_weight", + "backbone_s1.layers.23.attn.out_proj.bias", + "backbone_s1.layers.23.attn.out_proj.weight", + "backbone_s1.layers.23.ffn.layers.0.bias", + "backbone_s1.layers.23.ffn.layers.0.weight", + "backbone_s1.layers.23.ffn.layers.3.bias", + "backbone_s1.layers.23.ffn.layers.3.weight", + "backbone_s1.layers.23.norm1.bias", + "backbone_s1.layers.23.norm1.weight", + "backbone_s1.layers.23.norm2.bias", + "backbone_s1.layers.23.norm2.weight", + "backbone_s1.layers.3.attn.in_proj_bias", + "backbone_s1.layers.3.attn.in_proj_weight", + "backbone_s1.layers.3.attn.out_proj.bias", + "backbone_s1.layers.3.attn.out_proj.weight", + "backbone_s1.layers.3.ffn.layers.0.bias", + "backbone_s1.layers.3.ffn.layers.0.weight", + "backbone_s1.layers.3.ffn.layers.3.bias", + "backbone_s1.layers.3.ffn.layers.3.weight", + "backbone_s1.layers.3.norm1.bias", + "backbone_s1.layers.3.norm1.weight", + "backbone_s1.layers.3.norm2.bias", + "backbone_s1.layers.3.norm2.weight", + "backbone_s1.layers.4.attn.in_proj_bias", + "backbone_s1.layers.4.attn.in_proj_weight", + "backbone_s1.layers.4.attn.out_proj.bias", + "backbone_s1.layers.4.attn.out_proj.weight", + "backbone_s1.layers.4.ffn.layers.0.bias", + "backbone_s1.layers.4.ffn.layers.0.weight", + "backbone_s1.layers.4.ffn.layers.3.bias", + "backbone_s1.layers.4.ffn.layers.3.weight", + "backbone_s1.layers.4.norm1.bias", + "backbone_s1.layers.4.norm1.weight", + "backbone_s1.layers.4.norm2.bias", + "backbone_s1.layers.4.norm2.weight", + "backbone_s1.layers.5.attn.in_proj_bias", + "backbone_s1.layers.5.attn.in_proj_weight", + "backbone_s1.layers.5.attn.out_proj.bias", + "backbone_s1.layers.5.attn.out_proj.weight", + "backbone_s1.layers.5.ffn.layers.0.bias", + "backbone_s1.layers.5.ffn.layers.0.weight", + "backbone_s1.layers.5.ffn.layers.3.bias", + "backbone_s1.layers.5.ffn.layers.3.weight", + "backbone_s1.layers.5.norm1.bias", + "backbone_s1.layers.5.norm1.weight", + "backbone_s1.layers.5.norm2.bias", + "backbone_s1.layers.5.norm2.weight", + "backbone_s1.layers.6.attn.in_proj_bias", + "backbone_s1.layers.6.attn.in_proj_weight", + "backbone_s1.layers.6.attn.out_proj.bias", + "backbone_s1.layers.6.attn.out_proj.weight", + "backbone_s1.layers.6.ffn.layers.0.bias", + "backbone_s1.layers.6.ffn.layers.0.weight", + "backbone_s1.layers.6.ffn.layers.3.bias", + "backbone_s1.layers.6.ffn.layers.3.weight", + "backbone_s1.layers.6.norm1.bias", + "backbone_s1.layers.6.norm1.weight", + "backbone_s1.layers.6.norm2.bias", + "backbone_s1.layers.6.norm2.weight", + "backbone_s1.layers.7.attn.in_proj_bias", + "backbone_s1.layers.7.attn.in_proj_weight", + "backbone_s1.layers.7.attn.out_proj.bias", + "backbone_s1.layers.7.attn.out_proj.weight", + "backbone_s1.layers.7.ffn.layers.0.bias", + "backbone_s1.layers.7.ffn.layers.0.weight", + "backbone_s1.layers.7.ffn.layers.3.bias", + "backbone_s1.layers.7.ffn.layers.3.weight", + "backbone_s1.layers.7.norm1.bias", + "backbone_s1.layers.7.norm1.weight", + "backbone_s1.layers.7.norm2.bias", + "backbone_s1.layers.7.norm2.weight", + "backbone_s1.layers.8.attn.in_proj_bias", + "backbone_s1.layers.8.attn.in_proj_weight", + "backbone_s1.layers.8.attn.out_proj.bias", + "backbone_s1.layers.8.attn.out_proj.weight", + "backbone_s1.layers.8.ffn.layers.0.bias", + "backbone_s1.layers.8.ffn.layers.0.weight", + "backbone_s1.layers.8.ffn.layers.3.bias", + "backbone_s1.layers.8.ffn.layers.3.weight", + "backbone_s1.layers.8.norm1.bias", + "backbone_s1.layers.8.norm1.weight", + "backbone_s1.layers.8.norm2.bias", + "backbone_s1.layers.8.norm2.weight", + "backbone_s1.layers.9.attn.in_proj_bias", + "backbone_s1.layers.9.attn.in_proj_weight", + "backbone_s1.layers.9.attn.out_proj.bias", + "backbone_s1.layers.9.attn.out_proj.weight", + "backbone_s1.layers.9.ffn.layers.0.bias", + "backbone_s1.layers.9.ffn.layers.0.weight", + "backbone_s1.layers.9.ffn.layers.3.bias", + "backbone_s1.layers.9.ffn.layers.3.weight", + "backbone_s1.layers.9.norm1.bias", + "backbone_s1.layers.9.norm1.weight", + "backbone_s1.layers.9.norm2.bias", + "backbone_s1.layers.9.norm2.weight", + "backbone_s1.mask_token", + "backbone_s1.patch_embed.projection.bias", + "backbone_s1.patch_embed.projection.weight", + "backbone_s1.pos_embed", + "backbone_s1.vocabulary_token", + "backbone_s1.vocabulary_weight", + "backbone_s2.cls_token", + "backbone_s2.layers.0.attn.in_proj_bias", + "backbone_s2.layers.0.attn.in_proj_weight", + "backbone_s2.layers.0.attn.out_proj.bias", + "backbone_s2.layers.0.attn.out_proj.weight", + "backbone_s2.layers.0.ffn.layers.0.bias", + "backbone_s2.layers.0.ffn.layers.0.weight", + "backbone_s2.layers.0.ffn.layers.3.bias", + "backbone_s2.layers.0.ffn.layers.3.weight", + "backbone_s2.layers.0.norm1.bias", + "backbone_s2.layers.0.norm1.weight", + "backbone_s2.layers.0.norm2.bias", + "backbone_s2.layers.0.norm2.weight", + "backbone_s2.layers.1.attn.in_proj_bias", + "backbone_s2.layers.1.attn.in_proj_weight", + "backbone_s2.layers.1.attn.out_proj.bias", + "backbone_s2.layers.1.attn.out_proj.weight", + "backbone_s2.layers.1.ffn.layers.0.bias", + "backbone_s2.layers.1.ffn.layers.0.weight", + "backbone_s2.layers.1.ffn.layers.3.bias", + "backbone_s2.layers.1.ffn.layers.3.weight", + "backbone_s2.layers.1.norm1.bias", + "backbone_s2.layers.1.norm1.weight", + "backbone_s2.layers.1.norm2.bias", + "backbone_s2.layers.1.norm2.weight", + "backbone_s2.layers.10.attn.in_proj_bias", + "backbone_s2.layers.10.attn.in_proj_weight", + "backbone_s2.layers.10.attn.out_proj.bias", + "backbone_s2.layers.10.attn.out_proj.weight", + "backbone_s2.layers.10.ffn.layers.0.bias", + "backbone_s2.layers.10.ffn.layers.0.weight", + "backbone_s2.layers.10.ffn.layers.3.bias", + "backbone_s2.layers.10.ffn.layers.3.weight", + "backbone_s2.layers.10.norm1.bias", + "backbone_s2.layers.10.norm1.weight", + "backbone_s2.layers.10.norm2.bias", + "backbone_s2.layers.10.norm2.weight", + "backbone_s2.layers.11.attn.in_proj_bias", + "backbone_s2.layers.11.attn.in_proj_weight", + "backbone_s2.layers.11.attn.out_proj.bias", + "backbone_s2.layers.11.attn.out_proj.weight", + "backbone_s2.layers.11.ffn.layers.0.bias", + "backbone_s2.layers.11.ffn.layers.0.weight", + "backbone_s2.layers.11.ffn.layers.3.bias", + "backbone_s2.layers.11.ffn.layers.3.weight", + "backbone_s2.layers.11.norm1.bias", + "backbone_s2.layers.11.norm1.weight", + "backbone_s2.layers.11.norm2.bias", + "backbone_s2.layers.11.norm2.weight", + "backbone_s2.layers.12.attn.in_proj_bias", + "backbone_s2.layers.12.attn.in_proj_weight", + "backbone_s2.layers.12.attn.out_proj.bias", + "backbone_s2.layers.12.attn.out_proj.weight", + "backbone_s2.layers.12.ffn.layers.0.bias", + "backbone_s2.layers.12.ffn.layers.0.weight", + "backbone_s2.layers.12.ffn.layers.3.bias", + "backbone_s2.layers.12.ffn.layers.3.weight", + "backbone_s2.layers.12.norm1.bias", + "backbone_s2.layers.12.norm1.weight", + "backbone_s2.layers.12.norm2.bias", + "backbone_s2.layers.12.norm2.weight", + "backbone_s2.layers.13.attn.in_proj_bias", + "backbone_s2.layers.13.attn.in_proj_weight", + "backbone_s2.layers.13.attn.out_proj.bias", + "backbone_s2.layers.13.attn.out_proj.weight", + "backbone_s2.layers.13.ffn.layers.0.bias", + "backbone_s2.layers.13.ffn.layers.0.weight", + "backbone_s2.layers.13.ffn.layers.3.bias", + "backbone_s2.layers.13.ffn.layers.3.weight", + "backbone_s2.layers.13.norm1.bias", + "backbone_s2.layers.13.norm1.weight", + "backbone_s2.layers.13.norm2.bias", + "backbone_s2.layers.13.norm2.weight", + "backbone_s2.layers.14.attn.in_proj_bias", + "backbone_s2.layers.14.attn.in_proj_weight", + "backbone_s2.layers.14.attn.out_proj.bias", + "backbone_s2.layers.14.attn.out_proj.weight", + "backbone_s2.layers.14.ffn.layers.0.bias", + "backbone_s2.layers.14.ffn.layers.0.weight", + "backbone_s2.layers.14.ffn.layers.3.bias", + "backbone_s2.layers.14.ffn.layers.3.weight", + "backbone_s2.layers.14.norm1.bias", + "backbone_s2.layers.14.norm1.weight", + "backbone_s2.layers.14.norm2.bias", + "backbone_s2.layers.14.norm2.weight", + "backbone_s2.layers.15.attn.in_proj_bias", + "backbone_s2.layers.15.attn.in_proj_weight", + "backbone_s2.layers.15.attn.out_proj.bias", + "backbone_s2.layers.15.attn.out_proj.weight", + "backbone_s2.layers.15.ffn.layers.0.bias", + "backbone_s2.layers.15.ffn.layers.0.weight", + "backbone_s2.layers.15.ffn.layers.3.bias", + "backbone_s2.layers.15.ffn.layers.3.weight", + "backbone_s2.layers.15.norm1.bias", + "backbone_s2.layers.15.norm1.weight", + "backbone_s2.layers.15.norm2.bias", + "backbone_s2.layers.15.norm2.weight", + "backbone_s2.layers.16.attn.in_proj_bias", + "backbone_s2.layers.16.attn.in_proj_weight", + "backbone_s2.layers.16.attn.out_proj.bias", + "backbone_s2.layers.16.attn.out_proj.weight", + "backbone_s2.layers.16.ffn.layers.0.bias", + "backbone_s2.layers.16.ffn.layers.0.weight", + "backbone_s2.layers.16.ffn.layers.3.bias", + "backbone_s2.layers.16.ffn.layers.3.weight", + "backbone_s2.layers.16.norm1.bias", + "backbone_s2.layers.16.norm1.weight", + "backbone_s2.layers.16.norm2.bias", + "backbone_s2.layers.16.norm2.weight", + "backbone_s2.layers.17.attn.in_proj_bias", + "backbone_s2.layers.17.attn.in_proj_weight", + "backbone_s2.layers.17.attn.out_proj.bias", + "backbone_s2.layers.17.attn.out_proj.weight", + "backbone_s2.layers.17.ffn.layers.0.bias", + "backbone_s2.layers.17.ffn.layers.0.weight", + "backbone_s2.layers.17.ffn.layers.3.bias", + "backbone_s2.layers.17.ffn.layers.3.weight", + "backbone_s2.layers.17.norm1.bias", + "backbone_s2.layers.17.norm1.weight", + "backbone_s2.layers.17.norm2.bias", + "backbone_s2.layers.17.norm2.weight", + "backbone_s2.layers.18.attn.in_proj_bias", + "backbone_s2.layers.18.attn.in_proj_weight", + "backbone_s2.layers.18.attn.out_proj.bias", + "backbone_s2.layers.18.attn.out_proj.weight", + "backbone_s2.layers.18.ffn.layers.0.bias", + "backbone_s2.layers.18.ffn.layers.0.weight", + "backbone_s2.layers.18.ffn.layers.3.bias", + "backbone_s2.layers.18.ffn.layers.3.weight", + "backbone_s2.layers.18.norm1.bias", + "backbone_s2.layers.18.norm1.weight", + "backbone_s2.layers.18.norm2.bias", + "backbone_s2.layers.18.norm2.weight", + "backbone_s2.layers.19.attn.in_proj_bias", + "backbone_s2.layers.19.attn.in_proj_weight", + "backbone_s2.layers.19.attn.out_proj.bias", + "backbone_s2.layers.19.attn.out_proj.weight", + "backbone_s2.layers.19.ffn.layers.0.bias", + "backbone_s2.layers.19.ffn.layers.0.weight", + "backbone_s2.layers.19.ffn.layers.3.bias", + "backbone_s2.layers.19.ffn.layers.3.weight", + "backbone_s2.layers.19.norm1.bias", + "backbone_s2.layers.19.norm1.weight", + "backbone_s2.layers.19.norm2.bias", + "backbone_s2.layers.19.norm2.weight", + "backbone_s2.layers.2.attn.in_proj_bias", + "backbone_s2.layers.2.attn.in_proj_weight", + "backbone_s2.layers.2.attn.out_proj.bias", + "backbone_s2.layers.2.attn.out_proj.weight", + "backbone_s2.layers.2.ffn.layers.0.bias", + "backbone_s2.layers.2.ffn.layers.0.weight", + "backbone_s2.layers.2.ffn.layers.3.bias", + "backbone_s2.layers.2.ffn.layers.3.weight", + "backbone_s2.layers.2.norm1.bias", + "backbone_s2.layers.2.norm1.weight", + "backbone_s2.layers.2.norm2.bias", + "backbone_s2.layers.2.norm2.weight", + "backbone_s2.layers.20.attn.in_proj_bias", + "backbone_s2.layers.20.attn.in_proj_weight", + "backbone_s2.layers.20.attn.out_proj.bias", + "backbone_s2.layers.20.attn.out_proj.weight", + "backbone_s2.layers.20.ffn.layers.0.bias", + "backbone_s2.layers.20.ffn.layers.0.weight", + "backbone_s2.layers.20.ffn.layers.3.bias", + "backbone_s2.layers.20.ffn.layers.3.weight", + "backbone_s2.layers.20.norm1.bias", + "backbone_s2.layers.20.norm1.weight", + "backbone_s2.layers.20.norm2.bias", + "backbone_s2.layers.20.norm2.weight", + "backbone_s2.layers.21.attn.in_proj_bias", + "backbone_s2.layers.21.attn.in_proj_weight", + "backbone_s2.layers.21.attn.out_proj.bias", + "backbone_s2.layers.21.attn.out_proj.weight", + "backbone_s2.layers.21.ffn.layers.0.bias", + "backbone_s2.layers.21.ffn.layers.0.weight", + "backbone_s2.layers.21.ffn.layers.3.bias", + "backbone_s2.layers.21.ffn.layers.3.weight", + "backbone_s2.layers.21.norm1.bias", + "backbone_s2.layers.21.norm1.weight", + "backbone_s2.layers.21.norm2.bias", + "backbone_s2.layers.21.norm2.weight", + "backbone_s2.layers.22.attn.in_proj_bias", + "backbone_s2.layers.22.attn.in_proj_weight", + "backbone_s2.layers.22.attn.out_proj.bias", + "backbone_s2.layers.22.attn.out_proj.weight", + "backbone_s2.layers.22.ffn.layers.0.bias", + "backbone_s2.layers.22.ffn.layers.0.weight", + "backbone_s2.layers.22.ffn.layers.3.bias", + "backbone_s2.layers.22.ffn.layers.3.weight", + "backbone_s2.layers.22.norm1.bias", + "backbone_s2.layers.22.norm1.weight", + "backbone_s2.layers.22.norm2.bias", + "backbone_s2.layers.22.norm2.weight", + "backbone_s2.layers.23.attn.in_proj_bias", + "backbone_s2.layers.23.attn.in_proj_weight", + "backbone_s2.layers.23.attn.out_proj.bias", + "backbone_s2.layers.23.attn.out_proj.weight", + "backbone_s2.layers.23.ffn.layers.0.bias", + "backbone_s2.layers.23.ffn.layers.0.weight", + "backbone_s2.layers.23.ffn.layers.3.bias", + "backbone_s2.layers.23.ffn.layers.3.weight", + "backbone_s2.layers.23.norm1.bias", + "backbone_s2.layers.23.norm1.weight", + "backbone_s2.layers.23.norm2.bias", + "backbone_s2.layers.23.norm2.weight", + "backbone_s2.layers.3.attn.in_proj_bias", + "backbone_s2.layers.3.attn.in_proj_weight", + "backbone_s2.layers.3.attn.out_proj.bias", + "backbone_s2.layers.3.attn.out_proj.weight", + "backbone_s2.layers.3.ffn.layers.0.bias", + "backbone_s2.layers.3.ffn.layers.0.weight", + "backbone_s2.layers.3.ffn.layers.3.bias", + "backbone_s2.layers.3.ffn.layers.3.weight", + "backbone_s2.layers.3.norm1.bias", + "backbone_s2.layers.3.norm1.weight", + "backbone_s2.layers.3.norm2.bias", + "backbone_s2.layers.3.norm2.weight", + "backbone_s2.layers.4.attn.in_proj_bias", + "backbone_s2.layers.4.attn.in_proj_weight", + "backbone_s2.layers.4.attn.out_proj.bias", + "backbone_s2.layers.4.attn.out_proj.weight", + "backbone_s2.layers.4.ffn.layers.0.bias", + "backbone_s2.layers.4.ffn.layers.0.weight", + "backbone_s2.layers.4.ffn.layers.3.bias", + "backbone_s2.layers.4.ffn.layers.3.weight", + "backbone_s2.layers.4.norm1.bias", + "backbone_s2.layers.4.norm1.weight", + "backbone_s2.layers.4.norm2.bias", + "backbone_s2.layers.4.norm2.weight", + "backbone_s2.layers.5.attn.in_proj_bias", + "backbone_s2.layers.5.attn.in_proj_weight", + "backbone_s2.layers.5.attn.out_proj.bias", + "backbone_s2.layers.5.attn.out_proj.weight", + "backbone_s2.layers.5.ffn.layers.0.bias", + "backbone_s2.layers.5.ffn.layers.0.weight", + "backbone_s2.layers.5.ffn.layers.3.bias", + "backbone_s2.layers.5.ffn.layers.3.weight", + "backbone_s2.layers.5.norm1.bias", + "backbone_s2.layers.5.norm1.weight", + "backbone_s2.layers.5.norm2.bias", + "backbone_s2.layers.5.norm2.weight", + "backbone_s2.layers.6.attn.in_proj_bias", + "backbone_s2.layers.6.attn.in_proj_weight", + "backbone_s2.layers.6.attn.out_proj.bias", + "backbone_s2.layers.6.attn.out_proj.weight", + "backbone_s2.layers.6.ffn.layers.0.bias", + "backbone_s2.layers.6.ffn.layers.0.weight", + "backbone_s2.layers.6.ffn.layers.3.bias", + "backbone_s2.layers.6.ffn.layers.3.weight", + "backbone_s2.layers.6.norm1.bias", + "backbone_s2.layers.6.norm1.weight", + "backbone_s2.layers.6.norm2.bias", + "backbone_s2.layers.6.norm2.weight", + "backbone_s2.layers.7.attn.in_proj_bias", + "backbone_s2.layers.7.attn.in_proj_weight", + "backbone_s2.layers.7.attn.out_proj.bias", + "backbone_s2.layers.7.attn.out_proj.weight", + "backbone_s2.layers.7.ffn.layers.0.bias", + "backbone_s2.layers.7.ffn.layers.0.weight", + "backbone_s2.layers.7.ffn.layers.3.bias", + "backbone_s2.layers.7.ffn.layers.3.weight", + "backbone_s2.layers.7.norm1.bias", + "backbone_s2.layers.7.norm1.weight", + "backbone_s2.layers.7.norm2.bias", + "backbone_s2.layers.7.norm2.weight", + "backbone_s2.layers.8.attn.in_proj_bias", + "backbone_s2.layers.8.attn.in_proj_weight", + "backbone_s2.layers.8.attn.out_proj.bias", + "backbone_s2.layers.8.attn.out_proj.weight", + "backbone_s2.layers.8.ffn.layers.0.bias", + "backbone_s2.layers.8.ffn.layers.0.weight", + "backbone_s2.layers.8.ffn.layers.3.bias", + "backbone_s2.layers.8.ffn.layers.3.weight", + "backbone_s2.layers.8.norm1.bias", + "backbone_s2.layers.8.norm1.weight", + "backbone_s2.layers.8.norm2.bias", + "backbone_s2.layers.8.norm2.weight", + "backbone_s2.layers.9.attn.in_proj_bias", + "backbone_s2.layers.9.attn.in_proj_weight", + "backbone_s2.layers.9.attn.out_proj.bias", + "backbone_s2.layers.9.attn.out_proj.weight", + "backbone_s2.layers.9.ffn.layers.0.bias", + "backbone_s2.layers.9.ffn.layers.0.weight", + "backbone_s2.layers.9.ffn.layers.3.bias", + "backbone_s2.layers.9.ffn.layers.3.weight", + "backbone_s2.layers.9.norm1.bias", + "backbone_s2.layers.9.norm1.weight", + "backbone_s2.layers.9.norm2.bias", + "backbone_s2.layers.9.norm2.weight", + "backbone_s2.mask_token", + "backbone_s2.patch_embed.projection.bias", + "backbone_s2.patch_embed.projection.weight", + "backbone_s2.pos_embed", + "backbone_s2.vocabulary_token", + "backbone_s2.vocabulary_weight", + "fusion.cls_token", + "fusion.layers.0.attn.in_proj_bias", + "fusion.layers.0.attn.in_proj_weight", + "fusion.layers.0.attn.out_proj.bias", + "fusion.layers.0.attn.out_proj.weight", + "fusion.layers.0.ffn.layers.0.bias", + "fusion.layers.0.ffn.layers.0.weight", + "fusion.layers.0.ffn.layers.3.bias", + "fusion.layers.0.ffn.layers.3.weight", + "fusion.layers.0.norm1.bias", + "fusion.layers.0.norm1.weight", + "fusion.layers.0.norm2.bias", + "fusion.layers.0.norm2.weight", + "fusion.layers.1.attn.in_proj_bias", + "fusion.layers.1.attn.in_proj_weight", + "fusion.layers.1.attn.out_proj.bias", + "fusion.layers.1.attn.out_proj.weight", + "fusion.layers.1.ffn.layers.0.bias", + "fusion.layers.1.ffn.layers.0.weight", + "fusion.layers.1.ffn.layers.3.bias", + "fusion.layers.1.ffn.layers.3.weight", + "fusion.layers.1.norm1.bias", + "fusion.layers.1.norm1.weight", + "fusion.layers.1.norm2.bias", + "fusion.layers.1.norm2.weight", + "fusion.layers.10.attn.in_proj_bias", + "fusion.layers.10.attn.in_proj_weight", + "fusion.layers.10.attn.out_proj.bias", + "fusion.layers.10.attn.out_proj.weight", + "fusion.layers.10.ffn.layers.0.bias", + "fusion.layers.10.ffn.layers.0.weight", + "fusion.layers.10.ffn.layers.3.bias", + "fusion.layers.10.ffn.layers.3.weight", + "fusion.layers.10.norm1.bias", + "fusion.layers.10.norm1.weight", + "fusion.layers.10.norm2.bias", + "fusion.layers.10.norm2.weight", + "fusion.layers.11.attn.in_proj_bias", + "fusion.layers.11.attn.in_proj_weight", + "fusion.layers.11.attn.out_proj.bias", + "fusion.layers.11.attn.out_proj.weight", + "fusion.layers.11.ffn.layers.0.bias", + "fusion.layers.11.ffn.layers.0.weight", + "fusion.layers.11.ffn.layers.3.bias", + "fusion.layers.11.ffn.layers.3.weight", + "fusion.layers.11.norm1.bias", + "fusion.layers.11.norm1.weight", + "fusion.layers.11.norm2.bias", + "fusion.layers.11.norm2.weight", + "fusion.layers.12.attn.in_proj_bias", + "fusion.layers.12.attn.in_proj_weight", + "fusion.layers.12.attn.out_proj.bias", + "fusion.layers.12.attn.out_proj.weight", + "fusion.layers.12.ffn.layers.0.bias", + "fusion.layers.12.ffn.layers.0.weight", + "fusion.layers.12.ffn.layers.3.bias", + "fusion.layers.12.ffn.layers.3.weight", + "fusion.layers.12.norm1.bias", + "fusion.layers.12.norm1.weight", + "fusion.layers.12.norm2.bias", + "fusion.layers.12.norm2.weight", + "fusion.layers.13.attn.in_proj_bias", + "fusion.layers.13.attn.in_proj_weight", + "fusion.layers.13.attn.out_proj.bias", + "fusion.layers.13.attn.out_proj.weight", + "fusion.layers.13.ffn.layers.0.bias", + "fusion.layers.13.ffn.layers.0.weight", + "fusion.layers.13.ffn.layers.3.bias", + "fusion.layers.13.ffn.layers.3.weight", + "fusion.layers.13.norm1.bias", + "fusion.layers.13.norm1.weight", + "fusion.layers.13.norm2.bias", + "fusion.layers.13.norm2.weight", + "fusion.layers.14.attn.in_proj_bias", + "fusion.layers.14.attn.in_proj_weight", + "fusion.layers.14.attn.out_proj.bias", + "fusion.layers.14.attn.out_proj.weight", + "fusion.layers.14.ffn.layers.0.bias", + "fusion.layers.14.ffn.layers.0.weight", + "fusion.layers.14.ffn.layers.3.bias", + "fusion.layers.14.ffn.layers.3.weight", + "fusion.layers.14.norm1.bias", + "fusion.layers.14.norm1.weight", + "fusion.layers.14.norm2.bias", + "fusion.layers.14.norm2.weight", + "fusion.layers.15.attn.in_proj_bias", + "fusion.layers.15.attn.in_proj_weight", + "fusion.layers.15.attn.out_proj.bias", + "fusion.layers.15.attn.out_proj.weight", + "fusion.layers.15.ffn.layers.0.bias", + "fusion.layers.15.ffn.layers.0.weight", + "fusion.layers.15.ffn.layers.3.bias", + "fusion.layers.15.ffn.layers.3.weight", + "fusion.layers.15.norm1.bias", + "fusion.layers.15.norm1.weight", + "fusion.layers.15.norm2.bias", + "fusion.layers.15.norm2.weight", + "fusion.layers.16.attn.in_proj_bias", + "fusion.layers.16.attn.in_proj_weight", + "fusion.layers.16.attn.out_proj.bias", + "fusion.layers.16.attn.out_proj.weight", + "fusion.layers.16.ffn.layers.0.bias", + "fusion.layers.16.ffn.layers.0.weight", + "fusion.layers.16.ffn.layers.3.bias", + "fusion.layers.16.ffn.layers.3.weight", + "fusion.layers.16.norm1.bias", + "fusion.layers.16.norm1.weight", + "fusion.layers.16.norm2.bias", + "fusion.layers.16.norm2.weight", + "fusion.layers.17.attn.in_proj_bias", + "fusion.layers.17.attn.in_proj_weight", + "fusion.layers.17.attn.out_proj.bias", + "fusion.layers.17.attn.out_proj.weight", + "fusion.layers.17.ffn.layers.0.bias", + "fusion.layers.17.ffn.layers.0.weight", + "fusion.layers.17.ffn.layers.3.bias", + "fusion.layers.17.ffn.layers.3.weight", + "fusion.layers.17.norm1.bias", + "fusion.layers.17.norm1.weight", + "fusion.layers.17.norm2.bias", + "fusion.layers.17.norm2.weight", + "fusion.layers.18.attn.in_proj_bias", + "fusion.layers.18.attn.in_proj_weight", + "fusion.layers.18.attn.out_proj.bias", + "fusion.layers.18.attn.out_proj.weight", + "fusion.layers.18.ffn.layers.0.bias", + "fusion.layers.18.ffn.layers.0.weight", + "fusion.layers.18.ffn.layers.3.bias", + "fusion.layers.18.ffn.layers.3.weight", + "fusion.layers.18.norm1.bias", + "fusion.layers.18.norm1.weight", + "fusion.layers.18.norm2.bias", + "fusion.layers.18.norm2.weight", + "fusion.layers.19.attn.in_proj_bias", + "fusion.layers.19.attn.in_proj_weight", + "fusion.layers.19.attn.out_proj.bias", + "fusion.layers.19.attn.out_proj.weight", + "fusion.layers.19.ffn.layers.0.bias", + "fusion.layers.19.ffn.layers.0.weight", + "fusion.layers.19.ffn.layers.3.bias", + "fusion.layers.19.ffn.layers.3.weight", + "fusion.layers.19.norm1.bias", + "fusion.layers.19.norm1.weight", + "fusion.layers.19.norm2.bias", + "fusion.layers.19.norm2.weight", + "fusion.layers.2.attn.in_proj_bias", + "fusion.layers.2.attn.in_proj_weight", + "fusion.layers.2.attn.out_proj.bias", + "fusion.layers.2.attn.out_proj.weight", + "fusion.layers.2.ffn.layers.0.bias", + "fusion.layers.2.ffn.layers.0.weight", + "fusion.layers.2.ffn.layers.3.bias", + "fusion.layers.2.ffn.layers.3.weight", + "fusion.layers.2.norm1.bias", + "fusion.layers.2.norm1.weight", + "fusion.layers.2.norm2.bias", + "fusion.layers.2.norm2.weight", + "fusion.layers.20.attn.in_proj_bias", + "fusion.layers.20.attn.in_proj_weight", + "fusion.layers.20.attn.out_proj.bias", + "fusion.layers.20.attn.out_proj.weight", + "fusion.layers.20.ffn.layers.0.bias", + "fusion.layers.20.ffn.layers.0.weight", + "fusion.layers.20.ffn.layers.3.bias", + "fusion.layers.20.ffn.layers.3.weight", + "fusion.layers.20.norm1.bias", + "fusion.layers.20.norm1.weight", + "fusion.layers.20.norm2.bias", + "fusion.layers.20.norm2.weight", + "fusion.layers.21.attn.in_proj_bias", + "fusion.layers.21.attn.in_proj_weight", + "fusion.layers.21.attn.out_proj.bias", + "fusion.layers.21.attn.out_proj.weight", + "fusion.layers.21.ffn.layers.0.bias", + "fusion.layers.21.ffn.layers.0.weight", + "fusion.layers.21.ffn.layers.3.bias", + "fusion.layers.21.ffn.layers.3.weight", + "fusion.layers.21.norm1.bias", + "fusion.layers.21.norm1.weight", + "fusion.layers.21.norm2.bias", + "fusion.layers.21.norm2.weight", + "fusion.layers.22.attn.in_proj_bias", + "fusion.layers.22.attn.in_proj_weight", + "fusion.layers.22.attn.out_proj.bias", + "fusion.layers.22.attn.out_proj.weight", + "fusion.layers.22.ffn.layers.0.bias", + "fusion.layers.22.ffn.layers.0.weight", + "fusion.layers.22.ffn.layers.3.bias", + "fusion.layers.22.ffn.layers.3.weight", + "fusion.layers.22.norm1.bias", + "fusion.layers.22.norm1.weight", + "fusion.layers.22.norm2.bias", + "fusion.layers.22.norm2.weight", + "fusion.layers.23.attn.in_proj_bias", + "fusion.layers.23.attn.in_proj_weight", + "fusion.layers.23.attn.out_proj.bias", + "fusion.layers.23.attn.out_proj.weight", + "fusion.layers.23.ffn.layers.0.bias", + "fusion.layers.23.ffn.layers.0.weight", + "fusion.layers.23.ffn.layers.3.bias", + "fusion.layers.23.ffn.layers.3.weight", + "fusion.layers.23.norm1.bias", + "fusion.layers.23.norm1.weight", + "fusion.layers.23.norm2.bias", + "fusion.layers.23.norm2.weight", + "fusion.layers.3.attn.in_proj_bias", + "fusion.layers.3.attn.in_proj_weight", + "fusion.layers.3.attn.out_proj.bias", + "fusion.layers.3.attn.out_proj.weight", + "fusion.layers.3.ffn.layers.0.bias", + "fusion.layers.3.ffn.layers.0.weight", + "fusion.layers.3.ffn.layers.3.bias", + "fusion.layers.3.ffn.layers.3.weight", + "fusion.layers.3.norm1.bias", + "fusion.layers.3.norm1.weight", + "fusion.layers.3.norm2.bias", + "fusion.layers.3.norm2.weight", + "fusion.layers.4.attn.in_proj_bias", + "fusion.layers.4.attn.in_proj_weight", + "fusion.layers.4.attn.out_proj.bias", + "fusion.layers.4.attn.out_proj.weight", + "fusion.layers.4.ffn.layers.0.bias", + "fusion.layers.4.ffn.layers.0.weight", + "fusion.layers.4.ffn.layers.3.bias", + "fusion.layers.4.ffn.layers.3.weight", + "fusion.layers.4.norm1.bias", + "fusion.layers.4.norm1.weight", + "fusion.layers.4.norm2.bias", + "fusion.layers.4.norm2.weight", + "fusion.layers.5.attn.in_proj_bias", + "fusion.layers.5.attn.in_proj_weight", + "fusion.layers.5.attn.out_proj.bias", + "fusion.layers.5.attn.out_proj.weight", + "fusion.layers.5.ffn.layers.0.bias", + "fusion.layers.5.ffn.layers.0.weight", + "fusion.layers.5.ffn.layers.3.bias", + "fusion.layers.5.ffn.layers.3.weight", + "fusion.layers.5.norm1.bias", + "fusion.layers.5.norm1.weight", + "fusion.layers.5.norm2.bias", + "fusion.layers.5.norm2.weight", + "fusion.layers.6.attn.in_proj_bias", + "fusion.layers.6.attn.in_proj_weight", + "fusion.layers.6.attn.out_proj.bias", + "fusion.layers.6.attn.out_proj.weight", + "fusion.layers.6.ffn.layers.0.bias", + "fusion.layers.6.ffn.layers.0.weight", + "fusion.layers.6.ffn.layers.3.bias", + "fusion.layers.6.ffn.layers.3.weight", + "fusion.layers.6.norm1.bias", + "fusion.layers.6.norm1.weight", + "fusion.layers.6.norm2.bias", + "fusion.layers.6.norm2.weight", + "fusion.layers.7.attn.in_proj_bias", + "fusion.layers.7.attn.in_proj_weight", + "fusion.layers.7.attn.out_proj.bias", + "fusion.layers.7.attn.out_proj.weight", + "fusion.layers.7.ffn.layers.0.bias", + "fusion.layers.7.ffn.layers.0.weight", + "fusion.layers.7.ffn.layers.3.bias", + "fusion.layers.7.ffn.layers.3.weight", + "fusion.layers.7.norm1.bias", + "fusion.layers.7.norm1.weight", + "fusion.layers.7.norm2.bias", + "fusion.layers.7.norm2.weight", + "fusion.layers.8.attn.in_proj_bias", + "fusion.layers.8.attn.in_proj_weight", + "fusion.layers.8.attn.out_proj.bias", + "fusion.layers.8.attn.out_proj.weight", + "fusion.layers.8.ffn.layers.0.bias", + "fusion.layers.8.ffn.layers.0.weight", + "fusion.layers.8.ffn.layers.3.bias", + "fusion.layers.8.ffn.layers.3.weight", + "fusion.layers.8.norm1.bias", + "fusion.layers.8.norm1.weight", + "fusion.layers.8.norm2.bias", + "fusion.layers.8.norm2.weight", + "fusion.layers.9.attn.in_proj_bias", + "fusion.layers.9.attn.in_proj_weight", + "fusion.layers.9.attn.out_proj.bias", + "fusion.layers.9.attn.out_proj.weight", + "fusion.layers.9.ffn.layers.0.bias", + "fusion.layers.9.ffn.layers.0.weight", + "fusion.layers.9.ffn.layers.3.bias", + "fusion.layers.9.ffn.layers.3.weight", + "fusion.layers.9.norm1.bias", + "fusion.layers.9.norm1.weight", + "fusion.layers.9.norm2.bias", + "fusion.layers.9.norm2.weight", + "fusion.porj_linear.bias", + "fusion.porj_linear.weight", + "head_rec_hr.bottleneck.bn.bias", + "head_rec_hr.bottleneck.bn.num_batches_tracked", + "head_rec_hr.bottleneck.bn.running_mean", + "head_rec_hr.bottleneck.bn.running_var", + "head_rec_hr.bottleneck.bn.weight", + "head_rec_hr.bottleneck.conv.weight", + "head_rec_hr.conv_seg.bias", + "head_rec_hr.conv_seg.weight", + "head_rec_hr.fpn_bottleneck.bn.bias", + "head_rec_hr.fpn_bottleneck.bn.num_batches_tracked", + "head_rec_hr.fpn_bottleneck.bn.running_mean", + "head_rec_hr.fpn_bottleneck.bn.running_var", + "head_rec_hr.fpn_bottleneck.bn.weight", + "head_rec_hr.fpn_bottleneck.conv.weight", + "head_rec_hr.fpn_convs.0.bn.bias", + "head_rec_hr.fpn_convs.0.bn.num_batches_tracked", + "head_rec_hr.fpn_convs.0.bn.running_mean", + "head_rec_hr.fpn_convs.0.bn.running_var", + "head_rec_hr.fpn_convs.0.bn.weight", + "head_rec_hr.fpn_convs.0.conv.weight", + "head_rec_hr.fpn_convs.1.bn.bias", + "head_rec_hr.fpn_convs.1.bn.num_batches_tracked", + "head_rec_hr.fpn_convs.1.bn.running_mean", + "head_rec_hr.fpn_convs.1.bn.running_var", + "head_rec_hr.fpn_convs.1.bn.weight", + "head_rec_hr.fpn_convs.1.conv.weight", + "head_rec_hr.fpn_convs.2.bn.bias", + "head_rec_hr.fpn_convs.2.bn.num_batches_tracked", + "head_rec_hr.fpn_convs.2.bn.running_mean", + "head_rec_hr.fpn_convs.2.bn.running_var", + "head_rec_hr.fpn_convs.2.bn.weight", + "head_rec_hr.fpn_convs.2.conv.weight", + "head_rec_hr.fpn_convs.3.bn.bias", + "head_rec_hr.fpn_convs.3.bn.num_batches_tracked", + "head_rec_hr.fpn_convs.3.bn.running_mean", + "head_rec_hr.fpn_convs.3.bn.running_var", + "head_rec_hr.fpn_convs.3.bn.weight", + "head_rec_hr.fpn_convs.3.conv.weight", + "head_rec_hr.lateral_convs.0.bn.bias", + "head_rec_hr.lateral_convs.0.bn.num_batches_tracked", + "head_rec_hr.lateral_convs.0.bn.running_mean", + "head_rec_hr.lateral_convs.0.bn.running_var", + "head_rec_hr.lateral_convs.0.bn.weight", + "head_rec_hr.lateral_convs.0.conv.weight", + "head_rec_hr.lateral_convs.1.bn.bias", + "head_rec_hr.lateral_convs.1.bn.num_batches_tracked", + "head_rec_hr.lateral_convs.1.bn.running_mean", + "head_rec_hr.lateral_convs.1.bn.running_var", + "head_rec_hr.lateral_convs.1.bn.weight", + "head_rec_hr.lateral_convs.1.conv.weight", + "head_rec_hr.lateral_convs.2.bn.bias", + "head_rec_hr.lateral_convs.2.bn.num_batches_tracked", + "head_rec_hr.lateral_convs.2.bn.running_mean", + "head_rec_hr.lateral_convs.2.bn.running_var", + "head_rec_hr.lateral_convs.2.bn.weight", + "head_rec_hr.lateral_convs.2.conv.weight", + "head_rec_hr.lateral_convs.3.bn.bias", + "head_rec_hr.lateral_convs.3.bn.num_batches_tracked", + "head_rec_hr.lateral_convs.3.bn.running_mean", + "head_rec_hr.lateral_convs.3.bn.running_var", + "head_rec_hr.lateral_convs.3.bn.weight", + "head_rec_hr.lateral_convs.3.conv.weight", + "head_rec_hr.psp_modules.0.1.bn.bias", + "head_rec_hr.psp_modules.0.1.bn.num_batches_tracked", + "head_rec_hr.psp_modules.0.1.bn.running_mean", + "head_rec_hr.psp_modules.0.1.bn.running_var", + "head_rec_hr.psp_modules.0.1.bn.weight", + "head_rec_hr.psp_modules.0.1.conv.weight", + "head_rec_hr.psp_modules.1.1.bn.bias", + "head_rec_hr.psp_modules.1.1.bn.num_batches_tracked", + "head_rec_hr.psp_modules.1.1.bn.running_mean", + "head_rec_hr.psp_modules.1.1.bn.running_var", + "head_rec_hr.psp_modules.1.1.bn.weight", + "head_rec_hr.psp_modules.1.1.conv.weight", + "head_rec_hr.psp_modules.2.1.bn.bias", + "head_rec_hr.psp_modules.2.1.bn.num_batches_tracked", + "head_rec_hr.psp_modules.2.1.bn.running_mean", + "head_rec_hr.psp_modules.2.1.bn.running_var", + "head_rec_hr.psp_modules.2.1.bn.weight", + "head_rec_hr.psp_modules.2.1.conv.weight", + "head_rec_hr.psp_modules.3.1.bn.bias", + "head_rec_hr.psp_modules.3.1.bn.num_batches_tracked", + "head_rec_hr.psp_modules.3.1.bn.running_mean", + "head_rec_hr.psp_modules.3.1.bn.running_var", + "head_rec_hr.psp_modules.3.1.bn.weight", + "head_rec_hr.psp_modules.3.1.conv.weight", + "head_s1.decoder.0.bias", + "head_s1.decoder.0.weight", + "head_s2.decoder.0.bias", + "head_s2.decoder.0.weight", + "modality_vae.vae_hr.codebook.weight", + "modality_vae.vae_hr.conv1.bias", + "modality_vae.vae_hr.conv1.weight", + "modality_vae.vae_hr.conv2.bias", + "modality_vae.vae_hr.conv2.weight", + "modality_vae.vae_hr.dec_block1.0.bias", + "modality_vae.vae_hr.dec_block1.0.weight", + "modality_vae.vae_hr.dec_block1.2.bias", + "modality_vae.vae_hr.dec_block1.2.weight", + "modality_vae.vae_hr.dec_block2.0.bias", + "modality_vae.vae_hr.dec_block2.0.weight", + "modality_vae.vae_hr.dec_block2.2.bias", + "modality_vae.vae_hr.dec_block2.2.weight", + "modality_vae.vae_hr.enc_block1.0.bias", + "modality_vae.vae_hr.enc_block1.0.weight", + "modality_vae.vae_hr.enc_block1.2.bias", + "modality_vae.vae_hr.enc_block1.2.weight", + "modality_vae.vae_hr.enc_block2.0.bias", + "modality_vae.vae_hr.enc_block2.0.weight", + "modality_vae.vae_hr.enc_block2.2.bias", + "modality_vae.vae_hr.enc_block2.2.weight", + "modality_vae.vae_hr.gamma_1", + "modality_vae.vae_hr.gamma_2", + "modality_vae.vae_hr.gamma_3", + "modality_vae.vae_hr.gamma_4", + "modality_vae.vae_hr.logit_conv.bias", + "modality_vae.vae_hr.logit_conv.weight", + "modality_vae.vae_hr.rec_conv.bias", + "modality_vae.vae_hr.rec_conv.weight", + "modality_vae.vae_s1.codebook.weight", + "modality_vae.vae_s1.conv1.bias", + "modality_vae.vae_s1.conv1.weight", + "modality_vae.vae_s1.conv2.bias", + "modality_vae.vae_s1.conv2.weight", + "modality_vae.vae_s1.dec_block1.0.bias", + "modality_vae.vae_s1.dec_block1.0.weight", + "modality_vae.vae_s1.dec_block1.2.bias", + "modality_vae.vae_s1.dec_block1.2.weight", + "modality_vae.vae_s1.dec_block2.0.bias", + "modality_vae.vae_s1.dec_block2.0.weight", + "modality_vae.vae_s1.dec_block2.2.bias", + "modality_vae.vae_s1.dec_block2.2.weight", + "modality_vae.vae_s1.enc_block1.0.bias", + "modality_vae.vae_s1.enc_block1.0.weight", + "modality_vae.vae_s1.enc_block1.2.bias", + "modality_vae.vae_s1.enc_block1.2.weight", + "modality_vae.vae_s1.enc_block2.0.bias", + "modality_vae.vae_s1.enc_block2.0.weight", + "modality_vae.vae_s1.enc_block2.2.bias", + "modality_vae.vae_s1.enc_block2.2.weight", + "modality_vae.vae_s1.gamma_1", + "modality_vae.vae_s1.gamma_2", + "modality_vae.vae_s1.gamma_3", + "modality_vae.vae_s1.gamma_4", + "modality_vae.vae_s1.logit_conv.bias", + "modality_vae.vae_s1.logit_conv.weight", + "modality_vae.vae_s1.rec_conv.bias", + "modality_vae.vae_s1.rec_conv.weight", + "modality_vae.vae_s2.codebook.weight", + "modality_vae.vae_s2.conv1.bias", + "modality_vae.vae_s2.conv1.weight", + "modality_vae.vae_s2.conv2.bias", + "modality_vae.vae_s2.conv2.weight", + "modality_vae.vae_s2.dec_block1.0.bias", + "modality_vae.vae_s2.dec_block1.0.weight", + "modality_vae.vae_s2.dec_block1.2.bias", + "modality_vae.vae_s2.dec_block1.2.weight", + "modality_vae.vae_s2.dec_block2.0.bias", + "modality_vae.vae_s2.dec_block2.0.weight", + "modality_vae.vae_s2.dec_block2.2.bias", + "modality_vae.vae_s2.dec_block2.2.weight", + "modality_vae.vae_s2.enc_block1.0.bias", + "modality_vae.vae_s2.enc_block1.0.weight", + "modality_vae.vae_s2.enc_block1.2.bias", + "modality_vae.vae_s2.enc_block1.2.weight", + "modality_vae.vae_s2.enc_block2.0.bias", + "modality_vae.vae_s2.enc_block2.0.weight", + "modality_vae.vae_s2.enc_block2.2.bias", + "modality_vae.vae_s2.enc_block2.2.weight", + "modality_vae.vae_s2.gamma_1", + "modality_vae.vae_s2.gamma_2", + "modality_vae.vae_s2.gamma_3", + "modality_vae.vae_s2.gamma_4", + "modality_vae.vae_s2.logit_conv.bias", + "modality_vae.vae_s2.logit_conv.weight", + "modality_vae.vae_s2.rec_conv.bias", + "modality_vae.vae_s2.rec_conv.weight" + ] +} diff --git a/skysensepp-fewshot-release/model.safetensors b/skysensepp-fewshot-release/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..6b601ef0397d565a2451bd025cd8d1d02f15f70f --- /dev/null +++ b/skysensepp-fewshot-release/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e6d35f7768cce0bb019d3f7089265ff6099168d651ac3ae5b4abc0139dfb759c +size 7238917044 diff --git a/skysensepp-fewshot-release/modeling_skysensepp.py b/skysensepp-fewshot-release/modeling_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..b03c141ae6d702ad9d8cfc4cf19a966f928b1730 --- /dev/null +++ b/skysensepp-fewshot-release/modeling_skysensepp.py @@ -0,0 +1,214 @@ +"""Full SkySense++ model for few-shot / 1-shot segmentation.""" + +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PreTrainedModel +from transformers.modeling_outputs import ModelOutput + +from .configuration_skysensepp import SkySensePlusPlusConfig +from .modeling_utils import DropPath as _DropPath # noqa: F401 — bundled for remote code +from .modeling_skysensepp_components import ModalityCompletion, UPerHead, UPHead +from .modeling_skysensepp_fusion_neck import SkySensePlusPlusFusionNeckModel +from .modeling_skysensepp_swinv2_msl import SkySensePlusPlusSwinV2MSLModel +from .modeling_skysensepp_vit_msl import SkySensePlusPlusViTMSLModel + +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) + + +@dataclass +class SkySensePlusPlusOutput(ModelOutput): + logits: Optional[torch.FloatTensor] = None + mapped_targets: Optional[torch.LongTensor] = None + idx_2_color: Optional[dict] = None + mask_hr: Optional[torch.Tensor] = None + vae_out: Optional[dict] = None + + +class SkySensePlusPlusPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusConfig + base_model_prefix = "skysensepp" + supports_gradient_checkpointing = False + + +class SkySensePlusPlusModel(SkySensePlusPlusPreTrainedModel): + """End-to-end SkySense++ pipeline matching the released few-shot checkpoint.""" + + def __init__(self, config: SkySensePlusPlusConfig): + super().__init__(config) + self.sources = list(config.sources) + self.use_modal_vae = config.use_modal_vae + self.vocabulary_size = config.vocabulary_size + self.vocabulary = list(range(1, config.vocabulary_size + 1)) + + if "hr" in self.sources: + self.backbone_hr = SkySensePlusPlusSwinV2MSLModel(config.backbone_hr) + if "s2" in self.sources: + self.backbone_s2 = SkySensePlusPlusViTMSLModel(config.backbone_s2) + self.head_s2 = UPHead(config.head_s2.in_dim, config.head_s2.out_dim, config.head_s2.up_scale) + if "s1" in self.sources: + self.backbone_s1 = SkySensePlusPlusViTMSLModel(config.backbone_s1) + self.head_s1 = UPHead(config.head_s1.in_dim, config.head_s1.out_dim, config.head_s1.up_scale) + + self.fusion = SkySensePlusPlusFusionNeckModel(config.fusion) + if self.use_modal_vae: + self.modality_vae = ModalityCompletion( + input_shape_hr=tuple(config.modality_vae.input_shape_hr), + input_shape_s2=tuple(config.modality_vae.input_shape_s2), + input_shape_s1=tuple(config.modality_vae.input_shape_s1), + conv_dim=config.modality_vae.conv_dim, + z_dim=config.modality_vae.z_dim, + n_codebook=config.modality_vae.n_codebook, + ) + self.head_rec_hr = UPerHead( + in_channels=config.head_rec_hr.in_channels, + channels=config.head_rec_hr.channels, + num_classes=config.head_rec_hr.num_classes, + pool_scales=tuple(config.head_rec_hr.pool_scales), + dropout_ratio=config.head_rec_hr.dropout_ratio, + align_corners=config.head_rec_hr.align_corners, + ) + self.post_init() + + def convert_target(self, target: torch.Tensor): + mean = target.new_tensor(IMAGENET_MEAN).reshape(1, 3, 1, 1) + std = target.new_tensor(IMAGENET_STD).reshape(1, 3, 1, 1) + target = ((target * std + mean) * 255).to(torch.long) + target = target[:, 0] * 256 * 256 + target[:, 1] * 256 + target[:, 2] + target = target.type(torch.long) + unique_target = target.unique() + target_index = torch.searchsorted(unique_target, target) + no_bg = unique_target[0].item() > 0 + if no_bg: + target_index = target_index + 1 + target_index_unique = target_index.unique().tolist() + vocab = target.new_tensor([0] + self.vocabulary) + mapped_target = target_index.clone() + idx_2_color = {} + for value in target_index_unique: + mapped_target[target_index == value] = vocab[value] + idx_2_color[vocab[value].item()] = unique_target[value - 1 if no_bg else value].item() + return mapped_target, idx_2_color + + def forward( + self, + hr_img: Optional[torch.Tensor] = None, + s2_img: Optional[torch.Tensor] = None, + s1_img: Optional[torch.Tensor] = None, + targets: Optional[torch.Tensor] = None, + anno_mask: Optional[torch.Tensor] = None, + modality_flags: Optional[torch.Tensor] = None, + return_dict: Optional[bool] = None, + ) -> Union[Dict, SkySensePlusPlusOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + output: Dict = {} + + if targets is None: + raise ValueError("SkySense++ few-shot forward requires `targets` for annotation conditioning.") + anno_img, idx_2_color = self.convert_target(targets) + output["mapped_targets"] = anno_img + output["idx_2_color"] = idx_2_color + + anno_s2 = anno_img[:, 15::32, 15::32] + anno_s1 = anno_s2 + + if anno_mask is not None: + batch_size, mask_h, mask_w = anno_mask.shape + block_size = 32 + anno_mask_hr = ( + anno_mask.unsqueeze(-1) + .unsqueeze(-1) + .repeat(1, 1, 1, block_size, block_size) + .permute(0, 1, 3, 2, 4) + .reshape(batch_size, mask_h * block_size, mask_w * block_size) + .contiguous() + ) + else: + anno_mask_hr = None + + if "hr" in self.sources: + hr_features = self.backbone_hr(hr_img, anno_img, anno_mask_hr, return_dict=False) + output["mask_hr"] = anno_mask_hr + + batch_size = hr_img.shape[0] + seq_len_s2 = s2_img.shape[2] if s2_img is not None else 1 + seq_len_s1 = s1_img.shape[2] if s1_img is not None else 1 + + if "s2" in self.sources: + b, c, seq, h, w = s2_img.shape + s2_flat = s2_img.permute(0, 2, 1, 3, 4).reshape(b * seq, c, h, w).contiguous() + s2_features = self.backbone_s2(s2_flat, anno_s2, anno_mask, return_dict=False) + s2_features = self.head_s2(s2_features[-1]) + s2_features = [s2_features] + + if "s1" in self.sources: + b, c, seq, h, w = s1_img.shape + s1_flat = s1_img.permute(0, 2, 1, 3, 4).reshape(b * seq, c, h, w).contiguous() + s1_features = self.backbone_s1(s1_flat, anno_s1, anno_mask, return_dict=False) + s1_features = self.head_s1(s1_features[-1]) + s1_features = [s1_features] + + hr_features_stage3 = hr_features[-1] + s2_features_stage3 = s2_features[-1] + s1_features_stage3 = s1_features[-1] + + if modality_flags is None: + modality_flags = torch.tensor([[0, 0, 1]] * batch_size, device=hr_img.device, dtype=torch.float32) + + if self.use_modal_vae: + vae_out = self.modality_vae(hr_features_stage3, s2_features_stage3, s1_features_stage3, modality_flags) + hr_features_stage3 = vae_out["hr_out"] + s2_features_stage3 = vae_out["s2_out"] + s1_features_stage3 = vae_out["s1_out"] + output["vae_out"] = vae_out + + _, c3, h3, w3 = hr_features_stage3.shape + hr_tokens = hr_features_stage3.permute(0, 2, 3, 1).reshape(batch_size * h3 * w3, c3).unsqueeze(1) + + _, c3_s2, h3_s2, w3_s2 = s2_features_stage3.shape + s2_tokens = ( + s2_features_stage3.reshape(batch_size, seq_len_s2, c3_s2, h3_s2, w3_s2) + .permute(0, 3, 4, 1, 2) + .reshape(batch_size, h3_s2 * w3_s2, seq_len_s2, c3_s2) + .reshape(batch_size * h3_s2 * w3_s2, seq_len_s2, c3_s2) + .contiguous() + ) + features_stage3 = torch.cat((hr_tokens, s2_tokens), dim=1) + + _, c3_s1, h3_s1, w3_s1 = s1_features_stage3.shape + s1_tokens = ( + s1_features_stage3.reshape(batch_size, seq_len_s1, c3_s1, h3_s1, w3_s1) + .permute(0, 3, 4, 1, 2) + .reshape(batch_size, h3_s1 * w3_s1, seq_len_s1, c3_s1) + .reshape(batch_size * h3_s1 * w3_s1, seq_len_s1, c3_s1) + .contiguous() + ) + features_stage3 = torch.cat((features_stage3, s1_tokens), dim=1) + + fusion_out = self.fusion(features_stage3, return_dict=True) + cls_token = fusion_out.pooler_output.reshape(batch_size, h3, w3, -1).permute(0, 3, 1, 2).contiguous() + + hr_rec_inputs = list(hr_features) + feat_stage1 = hr_rec_inputs[0] + if feat_stage1.shape[-1] == feat_stage1.shape[-2]: + left, right = torch.split(feat_stage1, feat_stage1.shape[-1] // 2, dim=-1) + hr_rec_inputs[0] = torch.cat((left, right), dim=1) + + logits_hr = self.head_rec_hr([*hr_rec_inputs, cls_token]) + if self.config.upsample_results: + logits_hr = F.interpolate(logits_hr.float(), scale_factor=4, mode="bilinear", align_corners=True) + output["logits_hr"] = logits_hr + + if not return_dict: + return output + return SkySensePlusPlusOutput( + logits=logits_hr, + mapped_targets=output.get("mapped_targets"), + idx_2_color=output.get("idx_2_color"), + mask_hr=output.get("mask_hr"), + vae_out=output.get("vae_out"), + ) diff --git a/skysensepp-fewshot-release/modeling_skysensepp_components.py b/skysensepp-fewshot-release/modeling_skysensepp_components.py new file mode 100644 index 0000000000000000000000000000000000000000..9112b628d4a0f27ddb89ab70d50397d470718b5d --- /dev/null +++ b/skysensepp-fewshot-release/modeling_skysensepp_components.py @@ -0,0 +1,238 @@ +"""Shared heads and necks for the full SkySense++ model.""" + +from typing import List, Sequence, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def resize_tensor(x: torch.Tensor, size: Tuple[int, int], align_corners: bool = False) -> torch.Tensor: + return F.interpolate(x, size=size, mode="bilinear", align_corners=align_corners) + + +class ConvModule(nn.Module): + def __init__(self, in_channels: int, out_channels: int, kernel_size: int, padding: int = 0): + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding) + self.bn = nn.BatchNorm2d(out_channels) + self.relu = nn.ReLU(inplace=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.bn(self.conv(x))) + + +class UPHead(nn.Module): + def __init__(self, in_dim: int, out_dim: int, up_scale: int): + super().__init__() + self.decoder = nn.Sequential( + nn.Conv2d(in_dim, up_scale**2 * out_dim, kernel_size=1), + nn.PixelShuffle(up_scale), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.decoder(x) + + +class PPM(nn.ModuleList): + def __init__( + self, + pool_scales: Sequence[int], + in_channels: int, + channels: int, + align_corners: bool = False, + ): + super().__init__() + self.align_corners = align_corners + for pool_scale in pool_scales: + self.append( + nn.Sequential( + nn.AdaptiveAvgPool2d(pool_scale), + ConvModule(in_channels, channels, kernel_size=1), + ) + ) + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + outputs = [] + for module in self: + out = module(x) + out = resize_tensor(out, x.shape[2:], align_corners=self.align_corners) + outputs.append(out) + return outputs + + +class UPerHead(nn.Module): + def __init__( + self, + in_channels: Sequence[int] = (704, 704, 1408, 2816, 1024), + channels: int = 512, + num_classes: int = 65, + pool_scales: Sequence[int] = (1, 2, 3, 6), + dropout_ratio: float = 0.1, + align_corners: bool = False, + ): + super().__init__() + self.in_channels = list(in_channels) + self.channels = channels + self.align_corners = align_corners + self.psp_modules = PPM(pool_scales, self.in_channels[-1], channels, align_corners) + self.bottleneck = ConvModule( + self.in_channels[-1] + len(pool_scales) * channels, + channels, + kernel_size=3, + padding=1, + ) + self.lateral_convs = nn.ModuleList() + self.fpn_convs = nn.ModuleList() + for in_ch in self.in_channels[:-1]: + self.lateral_convs.append(ConvModule(in_ch, channels, kernel_size=1)) + self.fpn_convs.append(ConvModule(channels, channels, kernel_size=3, padding=1)) + self.fpn_bottleneck = ConvModule(len(self.in_channels) * channels, channels, kernel_size=3, padding=1) + self.dropout = nn.Dropout2d(dropout_ratio) if dropout_ratio > 0 else nn.Identity() + self.conv_seg = nn.Conv2d(channels, num_classes, kernel_size=1) + + def psp_forward(self, inputs: List[torch.Tensor]) -> torch.Tensor: + x = inputs[-1] + psp_outs = [x, *self.psp_modules(x)] + return self.bottleneck(torch.cat(psp_outs, dim=1)) + + def forward(self, inputs: List[torch.Tensor]) -> torch.Tensor: + laterals = [conv(inputs[i]) for i, conv in enumerate(self.lateral_convs)] + laterals.append(self.psp_forward(inputs)) + + for i in range(len(laterals) - 1, 0, -1): + laterals[i - 1] = laterals[i - 1] + resize_tensor( + laterals[i], laterals[i - 1].shape[2:], align_corners=self.align_corners + ) + + fpn_outs = [self.fpn_convs[i](laterals[i]) for i in range(len(laterals) - 1)] + fpn_outs.append(laterals[-1]) + for i in range(len(fpn_outs) - 1, 0, -1): + fpn_outs[i] = resize_tensor(fpn_outs[i], fpn_outs[0].shape[2:], align_corners=self.align_corners) + output = self.fpn_bottleneck(torch.cat(fpn_outs, dim=1)) + output = self.dropout(output) + return self.conv_seg(output) + + +class BFloat16UpsampleNearest2d(nn.Module): + def __init__(self, scale_factor: int, mode: str = "bilinear"): + super().__init__() + self.scale_factor = scale_factor + self.mode = mode + + def forward(self, x: torch.Tensor) -> torch.Tensor: + upsampled = F.interpolate(x.float(), scale_factor=self.scale_factor, mode=self.mode) + return upsampled.to(x.dtype) + + +class ConvVQVAEv2(nn.Module): + def __init__(self, input_shape: Tuple[int, int, int], conv_dim: int, z_dim: int, num_tokens: int = 8192, temp: float = 0.9): + super().__init__() + self.temp = temp + self.codebook = nn.Embedding(num_tokens, z_dim) + self.relu = nn.LeakyReLU() + self.pool = nn.AvgPool2d(2) + self.conv1 = nn.Conv2d(input_shape[0], conv_dim, 5, stride=1, padding=2) + self.enc_block1 = nn.Sequential( + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + ) + self.gamma_1 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1))) + self.enc_block2 = nn.Sequential( + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + ) + self.gamma_2 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1))) + self.logit_conv = nn.Conv2d(conv_dim, num_tokens, 1) + self.unpool = BFloat16UpsampleNearest2d(scale_factor=2) + self.conv2 = nn.Conv2d(z_dim, conv_dim, 3, stride=1, padding=1) + self.dec_block1 = nn.Sequential( + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + ) + self.gamma_3 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1))) + self.dec_block2 = nn.Sequential( + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + nn.Conv2d(conv_dim, conv_dim, 3, stride=1, padding=1), + nn.LeakyReLU(), + ) + self.gamma_4 = nn.Parameter(0.001 * torch.ones((1, conv_dim, 1, 1))) + self.rec_conv = nn.Conv2d(conv_dim, input_shape[0], 3, stride=1, padding=1) + + def forward_encoder(self, x: torch.Tensor) -> torch.Tensor: + x = self.relu(self.conv1(x)) + x = x + self.gamma_1 * self.enc_block1(x) + x = self.pool(x) + x = x + self.gamma_2 * self.enc_block2(x) + x = self.pool(x) + return self.logit_conv(x) + + def forward_decoder(self, logits: torch.Tensor): + soft_one_hot = F.softmax(logits * (self.temp * 10), dim=1) + sampled = torch.einsum("bnhw,nd->bdhw", soft_one_hot, self.codebook.weight) + x = self.relu(self.conv2(sampled)) + x = self.unpool(x) + x = x + self.gamma_3 * self.dec_block1(x) + x = self.unpool(x) + x = x + self.gamma_4 * self.dec_block2(x) + return self.rec_conv(x), soft_one_hot + + +class ModalityCompletion(nn.Module): + def __init__( + self, + input_shape_hr: Tuple[int, int, int] = (2816, 32, 16), + input_shape_s2: Tuple[int, int, int] = (2816, 32, 16), + input_shape_s1: Tuple[int, int, int] = (2816, 32, 16), + conv_dim: int = 256, + z_dim: int = 256, + n_codebook: int = 8192, + ): + super().__init__() + self.vae_hr = ConvVQVAEv2(input_shape_hr, conv_dim, z_dim, num_tokens=n_codebook) + self.vae_s2 = ConvVQVAEv2(input_shape_s2, conv_dim, z_dim, num_tokens=n_codebook) + self.vae_s1 = ConvVQVAEv2(input_shape_s1, conv_dim, z_dim, num_tokens=n_codebook) + + def forward( + self, + feat_hr: torch.Tensor, + feat_s2: torch.Tensor, + feat_s1: torch.Tensor, + modality_info: torch.Tensor, + ) -> dict[str, torch.Tensor]: + logits_hr = self.vae_hr.forward_encoder(feat_hr) + logits_s2 = self.vae_s2.forward_encoder(feat_s2) + logits_s1 = self.vae_s1.forward_encoder(feat_s1) + + flag_hr = modality_info[:, 0][:, None, None, None] + flag_s2 = modality_info[:, 1][:, None, None, None] + flag_s1 = modality_info[:, 2][:, None, None, None] + + mean_logits_hr_s2 = logits_hr * flag_hr + logits_s2 * flag_s2 + mean_logits_hr_s1 = logits_hr * flag_hr + logits_s1 * flag_s1 + mean_logits_s1_s2 = logits_s1 * flag_s1 + logits_s2 * flag_s2 + + logits_hr_rec = logits_hr * flag_hr + mean_logits_s1_s2 * (1.0 - flag_hr) + logits_s2_rec = logits_s2 * flag_s2 + mean_logits_hr_s1 * (1.0 - flag_s2) + logits_s1_rec = logits_s1 * flag_s1 + mean_logits_hr_s2 * (1.0 - flag_s1) + + g_hr, _ = self.vae_hr.forward_decoder(logits_hr_rec) + g_s2, _ = self.vae_s2.forward_decoder(logits_s2_rec) + g_s1, _ = self.vae_s1.forward_decoder(logits_s1_rec) + + inv_hr = 1.0 - flag_hr + inv_s2 = 1.0 - flag_s2 + inv_s1 = 1.0 - flag_s1 + + return { + "hr_out": feat_hr * flag_hr + g_hr * inv_hr, + "s2_out": feat_s2 * flag_s2 + g_s2 * inv_s2, + "s1_out": feat_s1 * flag_s1 + g_s1 * inv_s1, + } diff --git a/skysensepp-fewshot-release/modeling_skysensepp_fusion_neck.py b/skysensepp-fewshot-release/modeling_skysensepp_fusion_neck.py new file mode 100644 index 0000000000000000000000000000000000000000..5387a07d8eae4aef4485c8368046deba90ec2277 --- /dev/null +++ b/skysensepp-fewshot-release/modeling_skysensepp_fusion_neck.py @@ -0,0 +1,164 @@ +"""SkySense++ fusion neck (TransformerEncoder) — optional multi-modal fusion module.""" + +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPooling + +from .configuration_skysensepp import SkySensePlusPlusFusionNeckConfig +from .modeling_utils import DropPath, FFN + + +class FusionEncoderLayer(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + feedforward_channels: int, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + qkv_bias: bool = True, + with_cp: bool = False, + ): + super().__init__() + self.with_cp = with_cp + self.norm1 = nn.LayerNorm(embed_dims) + self.attn = nn.MultiheadAttention( + embed_dim=embed_dims, + num_heads=num_heads, + dropout=attn_drop_rate, + bias=qkv_bias, + batch_first=True, + ) + self.proj_drop = nn.Dropout(drop_rate) + self.norm2 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=2, + ffn_drop=drop_rate, + drop_path=drop_path_rate, + act_layer=nn.GELU, + add_identity=True, + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + def _inner_forward(x): + residual = x + x_norm = self.norm1(x) + attn_out, _ = self.attn(x_norm, x_norm, x_norm) + attn_out = self.proj_drop(attn_out) + x = residual + self.drop_path(attn_out) + return self.ffn(self.norm2(x), identity=x) + + if self.with_cp and x.requires_grad: + return cp.checkpoint(_inner_forward, x, use_reentrant=False) + return _inner_forward(x) + + +class SkySensePlusPlusFusionNeckPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusFusionNeckConfig + base_model_prefix = "skysensepp_fusion_neck" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusFusionNeckModel(SkySensePlusPlusFusionNeckPreTrainedModel): + """Fuses per-location multi-modal tokens into a cls-token representation. + + Input shape: ``(batch, num_modalities, input_dims)`` — e.g. concatenated + HR + S2 + S1 stage-3 features with ``input_dims=2816``. + + Output: cls token embedding ``(batch, embed_dims)`` when + ``output_cls_token=True`` (default). + """ + + def __init__(self, config: SkySensePlusPlusFusionNeckConfig): + super().__init__(config) + + # Original checkpoint uses the typo `porj_linear`. + self.porj_linear = nn.Linear(config.input_dims, config.embed_dims) + self.with_cls_token = config.with_cls_token + self.output_cls_token = config.output_cls_token + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + num_layers = config.num_layers + if num_layers > 1: + dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)] + else: + dpr = [0.0] + + self.layers = nn.ModuleList() + for i in range(config.num_layers): + self.layers.append( + FusionEncoderLayer( + embed_dims=config.embed_dims, + num_heads=config.num_heads, + feedforward_channels=config.mlp_ratio * config.embed_dims, + attn_drop_rate=config.attn_drop_rate, + drop_rate=config.drop_rate, + drop_path_rate=dpr[i], + qkv_bias=config.qkv_bias, + with_cp=config.with_cp, + ) + ) + + self.post_init() + + def forward( + self, + hidden_states: torch.Tensor, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + """Forward pass. + + Args: + hidden_states: ``(batch, seq_len, input_dims)`` fused modality tokens. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x = self.porj_linear(hidden_states) + cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + if not self.with_cls_token: + x = x[:, 1:] + + all_hidden_states = () if output_hidden_states else None + for layer in self.layers: + x = layer(x) + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if self.output_cls_token: + pooler = x[:, 0] + last_hidden = pooler.unsqueeze(1) + elif self.with_cls_token: + pooler = None + last_hidden = x[:, 1:] + else: + pooler = None + last_hidden = x + + if not return_dict: + return (last_hidden, pooler) if pooler is not None else (last_hidden,) + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden, + pooler_output=pooler, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-fewshot-release/modeling_skysensepp_swinv2_msl.py b/skysensepp-fewshot-release/modeling_skysensepp_swinv2_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0ab382892dfc61d579821983bc8f6b68b84d17 --- /dev/null +++ b/skysensepp-fewshot-release/modeling_skysensepp_swinv2_msl.py @@ -0,0 +1,343 @@ +"""SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace).""" + +from copy import deepcopy +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig +from .modeling_utils import ( + DropPath, + FFN, + PatchEmbed, + PatchMerging, + ShiftWindowMSA, + to_2tuple, +) + + +class SwinBlockV2(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int = 8, + shift: bool = False, + extra_norm: bool = False, + ffn_ratio: float = 4.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + with_cp: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.with_cp = with_cp + self.extra_norm = extra_norm + self.attn = ShiftWindowMSA( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=window_size, + shift_size=window_size // 2 if shift else 0, + drop_path=drop_path, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + self.norm1 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=int(embed_dims * ffn_ratio), + num_fcs=2, + drop_path=drop_path, + act_layer=nn.GELU, + add_identity=False, + ) + self.norm2 = nn.LayerNorm(embed_dims) + if self.extra_norm: + self.norm3 = nn.LayerNorm(embed_dims) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + def _inner_forward(x): + identity = x + x = self.attn(x, hw_shape) + x = self.norm1(x) + x = x + identity + + identity = x + x = self.ffn(x) + x = self.norm2(x) + x = x + identity + + if self.extra_norm: + x = self.norm3(x) + return x + + if self.with_cp and x.requires_grad: + x = cp.checkpoint(_inner_forward, x, use_reentrant=False) + else: + x = _inner_forward(x) + return x + + +class SwinBlockV2Sequence(nn.Module): + def __init__( + self, + embed_dims: int, + depth: int, + num_heads: int, + window_size: int = 8, + downsample: bool = False, + drop_paths: Union[Sequence[float], float] = 0.0, + with_cp: bool = False, + pad_small_map: bool = False, + extra_norm_every_n_blocks: int = 0, + pretrained_window_size: int = 0, + is_post_norm_downsample: bool = True, + ): + super().__init__() + if not isinstance(drop_paths, Sequence): + drop_paths = [drop_paths] * depth + + if downsample: + self.out_channels = 2 * embed_dims + self.downsample = PatchMerging( + in_channels=embed_dims, + out_channels=self.out_channels, + is_post_norm=is_post_norm_downsample, + ) + else: + self.out_channels = embed_dims + self.downsample = None + + self.blocks = nn.ModuleList() + for i in range(depth): + extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0 + self.blocks.append( + SwinBlockV2( + embed_dims=self.out_channels, + num_heads=num_heads, + window_size=window_size, + shift=(i % 2 == 1), + extra_norm=extra_norm, + drop_path=drop_paths[i], + with_cp=with_cp, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + ) + + def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + if self.downsample is not None: + x, out_shape = self.downsample(x, in_shape) + else: + out_shape = in_shape + + for block in self.blocks: + x = block(x, out_shape) + return x, out_shape + + +class ProjMHSA(nn.Module): + """Projected multi-head self-attention used in SkySense++ HR backbone.""" + + def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True): + super().__init__() + self.proj_in = nn.Linear(embed_dims, proj_dims) + self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias) + self.proj_out = nn.Linear(proj_dims, embed_dims) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj_in(x) + x, _ = self.attn(x, x, x) + return self.proj_out(x) + + +class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusSwinV2MSLConfig + base_model_prefix = "skysensepp_swinv2_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel): + """SkySense++ HR backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig): + super().__init__(config) + + self.num_layers = len(config.depths) + self.out_indices = config.out_indices + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.patch_size = config.patch_size + + if isinstance(config.window_size, int): + window_sizes = [config.window_size] * self.num_layers + else: + window_sizes = list(config.window_size) + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm, + input_size=config.img_size, + ) + + self.use_abs_pos_embed = config.use_abs_pos_embed + if self.use_abs_pos_embed: + patch_resolution = self.patch_embed.init_out_size + num_patches = patch_resolution[0] * patch_resolution[1] + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + total_depth = sum(config.depths) + if total_depth > 1: + dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)] + else: + dpr = [0.0] + + self.stages = nn.ModuleList() + embed_dims_list = [config.embed_dims] + for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)): + stage = SwinBlockV2Sequence( + embed_dims=embed_dims_list[-1], + depth=depth, + num_heads=num_heads, + window_size=window_sizes[i], + downsample=(i > 0), + drop_paths=dpr[:depth], + with_cp=config.with_cp, + pad_small_map=config.pad_small_map, + extra_norm_every_n_blocks=config.extra_norm_every_n_blocks, + pretrained_window_size=config.pretrained_window_sizes[i], + is_post_norm_downsample=config.is_post_norm_downsample, + ) + self.stages.append(stage) + dpr = dpr[depth:] + embed_dims_list.append(stage.out_channels) + + for i in self.out_indices: + self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1])) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = ProjMHSA(352, 256, num_heads=16) + self.attn2 = ProjMHSA(704, 512, num_heads=16) + self.attn3 = ProjMHSA(1408, 1024, num_heads=16) + self.norm_attn = nn.LayerNorm(1408) + + self.post_init() + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + if self.use_abs_pos_embed: + x = x + self.absolute_pos_embed + x = self.drop_after_pos(x) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, stage in enumerate(self.stages): + x, hw_shape = stage(x, hw_shape) + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + x = x + attention_blocks[i](x) + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + out = norm_layer(x) + out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous() + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-fewshot-release/modeling_skysensepp_vit_msl.py b/skysensepp-fewshot-release/modeling_skysensepp_vit_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..6e56ac39c612c2eb87baf00b1e5afc071dc4d060 --- /dev/null +++ b/skysensepp-fewshot-release/modeling_skysensepp_vit_msl.py @@ -0,0 +1,265 @@ +"""SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace).""" + +import math +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig +from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + feedforward_channels: int, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + num_fcs: int = 2, + qkv_bias: bool = True, + with_cp: bool = False, + ): + super().__init__() + self.with_cp = with_cp + self.norm1 = nn.LayerNorm(embed_dims) + self.attn = nn.MultiheadAttention( + embed_dim=embed_dims, + num_heads=num_heads, + dropout=attn_drop_rate, + bias=qkv_bias, + batch_first=True, + ) + self.proj_drop = nn.Dropout(drop_rate) + self.norm2 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=num_fcs, + ffn_drop=drop_rate, + drop_path=drop_path_rate, + act_layer=nn.GELU, + add_identity=True, + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + def _inner_forward(x): + residual = x + x_norm = self.norm1(x) + attn_out, _ = self.attn(x_norm, x_norm, x_norm) + attn_out = self.proj_drop(attn_out) + x = residual + self.drop_path(attn_out) + return self.ffn(self.norm2(x), identity=x) + + if self.with_cp and x.requires_grad: + return cp.checkpoint(_inner_forward, x, use_reentrant=False) + return _inner_forward(x) + + +class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusViTMSLConfig + base_model_prefix = "skysensepp_vit_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel): + """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusViTMSLConfig): + super().__init__(config) + + img_size = to_2tuple(config.img_size) + self.img_size = img_size + self.patch_size = config.patch_size + self.with_cls_token = config.with_cls_token + self.output_cls_token = config.output_cls_token + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.interpolate_mode = "bicubic" + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm if config.patch_norm else None, + ) + + num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size) + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + out_indices = list(config.out_indices) + self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices] + + num_layers = config.num_layers + if num_layers > 1: + dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)] + else: + dpr = [0.0] + + self.layers = nn.ModuleList() + for i in range(config.num_layers): + self.layers.append( + TransformerEncoderLayer( + embed_dims=config.embed_dims, + num_heads=config.num_heads, + feedforward_channels=config.mlp_ratio * config.embed_dims, + attn_drop_rate=config.attn_drop_rate, + drop_rate=config.drop_rate, + drop_path_rate=dpr[i], + num_fcs=2, + qkv_bias=config.qkv_bias, + with_cp=config.with_cp, + ) + ) + + self.final_norm = config.final_norm + if config.final_norm: + self.norm = nn.LayerNorm(config.embed_dims) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.norm_attn = nn.LayerNorm(config.embed_dims) + + self.post_init() + + @staticmethod + def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"): + pos_h, pos_w = pos_shape + pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :] + pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) + pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode) + return torch.flatten(pos_embed_weight, 2).transpose(1, 2) + + def _pos_embedding(self, patched_img, hw_shape, pos_embed): + x_len, pos_len = patched_img.shape[1], pos_embed.shape[1] + if x_len != pos_len: + pos_h = self.img_size[0] // self.patch_size + pos_w = self.img_size[1] // self.patch_size + pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode) + return self.drop_after_pos(patched_img + pos_embed) + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + x = self._pos_embedding(x, hw_shape, self.pos_embed) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, layer in enumerate(self.layers): + x = layer(x) + + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + attn_out, _ = attention_blocks[i](x, x, x) + x = x + attn_out + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm: + x = self.norm(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + out = x + out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous() + if self.output_cls_token: + out = [out, x[:, 0]] + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-fewshot-release/modeling_utils.py b/skysensepp-fewshot-release/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..93feae77f3a3c46e65167f2ed312a25b7cd4ad3a --- /dev/null +++ b/skysensepp-fewshot-release/modeling_utils.py @@ -0,0 +1,557 @@ +"""SkySense: Pure PyTorch + HuggingFace Transformers implementation. + +Shared utility modules used across SkySense model implementations. +""" + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def to_2tuple(x): + """Convert to a 2-tuple.""" + if isinstance(x, (list, tuple)): + return tuple(x) + return (x, x) + + +class DropPath(nn.Module): + """Drop paths (stochastic depth) per sample. + + Args: + drop_prob (float): Probability of dropping a path. Default: 0.0. + """ + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor = torch.floor(random_tensor + keep_prob) + output = x / keep_prob * random_tensor + return output + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding using Conv2d. + + Args: + in_channels (int): Number of input channels. Default: 3. + embed_dims (int): Embedding dimension. Default: 96. + kernel_size (int): Kernel size of the projection. Default: 4. + stride (int): Stride of the projection. Default: 4. + padding (int): Padding of the projection. Default: 0. + norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm. + input_size (int or tuple or None): Input resolution for calculating output size. + """ + + def __init__( + self, + in_channels: int = 3, + embed_dims: int = 96, + kernel_size: int = 4, + stride: int = 4, + padding: int = 0, + norm_layer: Optional[type] = nn.LayerNorm, + input_size: Optional[int] = None, + ): + super().__init__() + self.projection = nn.Conv2d( + in_channels, embed_dims, + kernel_size=kernel_size, stride=stride, padding=padding, + ) + self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity() + + # Compute init output size if input_size is given + if input_size is not None: + input_size = to_2tuple(input_size) + self.init_out_size = ( + (input_size[0] - kernel_size + 2 * padding) // stride + 1, + (input_size[1] - kernel_size + 2 * padding) // stride + 1, + ) + else: + self.init_out_size = None + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]: + x = self.projection(x) # (B, C, H, W) + out_size = (x.shape[2], x.shape[3]) + x = x.flatten(2).transpose(1, 2) # (B, H*W, C) + x = self.norm(x) + return x, out_size + + +class FFN(nn.Module): + """Feed-Forward Network. + + Args: + embed_dims (int): Input dimension. + feedforward_channels (int): Hidden dimension. + num_fcs (int): Number of FC layers. Default: 2. + ffn_drop (float): Dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + act_layer (nn.Module): Activation layer class. Default: nn.GELU. + add_identity (bool): Whether to add identity connection. Default: True. + """ + + def __init__( + self, + embed_dims: int, + feedforward_channels: int, + num_fcs: int = 2, + ffn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: type = nn.GELU, + add_identity: bool = True, + ): + super().__init__() + assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}" + self.embed_dims = embed_dims + self.feedforward_channels = feedforward_channels + self.add_identity = add_identity + + layers = [] + in_channels = embed_dims + for i in range(num_fcs - 1): + layers.append(nn.Linear(in_channels, feedforward_channels)) + layers.append(act_layer()) + layers.append(nn.Dropout(ffn_drop)) + in_channels = feedforward_channels + layers.append(nn.Linear(feedforward_channels, embed_dims)) + layers.append(nn.Dropout(ffn_drop)) + self.layers = nn.Sequential(*layers) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor: + out = self.layers(x) + out = self.drop_path(out) + if self.add_identity: + if identity is None: + identity = x + out = out + identity + return out + + +class WindowMSAV2(nn.Module): + """Window-based Multi-head Self-Attention for Swin Transformer V2. + + Uses cosine attention and log-spaced continuous position bias (log-CPB). + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): Window size (Wh, Ww). + pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0). + qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Output projection dropout rate. Default: 0.0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: Tuple[int, int], + pretrained_window_size: Tuple[int, int] = (0, 0), + qkv_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ): + super().__init__() + self.embed_dims = embed_dims + self.num_heads = num_heads + self.window_size = window_size + self.pretrained_window_size = pretrained_window_size + + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1)))) + + # MLP for continuous relative position bias (log-CPB) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) + + # Build relative coords table + self._build_relative_coords_table() + # Build relative position index + self._build_relative_position_index() + + self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False) + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(embed_dims)) + self.v_bias = nn.Parameter(torch.zeros(embed_dims)) + else: + self.q_bias = None + self.v_bias = None + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(embed_dims, embed_dims) + self.proj_drop = nn.Dropout(proj_drop) + self.softmax = nn.Softmax(dim=-1) + + def _build_relative_coords_table(self): + """Build the relative coordinates table for log-CPB.""" + Wh, Ww = self.window_size + # Table of relative coordinates + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2) + + # Normalize to [-1, 1] and apply log-scale + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 # normalize to -8, 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + self.register_buffer("relative_coords_table", coords_table) + + def _build_relative_position_index(self): + """Build the pairwise relative position index for each window token.""" + Wh, Ww = self.window_size + coords_h = torch.arange(Wh) + coords_w = torch.arange(Ww) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij')) + coords_flatten = coords.view(2, -1) + + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # (2, Wh*Ww, Wh*Ww) + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += Wh - 1 + relative_coords[:, :, 1] += Ww - 1 + relative_coords[:, :, 0] *= 2 * Ww - 1 + relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww) + self.register_buffer("relative_position_index", relative_position_index) + + def _compute_position_bias(self, N): + """Compute relative position bias, supporting dynamic window sizes. + + The log-CPB (Continuous Position Bias) MLP can generalize to any window + size by computing bias from normalized relative coordinates. + """ + init_N = self.window_size[0] * self.window_size[1] + if N == init_N: + # Use pre-built tables + relative_position_bias_table = self.cpb_mlp( + self.relative_coords_table + ).view(-1, self.num_heads) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view(N, N, -1) + else: + # Dynamic: compute for actual window size on-the-fly + Wh = Ww = int(math.sqrt(N)) + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + # Build position index for actual window size + ch = torch.arange(Wh, device=self.logit_scale.device) + cw = torch.arange(Ww, device=self.logit_scale.device) + coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij')) + coords_flat = coords.view(2, -1) + rel = coords_flat[:, :, None] - coords_flat[:, None, :] + rel = rel.permute(1, 2, 0).contiguous() + rel[:, :, 0] += Wh - 1 + rel[:, :, 1] += Ww - 1 + rel[:, :, 0] *= 2 * Ww - 1 + pos_index = rel.sum(-1) + + bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads) + relative_position_bias = bias_table[ + pos_index.view(-1) + ].view(N, N, -1) + + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + return relative_position_bias + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: (num_windows*B, N, C) where N = Wh*Ww + mask: (num_windows, N, N) or None + """ + B_, N, C = x.shape + + # Compute QKV with bias + if self.q_bias is not None: + qkv_bias = torch.cat( + (self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias)) + qkv = F.linear(x, self.qkv.weight, qkv_bias) + else: + qkv = self.qkv(x) + + qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + + # Cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=math.log(1.0 / 0.01) + ).exp() + attn = attn * logit_scale + + # Log-CPB relative position bias (supports dynamic window sizes) + relative_position_bias = self._compute_position_bias(N) + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + attn = attn + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + + attn = self.softmax(attn) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class ShiftWindowMSA(nn.Module): + """Shifted Window Multi-head Self-Attention. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. Default: 0. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Projection dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + pad_small_map (bool): Pad small feature maps to window size. Default: False. + pretrained_window_size (int): Pretrained window size. Default: 0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int, + shift_size: int = 0, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.window_size = window_size + self.shift_size = shift_size + self.pad_small_map = pad_small_map + + self.w_msa = WindowMSAV2( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=to_2tuple(window_size), + pretrained_window_size=to_2tuple(pretrained_window_size), + attn_drop=attn_drop, + proj_drop=proj_drop, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W, f"Input length {L} != H*W ({H}*{W})" + + x = x.view(B, H, W, C) + + window_size = self.window_size + shift_size = self.shift_size + + # Pad or shrink window + if self.pad_small_map: + pad_r = (window_size - W % window_size) % window_size + pad_b = (window_size - H % window_size) % window_size + x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b)) + _, Hp, Wp, _ = x.shape + else: + Hp, Wp = H, W + if window_size > Hp: + window_size = Hp + shift_size = 0 + if window_size > Wp: + window_size = Wp + shift_size = 0 + + # Compute attention mask for SW-MSA + attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device) + + # Cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) + + # Partition windows + x_windows = self._window_partition(x, window_size) + # (num_windows*B, window_size*window_size, C) + + # W-MSA/SW-MSA + attn_windows = self.w_msa(x_windows, mask=attn_mask) + + # Merge windows + x = self._window_reverse(attn_windows, window_size, Hp, Wp) + + # Reverse cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) + + if self.pad_small_map and (pad_r > 0 or pad_b > 0): + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + x = self.drop_path(x) + return x + + @staticmethod + def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor: + """Partition into non-overlapping windows.""" + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous() + windows = windows.view(-1, window_size * window_size, C) + return windows + + @staticmethod + def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor: + """Reverse window partition.""" + B_nW = windows.shape[0] + nH = H // window_size + nW = W // window_size + B = B_nW // (nH * nW) + x = windows.view(B, nH, nW, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous() + x = x.view(B, H, W, -1) + return x + + @staticmethod + def _compute_attn_mask(H, W, window_size, shift_size, device): + """Compute attention mask for shifted window attention.""" + if shift_size <= 0: + return None + img_mask = torch.zeros((1, H, W, 1), device=device) + h_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + w_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + # Partition mask + mask_windows = img_mask.view( + 1, H // window_size, window_size, W // window_size, window_size, 1 + ) + mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous() + mask_windows = mask_windows.view(-1, window_size * window_size) + + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0) + attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0) + return attn_mask + + +class PatchMerging(nn.Module): + """Patch Merging Layer for downsampling (2x). + + Args: + in_channels (int): Input channels. + out_channels (int): Output channels. + norm_layer (type): Normalization layer. Default: nn.LayerNorm. + is_post_norm (bool): Apply norm after linear. Default: True. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + norm_layer: type = nn.LayerNorm, + is_post_norm: bool = True, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.is_post_norm = is_post_norm + self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False) + if is_post_norm: + self.norm = norm_layer(out_channels) + else: + self.norm = norm_layer(4 * in_channels) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W + + x = x.view(B, H, W, C) + + # Pad if needed + pad_h = H % 2 + pad_w = W % 2 + if pad_h or pad_w: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], dim=-1) + + out_h = (H + pad_h) // 2 + out_w = (W + pad_w) // 2 + x = x.view(B, out_h * out_w, 4 * C) + + if self.is_post_norm: + x = self.reduction(x) + x = self.norm(x) + else: + x = self.norm(x) + x = self.reduction(x) + + return x, (out_h, out_w) diff --git a/skysensepp-fewshot-release/pipeline_skysensepp.py b/skysensepp-fewshot-release/pipeline_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9a5354361c9a08abd1a82b3df8d4ec678a209b --- /dev/null +++ b/skysensepp-fewshot-release/pipeline_skysensepp.py @@ -0,0 +1,86 @@ +"""Custom HuggingFace pipeline for SkySense++ MSL feature extraction.""" + +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline): + """Pipeline for SkySense++ MSL backbones. + + Expects image tensors plus semantic annotation maps (class indices). + """ + + def _sanitize_parameters( + self, + annotation=None, + mask=None, + output_hidden_states=None, + **kwargs, + ): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + + if annotation is not None: + preprocess_params["annotation"] = annotation + if mask is not None: + forward_params["mask"] = mask + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + + return preprocess_params, forward_params, postprocess_params + + def preprocess( + self, + pixel_values: Any, + annotation: Optional[Any] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if isinstance(pixel_values, dict): + annotation = pixel_values.get("annotation", annotation) + pixel_values = pixel_values.get("pixel_values", pixel_values) + + if isinstance(pixel_values, np.ndarray): + pixel_values = torch.from_numpy(pixel_values).float() + elif not isinstance(pixel_values, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}" + ) + + if annotation is None: + raise ValueError("SkySense++ MSL models require an `annotation` semantic map.") + + if isinstance(annotation, np.ndarray): + annotation = torch.from_numpy(annotation).long() + elif not isinstance(annotation, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for annotation, got {type(annotation)}" + ) + + if pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + if annotation.ndim == 2: + annotation = annotation.unsqueeze(0) + + return {"pixel_values": pixel_values, "annotation": annotation} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + pixel_values=model_inputs["pixel_values"], + annotation=model_inputs["annotation"], + mask=kwargs.get("mask"), + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"last_hidden_state": outputs.last_hidden_state} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-fewshot-release/pipeline_skysensepp_fewshot.py b/skysensepp-fewshot-release/pipeline_skysensepp_fewshot.py new file mode 100644 index 0000000000000000000000000000000000000000..696c4b999b45062efb5fd46aebdec0c76bb226a9 --- /dev/null +++ b/skysensepp-fewshot-release/pipeline_skysensepp_fewshot.py @@ -0,0 +1,132 @@ +"""HuggingFace pipeline for SkySense++ few-shot / 1-shot segmentation.""" + +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusFewShotPipeline(Pipeline): + """1-shot segmentation pipeline for the full SkySense++ release model. + + Expects vertically stacked prompt+query tensors: + - ``hr_img``: (B, 3, 1024, 512) — prompt on top, query on bottom + - ``s2_img`` / ``s1_img``: (B, C, seq, 32, 32) — stacked along height + - ``targets``: RGB annotation map (B, 3, 1024, 512), ImageNet-normalized + - ``anno_mask``: (B, 8, 4) with bottom half = 1 (query region) + """ + + def _sanitize_parameters( + self, + s2_img=None, + s1_img=None, + targets=None, + anno_mask=None, + modality_flags=None, + extract_query_only=None, + **kwargs, + ): + preprocess_params = {} + forward_params = {} + postprocess_params = {"extract_query_only": True if extract_query_only is None else extract_query_only} + + if s2_img is not None: + preprocess_params["s2_img"] = s2_img + if s1_img is not None: + preprocess_params["s1_img"] = s1_img + if targets is not None: + preprocess_params["targets"] = targets + if anno_mask is not None: + preprocess_params["anno_mask"] = anno_mask + if modality_flags is not None: + forward_params["modality_flags"] = modality_flags + + return preprocess_params, forward_params, postprocess_params + + def _to_tensor(self, value: Any, dtype: torch.dtype) -> torch.Tensor: + if isinstance(value, np.ndarray): + return torch.from_numpy(value).to(dtype=dtype) + if isinstance(value, torch.Tensor): + return value.to(dtype=dtype) + raise TypeError(f"Expected tensor or ndarray, got {type(value)}") + + def preprocess( + self, + hr_img: Any, + s2_img: Optional[Any] = None, + s1_img: Optional[Any] = None, + targets: Optional[Any] = None, + anno_mask: Optional[Any] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if isinstance(hr_img, dict): + payload = hr_img + hr_img = payload.get("hr_img", payload.get("pixel_values")) + s2_img = payload.get("s2_img", s2_img) + s1_img = payload.get("s1_img", s1_img) + targets = payload.get("targets", targets) + anno_mask = payload.get("anno_mask", anno_mask) + + hr_img = self._to_tensor(hr_img, torch.float32) + if hr_img.ndim == 3: + hr_img = hr_img.unsqueeze(0) + + if s2_img is None or s1_img is None or targets is None: + raise ValueError("Few-shot pipeline requires hr_img, s2_img, s1_img, and targets.") + + s2_img = self._to_tensor(s2_img, torch.float32) + s1_img = self._to_tensor(s1_img, torch.float32) + targets = self._to_tensor(targets, torch.float32) + + if s2_img.ndim == 4: + s2_img = s2_img.unsqueeze(0) + if s1_img.ndim == 4: + s1_img = s1_img.unsqueeze(0) + if targets.ndim == 3: + targets = targets.unsqueeze(0) + + if anno_mask is None: + batch_size = hr_img.shape[0] + anno_mask = torch.zeros(batch_size, 8, 4, dtype=torch.long) + anno_mask[:, 4:, :] = 1 + else: + anno_mask = self._to_tensor(anno_mask, torch.long) + if anno_mask.ndim == 2: + anno_mask = anno_mask.unsqueeze(0) + + return { + "hr_img": hr_img, + "s2_img": s2_img, + "s1_img": s1_img, + "targets": targets, + "anno_mask": anno_mask, + } + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + hr_img=model_inputs["hr_img"], + s2_img=model_inputs["s2_img"], + s1_img=model_inputs["s1_img"], + targets=model_inputs["targets"], + anno_mask=model_inputs["anno_mask"], + modality_flags=kwargs.get("modality_flags"), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess( + self, + model_outputs: Dict[str, Any], + extract_query_only: bool = True, + ) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + logits = outputs.logits + if extract_query_only and logits is not None: + logits = logits[:, :, logits.shape[2] // 2 :, :] + return { + "logits": logits, + "mapped_targets": outputs.mapped_targets, + "idx_2_color": outputs.idx_2_color, + } diff --git a/skysensepp-fewshot-release/pipeline_skysensepp_fusion.py b/skysensepp-fewshot-release/pipeline_skysensepp_fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..80f070fa9fd3f18a053caa1afef9dc142bce2597 --- /dev/null +++ b/skysensepp-fewshot-release/pipeline_skysensepp_fusion.py @@ -0,0 +1,53 @@ +"""Optional pipeline for SkySense++ fusion neck.""" + +from typing import Any, Dict + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusFusionNeckPipeline(Pipeline): + """Pipeline for the optional SkySense++ fusion neck module. + + Expects concatenated multi-modal tokens per spatial location: + ``(batch, num_modalities, input_dims)``. + """ + + def _sanitize_parameters(self, output_hidden_states=None, **kwargs): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + return preprocess_params, forward_params, postprocess_params + + def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]: + if isinstance(hidden_states, dict): + hidden_states = hidden_states["hidden_states"] + + if isinstance(hidden_states, np.ndarray): + hidden_states = torch.from_numpy(hidden_states).float() + elif not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}" + ) + if hidden_states.ndim == 2: + hidden_states = hidden_states.unsqueeze(0) + return {"hidden_states": hidden_states} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + hidden_states=model_inputs["hidden_states"], + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"pooler_output": outputs.pooler_output} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-fusion-neck/config.json b/skysensepp-fusion-neck/config.json new file mode 100644 index 0000000000000000000000000000000000000000..26a9e4e3fd42edb187564e723843d2653d5e7cc5 --- /dev/null +++ b/skysensepp-fusion-neck/config.json @@ -0,0 +1,47 @@ +{ + "return_dict": true, + "output_hidden_states": false, + "dtype": "float32", + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": [ + "SkySensePlusPlusFusionNeckModel" + ], + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "transformers_version": "5.0.0", + "input_dims": 2816, + "embed_dims": 1024, + "num_layers": 24, + "num_heads": 16, + "mlp_ratio": 4, + "qkv_bias": true, + "drop_rate": 0.0, + "attn_drop_rate": 0.0, + "drop_path_rate": 0.3, + "with_cls_token": true, + "output_cls_token": true, + "with_cp": false, + "model_type": "skysensepp_fusion_neck", + "output_attentions": false, + "auto_map": { + "AutoConfig": "configuration_skysensepp.SkySensePlusPlusFusionNeckConfig", + "AutoModel": "modeling_skysensepp_fusion_neck.SkySensePlusPlusFusionNeckModel" + }, + "custom_pipelines": { + "skysensepp-fusion": { + "impl": "pipeline_skysensepp_fusion.SkySensePlusPlusFusionNeckPipeline", + "pt": [ + "AutoModel" + ] + } + } +} diff --git a/skysensepp-fusion-neck/configuration_skysensepp.py b/skysensepp-fusion-neck/configuration_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..ca9ce69e042882c28ad5046d81e33a4fe99383ce --- /dev/null +++ b/skysensepp-fusion-neck/configuration_skysensepp.py @@ -0,0 +1,165 @@ +"""Configuration classes for SkySense++ MSL backbones.""" + +from transformers import PretrainedConfig + + +class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig): + """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical).""" + + model_type = "skysensepp_swinv2_msl" + + arch_zoo = { + "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0}, + "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0}, + "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6}, + "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6}, + } + + def __init__( + self, + arch="huge", + img_size=512, + patch_size=4, + in_channels=3, + window_size=8, + drop_rate=0.0, + drop_path_rate=0.2, + out_indices=(0, 1, 2, 3), + use_abs_pos_embed=False, + with_cp=False, + pad_small_map=False, + pretrained_window_sizes=(0, 0, 0, 0), + is_post_norm_downsample=True, + vocabulary_size=64, + merge_stage=2, + use_attn=True, + **kwargs, + ): + super().__init__(**kwargs) + + arch = arch.lower() + if arch not in self.arch_zoo: + raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}") + arch_settings = self.arch_zoo[arch] + + self.arch = arch + self.embed_dims = arch_settings["embed_dims"] + self.depths = arch_settings["depths"] + self.num_heads = arch_settings["num_heads"] + self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"] + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.window_size = window_size + self.drop_rate = drop_rate + self.drop_path_rate = drop_path_rate + self.out_indices = list(out_indices) + self.use_abs_pos_embed = use_abs_pos_embed + self.with_cp = with_cp + self.pad_small_map = pad_small_map + self.pretrained_window_sizes = list(pretrained_window_sizes) + self.is_post_norm_downsample = is_post_norm_downsample + + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + + +class SkySensePlusPlusViTMSLConfig(PretrainedConfig): + """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1).""" + + model_type = "skysensepp_vit_msl" + + def __init__( + self, + img_size=16, + patch_size=4, + in_channels=10, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + out_indices=(5, 11, 17, 23), + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=False, + output_cls_token=False, + patch_norm=False, + final_norm=False, + with_cp=False, + vocabulary_size=64, + merge_stage=4, + use_attn=False, + modality="s2", + **kwargs, + ): + super().__init__(**kwargs) + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.out_indices = list(out_indices) + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.patch_norm = patch_norm + self.final_norm = final_norm + self.with_cp = with_cp + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + self.modality = modality + + +class SkySensePlusPlusFusionNeckConfig(PretrainedConfig): + """Configuration for SkySense++ multi-modal fusion neck (TransformerEncoder). + + Optional component — not used by default backbone checkpoints. + Fuses concatenated HR/S2/S1 stage-3 features (2816-dim) via a ViT encoder + with cls token output (1024-dim). + """ + + model_type = "skysensepp_fusion_neck" + + def __init__( + self, + input_dims=2816, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=True, + output_cls_token=True, + with_cp=False, + **kwargs, + ): + super().__init__(**kwargs) + self.input_dims = input_dims + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.with_cp = with_cp diff --git a/skysensepp-fusion-neck/conversion_manifest.json b/skysensepp-fusion-neck/conversion_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..6fbf0c0e8354a82ff005303b5c7466ab73c6d695 --- /dev/null +++ b/skysensepp-fusion-neck/conversion_manifest.json @@ -0,0 +1,301 @@ +{ + "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release.ckpt", + "modality": "fusion", + "model_class": "SkySensePlusPlusFusionNeckModel", + "num_tensors": 291, + "missing_keys": [], + "unexpected_keys": [], + "tensor_names": [ + "cls_token", + "layers.0.attn.in_proj_bias", + "layers.0.attn.in_proj_weight", + "layers.0.attn.out_proj.bias", + "layers.0.attn.out_proj.weight", + "layers.0.ffn.layers.0.bias", + "layers.0.ffn.layers.0.weight", + "layers.0.ffn.layers.3.bias", + "layers.0.ffn.layers.3.weight", + "layers.0.norm1.bias", + "layers.0.norm1.weight", + "layers.0.norm2.bias", + "layers.0.norm2.weight", + "layers.1.attn.in_proj_bias", + "layers.1.attn.in_proj_weight", + "layers.1.attn.out_proj.bias", + "layers.1.attn.out_proj.weight", + "layers.1.ffn.layers.0.bias", + "layers.1.ffn.layers.0.weight", + "layers.1.ffn.layers.3.bias", + "layers.1.ffn.layers.3.weight", + "layers.1.norm1.bias", + "layers.1.norm1.weight", + "layers.1.norm2.bias", + "layers.1.norm2.weight", + "layers.10.attn.in_proj_bias", + "layers.10.attn.in_proj_weight", + "layers.10.attn.out_proj.bias", + "layers.10.attn.out_proj.weight", + "layers.10.ffn.layers.0.bias", + "layers.10.ffn.layers.0.weight", + "layers.10.ffn.layers.3.bias", + "layers.10.ffn.layers.3.weight", + "layers.10.norm1.bias", + "layers.10.norm1.weight", + "layers.10.norm2.bias", + "layers.10.norm2.weight", + "layers.11.attn.in_proj_bias", + "layers.11.attn.in_proj_weight", + "layers.11.attn.out_proj.bias", + "layers.11.attn.out_proj.weight", + "layers.11.ffn.layers.0.bias", + "layers.11.ffn.layers.0.weight", + "layers.11.ffn.layers.3.bias", + "layers.11.ffn.layers.3.weight", + "layers.11.norm1.bias", + "layers.11.norm1.weight", + "layers.11.norm2.bias", + "layers.11.norm2.weight", + "layers.12.attn.in_proj_bias", + "layers.12.attn.in_proj_weight", + "layers.12.attn.out_proj.bias", + "layers.12.attn.out_proj.weight", + "layers.12.ffn.layers.0.bias", + "layers.12.ffn.layers.0.weight", + "layers.12.ffn.layers.3.bias", + "layers.12.ffn.layers.3.weight", + "layers.12.norm1.bias", + "layers.12.norm1.weight", + "layers.12.norm2.bias", + "layers.12.norm2.weight", + "layers.13.attn.in_proj_bias", + "layers.13.attn.in_proj_weight", + "layers.13.attn.out_proj.bias", + "layers.13.attn.out_proj.weight", + "layers.13.ffn.layers.0.bias", + "layers.13.ffn.layers.0.weight", + "layers.13.ffn.layers.3.bias", + "layers.13.ffn.layers.3.weight", + "layers.13.norm1.bias", + "layers.13.norm1.weight", + "layers.13.norm2.bias", + "layers.13.norm2.weight", + "layers.14.attn.in_proj_bias", + "layers.14.attn.in_proj_weight", + "layers.14.attn.out_proj.bias", + "layers.14.attn.out_proj.weight", + "layers.14.ffn.layers.0.bias", + "layers.14.ffn.layers.0.weight", + "layers.14.ffn.layers.3.bias", + "layers.14.ffn.layers.3.weight", + "layers.14.norm1.bias", + "layers.14.norm1.weight", + "layers.14.norm2.bias", + "layers.14.norm2.weight", + "layers.15.attn.in_proj_bias", + "layers.15.attn.in_proj_weight", + "layers.15.attn.out_proj.bias", + "layers.15.attn.out_proj.weight", + "layers.15.ffn.layers.0.bias", + "layers.15.ffn.layers.0.weight", + "layers.15.ffn.layers.3.bias", + "layers.15.ffn.layers.3.weight", + "layers.15.norm1.bias", + "layers.15.norm1.weight", + "layers.15.norm2.bias", + "layers.15.norm2.weight", + "layers.16.attn.in_proj_bias", + "layers.16.attn.in_proj_weight", + "layers.16.attn.out_proj.bias", + "layers.16.attn.out_proj.weight", + "layers.16.ffn.layers.0.bias", + "layers.16.ffn.layers.0.weight", + "layers.16.ffn.layers.3.bias", + "layers.16.ffn.layers.3.weight", + "layers.16.norm1.bias", + "layers.16.norm1.weight", + "layers.16.norm2.bias", + "layers.16.norm2.weight", + "layers.17.attn.in_proj_bias", + "layers.17.attn.in_proj_weight", + "layers.17.attn.out_proj.bias", + "layers.17.attn.out_proj.weight", + "layers.17.ffn.layers.0.bias", + "layers.17.ffn.layers.0.weight", + "layers.17.ffn.layers.3.bias", + "layers.17.ffn.layers.3.weight", + "layers.17.norm1.bias", + "layers.17.norm1.weight", + "layers.17.norm2.bias", + "layers.17.norm2.weight", + "layers.18.attn.in_proj_bias", + "layers.18.attn.in_proj_weight", + "layers.18.attn.out_proj.bias", + "layers.18.attn.out_proj.weight", + "layers.18.ffn.layers.0.bias", + "layers.18.ffn.layers.0.weight", + "layers.18.ffn.layers.3.bias", + "layers.18.ffn.layers.3.weight", + "layers.18.norm1.bias", + "layers.18.norm1.weight", + "layers.18.norm2.bias", + "layers.18.norm2.weight", + "layers.19.attn.in_proj_bias", + "layers.19.attn.in_proj_weight", + "layers.19.attn.out_proj.bias", + "layers.19.attn.out_proj.weight", + "layers.19.ffn.layers.0.bias", + "layers.19.ffn.layers.0.weight", + "layers.19.ffn.layers.3.bias", + "layers.19.ffn.layers.3.weight", + "layers.19.norm1.bias", + "layers.19.norm1.weight", + "layers.19.norm2.bias", + "layers.19.norm2.weight", + "layers.2.attn.in_proj_bias", + "layers.2.attn.in_proj_weight", + "layers.2.attn.out_proj.bias", + "layers.2.attn.out_proj.weight", + "layers.2.ffn.layers.0.bias", + "layers.2.ffn.layers.0.weight", + "layers.2.ffn.layers.3.bias", + "layers.2.ffn.layers.3.weight", + "layers.2.norm1.bias", + "layers.2.norm1.weight", + "layers.2.norm2.bias", + "layers.2.norm2.weight", + "layers.20.attn.in_proj_bias", + "layers.20.attn.in_proj_weight", + "layers.20.attn.out_proj.bias", + "layers.20.attn.out_proj.weight", + "layers.20.ffn.layers.0.bias", + "layers.20.ffn.layers.0.weight", + "layers.20.ffn.layers.3.bias", + "layers.20.ffn.layers.3.weight", + "layers.20.norm1.bias", + "layers.20.norm1.weight", + "layers.20.norm2.bias", + "layers.20.norm2.weight", + "layers.21.attn.in_proj_bias", + "layers.21.attn.in_proj_weight", + "layers.21.attn.out_proj.bias", + "layers.21.attn.out_proj.weight", + "layers.21.ffn.layers.0.bias", + "layers.21.ffn.layers.0.weight", + "layers.21.ffn.layers.3.bias", + "layers.21.ffn.layers.3.weight", + "layers.21.norm1.bias", + "layers.21.norm1.weight", + "layers.21.norm2.bias", + "layers.21.norm2.weight", + "layers.22.attn.in_proj_bias", + "layers.22.attn.in_proj_weight", + "layers.22.attn.out_proj.bias", + "layers.22.attn.out_proj.weight", + "layers.22.ffn.layers.0.bias", + "layers.22.ffn.layers.0.weight", + "layers.22.ffn.layers.3.bias", + "layers.22.ffn.layers.3.weight", + "layers.22.norm1.bias", + "layers.22.norm1.weight", + "layers.22.norm2.bias", + "layers.22.norm2.weight", + "layers.23.attn.in_proj_bias", + "layers.23.attn.in_proj_weight", + "layers.23.attn.out_proj.bias", + "layers.23.attn.out_proj.weight", + "layers.23.ffn.layers.0.bias", + "layers.23.ffn.layers.0.weight", + "layers.23.ffn.layers.3.bias", + "layers.23.ffn.layers.3.weight", + "layers.23.norm1.bias", + "layers.23.norm1.weight", + "layers.23.norm2.bias", + "layers.23.norm2.weight", + "layers.3.attn.in_proj_bias", + "layers.3.attn.in_proj_weight", + "layers.3.attn.out_proj.bias", + "layers.3.attn.out_proj.weight", + "layers.3.ffn.layers.0.bias", + "layers.3.ffn.layers.0.weight", + "layers.3.ffn.layers.3.bias", + "layers.3.ffn.layers.3.weight", + "layers.3.norm1.bias", + "layers.3.norm1.weight", + "layers.3.norm2.bias", + "layers.3.norm2.weight", + "layers.4.attn.in_proj_bias", + "layers.4.attn.in_proj_weight", + "layers.4.attn.out_proj.bias", + "layers.4.attn.out_proj.weight", + "layers.4.ffn.layers.0.bias", + "layers.4.ffn.layers.0.weight", + "layers.4.ffn.layers.3.bias", + "layers.4.ffn.layers.3.weight", + "layers.4.norm1.bias", + "layers.4.norm1.weight", + "layers.4.norm2.bias", + "layers.4.norm2.weight", + "layers.5.attn.in_proj_bias", + "layers.5.attn.in_proj_weight", + "layers.5.attn.out_proj.bias", + "layers.5.attn.out_proj.weight", + "layers.5.ffn.layers.0.bias", + "layers.5.ffn.layers.0.weight", + "layers.5.ffn.layers.3.bias", + "layers.5.ffn.layers.3.weight", + "layers.5.norm1.bias", + "layers.5.norm1.weight", + "layers.5.norm2.bias", + "layers.5.norm2.weight", + "layers.6.attn.in_proj_bias", + "layers.6.attn.in_proj_weight", + "layers.6.attn.out_proj.bias", + "layers.6.attn.out_proj.weight", + "layers.6.ffn.layers.0.bias", + "layers.6.ffn.layers.0.weight", + "layers.6.ffn.layers.3.bias", + "layers.6.ffn.layers.3.weight", + "layers.6.norm1.bias", + "layers.6.norm1.weight", + "layers.6.norm2.bias", + "layers.6.norm2.weight", + "layers.7.attn.in_proj_bias", + "layers.7.attn.in_proj_weight", + "layers.7.attn.out_proj.bias", + "layers.7.attn.out_proj.weight", + "layers.7.ffn.layers.0.bias", + "layers.7.ffn.layers.0.weight", + "layers.7.ffn.layers.3.bias", + "layers.7.ffn.layers.3.weight", + "layers.7.norm1.bias", + "layers.7.norm1.weight", + "layers.7.norm2.bias", + "layers.7.norm2.weight", + "layers.8.attn.in_proj_bias", + "layers.8.attn.in_proj_weight", + "layers.8.attn.out_proj.bias", + "layers.8.attn.out_proj.weight", + "layers.8.ffn.layers.0.bias", + "layers.8.ffn.layers.0.weight", + "layers.8.ffn.layers.3.bias", + "layers.8.ffn.layers.3.weight", + "layers.8.norm1.bias", + "layers.8.norm1.weight", + "layers.8.norm2.bias", + "layers.8.norm2.weight", + "layers.9.attn.in_proj_bias", + "layers.9.attn.in_proj_weight", + "layers.9.attn.out_proj.bias", + "layers.9.attn.out_proj.weight", + "layers.9.ffn.layers.0.bias", + "layers.9.ffn.layers.0.weight", + "layers.9.ffn.layers.3.bias", + "layers.9.ffn.layers.3.weight", + "layers.9.norm1.bias", + "layers.9.norm1.weight", + "layers.9.norm2.bias", + "layers.9.norm2.weight", + "porj_linear.bias", + "porj_linear.weight" + ] +} diff --git a/skysensepp-fusion-neck/model.safetensors b/skysensepp-fusion-neck/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..66911f73bd59f2cdcff5d72082b3174702a482eb --- /dev/null +++ b/skysensepp-fusion-neck/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:18605dbf0082ad045575b7d115604ab048506e785585e4fbdef7dcf180ac1cb7 +size 1220808632 diff --git a/skysensepp-fusion-neck/modeling_skysensepp_fusion_neck.py b/skysensepp-fusion-neck/modeling_skysensepp_fusion_neck.py new file mode 100644 index 0000000000000000000000000000000000000000..5387a07d8eae4aef4485c8368046deba90ec2277 --- /dev/null +++ b/skysensepp-fusion-neck/modeling_skysensepp_fusion_neck.py @@ -0,0 +1,164 @@ +"""SkySense++ fusion neck (TransformerEncoder) — optional multi-modal fusion module.""" + +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPooling + +from .configuration_skysensepp import SkySensePlusPlusFusionNeckConfig +from .modeling_utils import DropPath, FFN + + +class FusionEncoderLayer(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + feedforward_channels: int, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + qkv_bias: bool = True, + with_cp: bool = False, + ): + super().__init__() + self.with_cp = with_cp + self.norm1 = nn.LayerNorm(embed_dims) + self.attn = nn.MultiheadAttention( + embed_dim=embed_dims, + num_heads=num_heads, + dropout=attn_drop_rate, + bias=qkv_bias, + batch_first=True, + ) + self.proj_drop = nn.Dropout(drop_rate) + self.norm2 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=2, + ffn_drop=drop_rate, + drop_path=drop_path_rate, + act_layer=nn.GELU, + add_identity=True, + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + def _inner_forward(x): + residual = x + x_norm = self.norm1(x) + attn_out, _ = self.attn(x_norm, x_norm, x_norm) + attn_out = self.proj_drop(attn_out) + x = residual + self.drop_path(attn_out) + return self.ffn(self.norm2(x), identity=x) + + if self.with_cp and x.requires_grad: + return cp.checkpoint(_inner_forward, x, use_reentrant=False) + return _inner_forward(x) + + +class SkySensePlusPlusFusionNeckPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusFusionNeckConfig + base_model_prefix = "skysensepp_fusion_neck" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusFusionNeckModel(SkySensePlusPlusFusionNeckPreTrainedModel): + """Fuses per-location multi-modal tokens into a cls-token representation. + + Input shape: ``(batch, num_modalities, input_dims)`` — e.g. concatenated + HR + S2 + S1 stage-3 features with ``input_dims=2816``. + + Output: cls token embedding ``(batch, embed_dims)`` when + ``output_cls_token=True`` (default). + """ + + def __init__(self, config: SkySensePlusPlusFusionNeckConfig): + super().__init__(config) + + # Original checkpoint uses the typo `porj_linear`. + self.porj_linear = nn.Linear(config.input_dims, config.embed_dims) + self.with_cls_token = config.with_cls_token + self.output_cls_token = config.output_cls_token + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + num_layers = config.num_layers + if num_layers > 1: + dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)] + else: + dpr = [0.0] + + self.layers = nn.ModuleList() + for i in range(config.num_layers): + self.layers.append( + FusionEncoderLayer( + embed_dims=config.embed_dims, + num_heads=config.num_heads, + feedforward_channels=config.mlp_ratio * config.embed_dims, + attn_drop_rate=config.attn_drop_rate, + drop_rate=config.drop_rate, + drop_path_rate=dpr[i], + qkv_bias=config.qkv_bias, + with_cp=config.with_cp, + ) + ) + + self.post_init() + + def forward( + self, + hidden_states: torch.Tensor, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + """Forward pass. + + Args: + hidden_states: ``(batch, seq_len, input_dims)`` fused modality tokens. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x = self.porj_linear(hidden_states) + cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) + x = torch.cat((cls_tokens, x), dim=1) + if not self.with_cls_token: + x = x[:, 1:] + + all_hidden_states = () if output_hidden_states else None + for layer in self.layers: + x = layer(x) + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if self.output_cls_token: + pooler = x[:, 0] + last_hidden = pooler.unsqueeze(1) + elif self.with_cls_token: + pooler = None + last_hidden = x[:, 1:] + else: + pooler = None + last_hidden = x + + if not return_dict: + return (last_hidden, pooler) if pooler is not None else (last_hidden,) + + return BaseModelOutputWithPooling( + last_hidden_state=last_hidden, + pooler_output=pooler, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-fusion-neck/modeling_utils.py b/skysensepp-fusion-neck/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..93feae77f3a3c46e65167f2ed312a25b7cd4ad3a --- /dev/null +++ b/skysensepp-fusion-neck/modeling_utils.py @@ -0,0 +1,557 @@ +"""SkySense: Pure PyTorch + HuggingFace Transformers implementation. + +Shared utility modules used across SkySense model implementations. +""" + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def to_2tuple(x): + """Convert to a 2-tuple.""" + if isinstance(x, (list, tuple)): + return tuple(x) + return (x, x) + + +class DropPath(nn.Module): + """Drop paths (stochastic depth) per sample. + + Args: + drop_prob (float): Probability of dropping a path. Default: 0.0. + """ + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor = torch.floor(random_tensor + keep_prob) + output = x / keep_prob * random_tensor + return output + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding using Conv2d. + + Args: + in_channels (int): Number of input channels. Default: 3. + embed_dims (int): Embedding dimension. Default: 96. + kernel_size (int): Kernel size of the projection. Default: 4. + stride (int): Stride of the projection. Default: 4. + padding (int): Padding of the projection. Default: 0. + norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm. + input_size (int or tuple or None): Input resolution for calculating output size. + """ + + def __init__( + self, + in_channels: int = 3, + embed_dims: int = 96, + kernel_size: int = 4, + stride: int = 4, + padding: int = 0, + norm_layer: Optional[type] = nn.LayerNorm, + input_size: Optional[int] = None, + ): + super().__init__() + self.projection = nn.Conv2d( + in_channels, embed_dims, + kernel_size=kernel_size, stride=stride, padding=padding, + ) + self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity() + + # Compute init output size if input_size is given + if input_size is not None: + input_size = to_2tuple(input_size) + self.init_out_size = ( + (input_size[0] - kernel_size + 2 * padding) // stride + 1, + (input_size[1] - kernel_size + 2 * padding) // stride + 1, + ) + else: + self.init_out_size = None + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]: + x = self.projection(x) # (B, C, H, W) + out_size = (x.shape[2], x.shape[3]) + x = x.flatten(2).transpose(1, 2) # (B, H*W, C) + x = self.norm(x) + return x, out_size + + +class FFN(nn.Module): + """Feed-Forward Network. + + Args: + embed_dims (int): Input dimension. + feedforward_channels (int): Hidden dimension. + num_fcs (int): Number of FC layers. Default: 2. + ffn_drop (float): Dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + act_layer (nn.Module): Activation layer class. Default: nn.GELU. + add_identity (bool): Whether to add identity connection. Default: True. + """ + + def __init__( + self, + embed_dims: int, + feedforward_channels: int, + num_fcs: int = 2, + ffn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: type = nn.GELU, + add_identity: bool = True, + ): + super().__init__() + assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}" + self.embed_dims = embed_dims + self.feedforward_channels = feedforward_channels + self.add_identity = add_identity + + layers = [] + in_channels = embed_dims + for i in range(num_fcs - 1): + layers.append(nn.Linear(in_channels, feedforward_channels)) + layers.append(act_layer()) + layers.append(nn.Dropout(ffn_drop)) + in_channels = feedforward_channels + layers.append(nn.Linear(feedforward_channels, embed_dims)) + layers.append(nn.Dropout(ffn_drop)) + self.layers = nn.Sequential(*layers) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor: + out = self.layers(x) + out = self.drop_path(out) + if self.add_identity: + if identity is None: + identity = x + out = out + identity + return out + + +class WindowMSAV2(nn.Module): + """Window-based Multi-head Self-Attention for Swin Transformer V2. + + Uses cosine attention and log-spaced continuous position bias (log-CPB). + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): Window size (Wh, Ww). + pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0). + qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Output projection dropout rate. Default: 0.0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: Tuple[int, int], + pretrained_window_size: Tuple[int, int] = (0, 0), + qkv_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ): + super().__init__() + self.embed_dims = embed_dims + self.num_heads = num_heads + self.window_size = window_size + self.pretrained_window_size = pretrained_window_size + + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1)))) + + # MLP for continuous relative position bias (log-CPB) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) + + # Build relative coords table + self._build_relative_coords_table() + # Build relative position index + self._build_relative_position_index() + + self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False) + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(embed_dims)) + self.v_bias = nn.Parameter(torch.zeros(embed_dims)) + else: + self.q_bias = None + self.v_bias = None + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(embed_dims, embed_dims) + self.proj_drop = nn.Dropout(proj_drop) + self.softmax = nn.Softmax(dim=-1) + + def _build_relative_coords_table(self): + """Build the relative coordinates table for log-CPB.""" + Wh, Ww = self.window_size + # Table of relative coordinates + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2) + + # Normalize to [-1, 1] and apply log-scale + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 # normalize to -8, 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + self.register_buffer("relative_coords_table", coords_table) + + def _build_relative_position_index(self): + """Build the pairwise relative position index for each window token.""" + Wh, Ww = self.window_size + coords_h = torch.arange(Wh) + coords_w = torch.arange(Ww) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij')) + coords_flatten = coords.view(2, -1) + + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # (2, Wh*Ww, Wh*Ww) + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += Wh - 1 + relative_coords[:, :, 1] += Ww - 1 + relative_coords[:, :, 0] *= 2 * Ww - 1 + relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww) + self.register_buffer("relative_position_index", relative_position_index) + + def _compute_position_bias(self, N): + """Compute relative position bias, supporting dynamic window sizes. + + The log-CPB (Continuous Position Bias) MLP can generalize to any window + size by computing bias from normalized relative coordinates. + """ + init_N = self.window_size[0] * self.window_size[1] + if N == init_N: + # Use pre-built tables + relative_position_bias_table = self.cpb_mlp( + self.relative_coords_table + ).view(-1, self.num_heads) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view(N, N, -1) + else: + # Dynamic: compute for actual window size on-the-fly + Wh = Ww = int(math.sqrt(N)) + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + # Build position index for actual window size + ch = torch.arange(Wh, device=self.logit_scale.device) + cw = torch.arange(Ww, device=self.logit_scale.device) + coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij')) + coords_flat = coords.view(2, -1) + rel = coords_flat[:, :, None] - coords_flat[:, None, :] + rel = rel.permute(1, 2, 0).contiguous() + rel[:, :, 0] += Wh - 1 + rel[:, :, 1] += Ww - 1 + rel[:, :, 0] *= 2 * Ww - 1 + pos_index = rel.sum(-1) + + bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads) + relative_position_bias = bias_table[ + pos_index.view(-1) + ].view(N, N, -1) + + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + return relative_position_bias + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: (num_windows*B, N, C) where N = Wh*Ww + mask: (num_windows, N, N) or None + """ + B_, N, C = x.shape + + # Compute QKV with bias + if self.q_bias is not None: + qkv_bias = torch.cat( + (self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias)) + qkv = F.linear(x, self.qkv.weight, qkv_bias) + else: + qkv = self.qkv(x) + + qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + + # Cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=math.log(1.0 / 0.01) + ).exp() + attn = attn * logit_scale + + # Log-CPB relative position bias (supports dynamic window sizes) + relative_position_bias = self._compute_position_bias(N) + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + attn = attn + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + + attn = self.softmax(attn) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class ShiftWindowMSA(nn.Module): + """Shifted Window Multi-head Self-Attention. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. Default: 0. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Projection dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + pad_small_map (bool): Pad small feature maps to window size. Default: False. + pretrained_window_size (int): Pretrained window size. Default: 0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int, + shift_size: int = 0, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.window_size = window_size + self.shift_size = shift_size + self.pad_small_map = pad_small_map + + self.w_msa = WindowMSAV2( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=to_2tuple(window_size), + pretrained_window_size=to_2tuple(pretrained_window_size), + attn_drop=attn_drop, + proj_drop=proj_drop, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W, f"Input length {L} != H*W ({H}*{W})" + + x = x.view(B, H, W, C) + + window_size = self.window_size + shift_size = self.shift_size + + # Pad or shrink window + if self.pad_small_map: + pad_r = (window_size - W % window_size) % window_size + pad_b = (window_size - H % window_size) % window_size + x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b)) + _, Hp, Wp, _ = x.shape + else: + Hp, Wp = H, W + if window_size > Hp: + window_size = Hp + shift_size = 0 + if window_size > Wp: + window_size = Wp + shift_size = 0 + + # Compute attention mask for SW-MSA + attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device) + + # Cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) + + # Partition windows + x_windows = self._window_partition(x, window_size) + # (num_windows*B, window_size*window_size, C) + + # W-MSA/SW-MSA + attn_windows = self.w_msa(x_windows, mask=attn_mask) + + # Merge windows + x = self._window_reverse(attn_windows, window_size, Hp, Wp) + + # Reverse cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) + + if self.pad_small_map and (pad_r > 0 or pad_b > 0): + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + x = self.drop_path(x) + return x + + @staticmethod + def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor: + """Partition into non-overlapping windows.""" + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous() + windows = windows.view(-1, window_size * window_size, C) + return windows + + @staticmethod + def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor: + """Reverse window partition.""" + B_nW = windows.shape[0] + nH = H // window_size + nW = W // window_size + B = B_nW // (nH * nW) + x = windows.view(B, nH, nW, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous() + x = x.view(B, H, W, -1) + return x + + @staticmethod + def _compute_attn_mask(H, W, window_size, shift_size, device): + """Compute attention mask for shifted window attention.""" + if shift_size <= 0: + return None + img_mask = torch.zeros((1, H, W, 1), device=device) + h_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + w_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + # Partition mask + mask_windows = img_mask.view( + 1, H // window_size, window_size, W // window_size, window_size, 1 + ) + mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous() + mask_windows = mask_windows.view(-1, window_size * window_size) + + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0) + attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0) + return attn_mask + + +class PatchMerging(nn.Module): + """Patch Merging Layer for downsampling (2x). + + Args: + in_channels (int): Input channels. + out_channels (int): Output channels. + norm_layer (type): Normalization layer. Default: nn.LayerNorm. + is_post_norm (bool): Apply norm after linear. Default: True. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + norm_layer: type = nn.LayerNorm, + is_post_norm: bool = True, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.is_post_norm = is_post_norm + self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False) + if is_post_norm: + self.norm = norm_layer(out_channels) + else: + self.norm = norm_layer(4 * in_channels) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W + + x = x.view(B, H, W, C) + + # Pad if needed + pad_h = H % 2 + pad_w = W % 2 + if pad_h or pad_w: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], dim=-1) + + out_h = (H + pad_h) // 2 + out_w = (W + pad_w) // 2 + x = x.view(B, out_h * out_w, 4 * C) + + if self.is_post_norm: + x = self.reduction(x) + x = self.norm(x) + else: + x = self.norm(x) + x = self.reduction(x) + + return x, (out_h, out_w) diff --git a/skysensepp-fusion-neck/pipeline_skysensepp.py b/skysensepp-fusion-neck/pipeline_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9a5354361c9a08abd1a82b3df8d4ec678a209b --- /dev/null +++ b/skysensepp-fusion-neck/pipeline_skysensepp.py @@ -0,0 +1,86 @@ +"""Custom HuggingFace pipeline for SkySense++ MSL feature extraction.""" + +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline): + """Pipeline for SkySense++ MSL backbones. + + Expects image tensors plus semantic annotation maps (class indices). + """ + + def _sanitize_parameters( + self, + annotation=None, + mask=None, + output_hidden_states=None, + **kwargs, + ): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + + if annotation is not None: + preprocess_params["annotation"] = annotation + if mask is not None: + forward_params["mask"] = mask + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + + return preprocess_params, forward_params, postprocess_params + + def preprocess( + self, + pixel_values: Any, + annotation: Optional[Any] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if isinstance(pixel_values, dict): + annotation = pixel_values.get("annotation", annotation) + pixel_values = pixel_values.get("pixel_values", pixel_values) + + if isinstance(pixel_values, np.ndarray): + pixel_values = torch.from_numpy(pixel_values).float() + elif not isinstance(pixel_values, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}" + ) + + if annotation is None: + raise ValueError("SkySense++ MSL models require an `annotation` semantic map.") + + if isinstance(annotation, np.ndarray): + annotation = torch.from_numpy(annotation).long() + elif not isinstance(annotation, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for annotation, got {type(annotation)}" + ) + + if pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + if annotation.ndim == 2: + annotation = annotation.unsqueeze(0) + + return {"pixel_values": pixel_values, "annotation": annotation} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + pixel_values=model_inputs["pixel_values"], + annotation=model_inputs["annotation"], + mask=kwargs.get("mask"), + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"last_hidden_state": outputs.last_hidden_state} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-fusion-neck/pipeline_skysensepp_fusion.py b/skysensepp-fusion-neck/pipeline_skysensepp_fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..80f070fa9fd3f18a053caa1afef9dc142bce2597 --- /dev/null +++ b/skysensepp-fusion-neck/pipeline_skysensepp_fusion.py @@ -0,0 +1,53 @@ +"""Optional pipeline for SkySense++ fusion neck.""" + +from typing import Any, Dict + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusFusionNeckPipeline(Pipeline): + """Pipeline for the optional SkySense++ fusion neck module. + + Expects concatenated multi-modal tokens per spatial location: + ``(batch, num_modalities, input_dims)``. + """ + + def _sanitize_parameters(self, output_hidden_states=None, **kwargs): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + return preprocess_params, forward_params, postprocess_params + + def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]: + if isinstance(hidden_states, dict): + hidden_states = hidden_states["hidden_states"] + + if isinstance(hidden_states, np.ndarray): + hidden_states = torch.from_numpy(hidden_states).float() + elif not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}" + ) + if hidden_states.ndim == 2: + hidden_states = hidden_states.unsqueeze(0) + return {"hidden_states": hidden_states} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + hidden_states=model_inputs["hidden_states"], + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"pooler_output": outputs.pooler_output} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-swinv2-msl-hr/__init__.py b/skysensepp-swinv2-msl-hr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d49c51cc97fc8d74fa2f98f3c4677a5b59f8bc2 --- /dev/null +++ b/skysensepp-swinv2-msl-hr/__init__.py @@ -0,0 +1,25 @@ +"""SkySense++: Multi-Modal Remote Sensing Foundation Model (HuggingFace).""" + +from .configuration_skysensepp import ( + SkySensePlusPlusSwinV2MSLConfig, + SkySensePlusPlusViTMSLConfig, +) +from .modeling_skysensepp_swinv2_msl import ( + SkySensePlusPlusSwinV2MSLModel, + SkySensePlusPlusSwinV2MSLPreTrainedModel, +) +from .modeling_skysensepp_vit_msl import ( + SkySensePlusPlusViTMSLModel, + SkySensePlusPlusViTMSLPreTrainedModel, +) +from .pipeline_skysensepp import SkySensePlusPlusMSLFeatureExtractionPipeline + +__all__ = [ + "SkySensePlusPlusSwinV2MSLConfig", + "SkySensePlusPlusViTMSLConfig", + "SkySensePlusPlusSwinV2MSLModel", + "SkySensePlusPlusSwinV2MSLPreTrainedModel", + "SkySensePlusPlusViTMSLModel", + "SkySensePlusPlusViTMSLPreTrainedModel", + "SkySensePlusPlusMSLFeatureExtractionPipeline", +] diff --git a/skysensepp-swinv2-msl-hr/config.json b/skysensepp-swinv2-msl-hr/config.json new file mode 100644 index 0000000000000000000000000000000000000000..8341675df5f2c19a264070a9e68ccb0e30c1a71c --- /dev/null +++ b/skysensepp-swinv2-msl-hr/config.json @@ -0,0 +1,82 @@ +{ + "return_dict": true, + "output_hidden_states": false, + "dtype": "float32", + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": [ + "SkySensePlusPlusSwinV2MSLModel" + ], + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "transformers_version": "5.0.0", + "arch": "huge", + "embed_dims": 352, + "depths": [ + 2, + 2, + 18, + 2 + ], + "num_heads": [ + 8, + 16, + 32, + 64 + ], + "extra_norm_every_n_blocks": 6, + "img_size": 512, + "patch_size": 4, + "in_channels": 3, + "window_size": 8, + "drop_rate": 0.0, + "drop_path_rate": 0.2, + "out_indices": [ + 0, + 1, + 2, + 3 + ], + "use_abs_pos_embed": false, + "with_cp": false, + "pad_small_map": false, + "pretrained_window_sizes": [ + 0, + 0, + 0, + 0 + ], + "is_post_norm_downsample": true, + "vocabulary_size": 64, + "num_vocabulary_tokens": 65, + "merge_stage": 2, + "use_attn": true, + "model_type": "skysensepp_swinv2_msl", + "output_attentions": false, + "auto_map": { + "AutoConfig": "configuration_skysensepp.SkySensePlusPlusSwinV2MSLConfig", + "AutoModel": "modeling_skysensepp_swinv2_msl.SkySensePlusPlusSwinV2MSLModel" + }, + "custom_pipelines": { + "skysensepp-feature-extraction": { + "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + }, + "image-feature-extraction": { + "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + } +} \ No newline at end of file diff --git a/skysensepp-swinv2-msl-hr/configuration_skysensepp.py b/skysensepp-swinv2-msl-hr/configuration_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..e086dd1de7d378e81cc124a89a36e33f45bc4b3f --- /dev/null +++ b/skysensepp-swinv2-msl-hr/configuration_skysensepp.py @@ -0,0 +1,124 @@ +"""Configuration classes for SkySense++ MSL backbones.""" + +from transformers import PretrainedConfig + + +class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig): + """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical).""" + + model_type = "skysensepp_swinv2_msl" + + arch_zoo = { + "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0}, + "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0}, + "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6}, + "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6}, + } + + def __init__( + self, + arch="huge", + img_size=224, + patch_size=4, + in_channels=3, + window_size=8, + drop_rate=0.0, + drop_path_rate=0.2, + out_indices=(0, 1, 2, 3), + use_abs_pos_embed=False, + with_cp=False, + pad_small_map=False, + pretrained_window_sizes=(0, 0, 0, 0), + is_post_norm_downsample=True, + vocabulary_size=64, + merge_stage=2, + use_attn=True, + **kwargs, + ): + super().__init__(**kwargs) + + arch = arch.lower() + if arch not in self.arch_zoo: + raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}") + arch_settings = self.arch_zoo[arch] + + self.arch = arch + self.embed_dims = arch_settings["embed_dims"] + self.depths = arch_settings["depths"] + self.num_heads = arch_settings["num_heads"] + self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"] + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.window_size = window_size + self.drop_rate = drop_rate + self.drop_path_rate = drop_path_rate + self.out_indices = list(out_indices) + self.use_abs_pos_embed = use_abs_pos_embed + self.with_cp = with_cp + self.pad_small_map = pad_small_map + self.pretrained_window_sizes = list(pretrained_window_sizes) + self.is_post_norm_downsample = is_post_norm_downsample + + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + + +class SkySensePlusPlusViTMSLConfig(PretrainedConfig): + """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1).""" + + model_type = "skysensepp_vit_msl" + + def __init__( + self, + img_size=16, + patch_size=4, + in_channels=10, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + out_indices=(5, 11, 17, 23), + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=False, + output_cls_token=False, + patch_norm=False, + final_norm=False, + with_cp=False, + vocabulary_size=64, + merge_stage=4, + use_attn=False, + modality="s2", + **kwargs, + ): + super().__init__(**kwargs) + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.out_indices = list(out_indices) + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.patch_norm = patch_norm + self.final_norm = final_norm + self.with_cp = with_cp + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + self.modality = modality diff --git a/skysensepp-swinv2-msl-hr/conversion_manifest.json b/skysensepp-swinv2-msl-hr/conversion_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..fb7878619332f960f2928a5940306b269e967d2f --- /dev/null +++ b/skysensepp-swinv2-msl-hr/conversion_manifest.json @@ -0,0 +1,523 @@ +{ + "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release_hr.pth", + "modality": "hr", + "model_class": "SkySensePlusPlusSwinV2MSLModel", + "num_tensors": 464, + "missing_keys": [ + "stages.0.blocks.0.attn.w_msa.relative_coords_table", + "stages.0.blocks.0.attn.w_msa.relative_position_index", + "stages.0.blocks.1.attn.w_msa.relative_coords_table", + "stages.0.blocks.1.attn.w_msa.relative_position_index", + "stages.1.blocks.0.attn.w_msa.relative_coords_table", + "stages.1.blocks.0.attn.w_msa.relative_position_index", + "stages.1.blocks.1.attn.w_msa.relative_coords_table", + "stages.1.blocks.1.attn.w_msa.relative_position_index", + "stages.2.blocks.0.attn.w_msa.relative_coords_table", + "stages.2.blocks.0.attn.w_msa.relative_position_index", + "stages.2.blocks.1.attn.w_msa.relative_coords_table", + "stages.2.blocks.1.attn.w_msa.relative_position_index", + "stages.2.blocks.2.attn.w_msa.relative_coords_table", + "stages.2.blocks.2.attn.w_msa.relative_position_index", + "stages.2.blocks.3.attn.w_msa.relative_coords_table", + "stages.2.blocks.3.attn.w_msa.relative_position_index", + "stages.2.blocks.4.attn.w_msa.relative_coords_table", + "stages.2.blocks.4.attn.w_msa.relative_position_index", + "stages.2.blocks.5.attn.w_msa.relative_coords_table", + "stages.2.blocks.5.attn.w_msa.relative_position_index", + "stages.2.blocks.6.attn.w_msa.relative_coords_table", + "stages.2.blocks.6.attn.w_msa.relative_position_index", + "stages.2.blocks.7.attn.w_msa.relative_coords_table", + "stages.2.blocks.7.attn.w_msa.relative_position_index", + "stages.2.blocks.8.attn.w_msa.relative_coords_table", + "stages.2.blocks.8.attn.w_msa.relative_position_index", + "stages.2.blocks.9.attn.w_msa.relative_coords_table", + "stages.2.blocks.9.attn.w_msa.relative_position_index", + "stages.2.blocks.10.attn.w_msa.relative_coords_table", + "stages.2.blocks.10.attn.w_msa.relative_position_index", + "stages.2.blocks.11.attn.w_msa.relative_coords_table", + "stages.2.blocks.11.attn.w_msa.relative_position_index", + "stages.2.blocks.12.attn.w_msa.relative_coords_table", + "stages.2.blocks.12.attn.w_msa.relative_position_index", + "stages.2.blocks.13.attn.w_msa.relative_coords_table", + "stages.2.blocks.13.attn.w_msa.relative_position_index", + "stages.2.blocks.14.attn.w_msa.relative_coords_table", + "stages.2.blocks.14.attn.w_msa.relative_position_index", + "stages.2.blocks.15.attn.w_msa.relative_coords_table", + "stages.2.blocks.15.attn.w_msa.relative_position_index", + "stages.2.blocks.16.attn.w_msa.relative_coords_table", + "stages.2.blocks.16.attn.w_msa.relative_position_index", + "stages.2.blocks.17.attn.w_msa.relative_coords_table", + "stages.2.blocks.17.attn.w_msa.relative_position_index", + "stages.3.blocks.0.attn.w_msa.relative_coords_table", + "stages.3.blocks.0.attn.w_msa.relative_position_index", + "stages.3.blocks.1.attn.w_msa.relative_coords_table", + "stages.3.blocks.1.attn.w_msa.relative_position_index" + ], + "unexpected_keys": [], + "tensor_names": [ + "attn1.attn.in_proj_bias", + "attn1.attn.in_proj_weight", + "attn1.attn.out_proj.bias", + "attn1.attn.out_proj.weight", + "attn1.proj_in.bias", + "attn1.proj_in.weight", + "attn1.proj_out.bias", + "attn1.proj_out.weight", + "attn2.attn.in_proj_bias", + "attn2.attn.in_proj_weight", + "attn2.attn.out_proj.bias", + "attn2.attn.out_proj.weight", + "attn2.proj_in.bias", + "attn2.proj_in.weight", + "attn2.proj_out.bias", + "attn2.proj_out.weight", + "attn3.attn.in_proj_bias", + "attn3.attn.in_proj_weight", + "attn3.attn.out_proj.bias", + "attn3.attn.out_proj.weight", + "attn3.proj_in.bias", + "attn3.proj_in.weight", + "attn3.proj_out.bias", + "attn3.proj_out.weight", + "mask_token", + "norm0.bias", + "norm0.weight", + "norm1.bias", + "norm1.weight", + "norm2.bias", + "norm2.weight", + "norm3.bias", + "norm3.weight", + "norm_attn.bias", + "norm_attn.weight", + "patch_embed.norm.bias", + "patch_embed.norm.weight", + "patch_embed.projection.bias", + "patch_embed.projection.weight", + "stages.0.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "stages.0.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "stages.0.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "stages.0.blocks.0.attn.w_msa.logit_scale", + "stages.0.blocks.0.attn.w_msa.proj.bias", + "stages.0.blocks.0.attn.w_msa.proj.weight", + "stages.0.blocks.0.attn.w_msa.q_bias", + "stages.0.blocks.0.attn.w_msa.qkv.weight", + "stages.0.blocks.0.attn.w_msa.v_bias", + "stages.0.blocks.0.ffn.layers.0.bias", + "stages.0.blocks.0.ffn.layers.0.weight", + "stages.0.blocks.0.ffn.layers.3.bias", + "stages.0.blocks.0.ffn.layers.3.weight", + "stages.0.blocks.0.norm1.bias", + "stages.0.blocks.0.norm1.weight", + "stages.0.blocks.0.norm2.bias", + "stages.0.blocks.0.norm2.weight", + "stages.0.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "stages.0.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "stages.0.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "stages.0.blocks.1.attn.w_msa.logit_scale", + "stages.0.blocks.1.attn.w_msa.proj.bias", + "stages.0.blocks.1.attn.w_msa.proj.weight", + "stages.0.blocks.1.attn.w_msa.q_bias", + "stages.0.blocks.1.attn.w_msa.qkv.weight", + "stages.0.blocks.1.attn.w_msa.v_bias", + "stages.0.blocks.1.ffn.layers.0.bias", + "stages.0.blocks.1.ffn.layers.0.weight", + "stages.0.blocks.1.ffn.layers.3.bias", + "stages.0.blocks.1.ffn.layers.3.weight", + "stages.0.blocks.1.norm1.bias", + "stages.0.blocks.1.norm1.weight", + "stages.0.blocks.1.norm2.bias", + "stages.0.blocks.1.norm2.weight", + "stages.1.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "stages.1.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "stages.1.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "stages.1.blocks.0.attn.w_msa.logit_scale", + "stages.1.blocks.0.attn.w_msa.proj.bias", + "stages.1.blocks.0.attn.w_msa.proj.weight", + "stages.1.blocks.0.attn.w_msa.q_bias", + "stages.1.blocks.0.attn.w_msa.qkv.weight", + "stages.1.blocks.0.attn.w_msa.v_bias", + "stages.1.blocks.0.ffn.layers.0.bias", + "stages.1.blocks.0.ffn.layers.0.weight", + "stages.1.blocks.0.ffn.layers.3.bias", + "stages.1.blocks.0.ffn.layers.3.weight", + "stages.1.blocks.0.norm1.bias", + "stages.1.blocks.0.norm1.weight", + "stages.1.blocks.0.norm2.bias", + "stages.1.blocks.0.norm2.weight", + "stages.1.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "stages.1.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "stages.1.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "stages.1.blocks.1.attn.w_msa.logit_scale", + "stages.1.blocks.1.attn.w_msa.proj.bias", + "stages.1.blocks.1.attn.w_msa.proj.weight", + "stages.1.blocks.1.attn.w_msa.q_bias", + "stages.1.blocks.1.attn.w_msa.qkv.weight", + "stages.1.blocks.1.attn.w_msa.v_bias", + "stages.1.blocks.1.ffn.layers.0.bias", + "stages.1.blocks.1.ffn.layers.0.weight", + "stages.1.blocks.1.ffn.layers.3.bias", + "stages.1.blocks.1.ffn.layers.3.weight", + "stages.1.blocks.1.norm1.bias", + "stages.1.blocks.1.norm1.weight", + "stages.1.blocks.1.norm2.bias", + "stages.1.blocks.1.norm2.weight", + "stages.1.downsample.norm.bias", + "stages.1.downsample.norm.weight", + "stages.1.downsample.reduction.weight", + "stages.2.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.0.attn.w_msa.logit_scale", + "stages.2.blocks.0.attn.w_msa.proj.bias", + "stages.2.blocks.0.attn.w_msa.proj.weight", + "stages.2.blocks.0.attn.w_msa.q_bias", + "stages.2.blocks.0.attn.w_msa.qkv.weight", + "stages.2.blocks.0.attn.w_msa.v_bias", + "stages.2.blocks.0.ffn.layers.0.bias", + "stages.2.blocks.0.ffn.layers.0.weight", + "stages.2.blocks.0.ffn.layers.3.bias", + "stages.2.blocks.0.ffn.layers.3.weight", + "stages.2.blocks.0.norm1.bias", + "stages.2.blocks.0.norm1.weight", + "stages.2.blocks.0.norm2.bias", + "stages.2.blocks.0.norm2.weight", + "stages.2.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.1.attn.w_msa.logit_scale", + "stages.2.blocks.1.attn.w_msa.proj.bias", + "stages.2.blocks.1.attn.w_msa.proj.weight", + "stages.2.blocks.1.attn.w_msa.q_bias", + "stages.2.blocks.1.attn.w_msa.qkv.weight", + "stages.2.blocks.1.attn.w_msa.v_bias", + "stages.2.blocks.1.ffn.layers.0.bias", + "stages.2.blocks.1.ffn.layers.0.weight", + "stages.2.blocks.1.ffn.layers.3.bias", + "stages.2.blocks.1.ffn.layers.3.weight", + "stages.2.blocks.1.norm1.bias", + "stages.2.blocks.1.norm1.weight", + "stages.2.blocks.1.norm2.bias", + "stages.2.blocks.1.norm2.weight", + "stages.2.blocks.10.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.10.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.10.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.10.attn.w_msa.logit_scale", + "stages.2.blocks.10.attn.w_msa.proj.bias", + "stages.2.blocks.10.attn.w_msa.proj.weight", + "stages.2.blocks.10.attn.w_msa.q_bias", + "stages.2.blocks.10.attn.w_msa.qkv.weight", + "stages.2.blocks.10.attn.w_msa.v_bias", + "stages.2.blocks.10.ffn.layers.0.bias", + "stages.2.blocks.10.ffn.layers.0.weight", + "stages.2.blocks.10.ffn.layers.3.bias", + "stages.2.blocks.10.ffn.layers.3.weight", + "stages.2.blocks.10.norm1.bias", + "stages.2.blocks.10.norm1.weight", + "stages.2.blocks.10.norm2.bias", + "stages.2.blocks.10.norm2.weight", + "stages.2.blocks.11.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.11.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.11.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.11.attn.w_msa.logit_scale", + "stages.2.blocks.11.attn.w_msa.proj.bias", + "stages.2.blocks.11.attn.w_msa.proj.weight", + "stages.2.blocks.11.attn.w_msa.q_bias", + "stages.2.blocks.11.attn.w_msa.qkv.weight", + "stages.2.blocks.11.attn.w_msa.v_bias", + "stages.2.blocks.11.ffn.layers.0.bias", + "stages.2.blocks.11.ffn.layers.0.weight", + "stages.2.blocks.11.ffn.layers.3.bias", + "stages.2.blocks.11.ffn.layers.3.weight", + "stages.2.blocks.11.norm1.bias", + "stages.2.blocks.11.norm1.weight", + "stages.2.blocks.11.norm2.bias", + "stages.2.blocks.11.norm2.weight", + "stages.2.blocks.11.norm3.bias", + "stages.2.blocks.11.norm3.weight", + "stages.2.blocks.12.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.12.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.12.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.12.attn.w_msa.logit_scale", + "stages.2.blocks.12.attn.w_msa.proj.bias", + "stages.2.blocks.12.attn.w_msa.proj.weight", + "stages.2.blocks.12.attn.w_msa.q_bias", + "stages.2.blocks.12.attn.w_msa.qkv.weight", + "stages.2.blocks.12.attn.w_msa.v_bias", + "stages.2.blocks.12.ffn.layers.0.bias", + "stages.2.blocks.12.ffn.layers.0.weight", + "stages.2.blocks.12.ffn.layers.3.bias", + "stages.2.blocks.12.ffn.layers.3.weight", + "stages.2.blocks.12.norm1.bias", + "stages.2.blocks.12.norm1.weight", + "stages.2.blocks.12.norm2.bias", + "stages.2.blocks.12.norm2.weight", + "stages.2.blocks.13.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.13.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.13.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.13.attn.w_msa.logit_scale", + "stages.2.blocks.13.attn.w_msa.proj.bias", + "stages.2.blocks.13.attn.w_msa.proj.weight", + "stages.2.blocks.13.attn.w_msa.q_bias", + "stages.2.blocks.13.attn.w_msa.qkv.weight", + "stages.2.blocks.13.attn.w_msa.v_bias", + "stages.2.blocks.13.ffn.layers.0.bias", + "stages.2.blocks.13.ffn.layers.0.weight", + "stages.2.blocks.13.ffn.layers.3.bias", + "stages.2.blocks.13.ffn.layers.3.weight", + "stages.2.blocks.13.norm1.bias", + "stages.2.blocks.13.norm1.weight", + "stages.2.blocks.13.norm2.bias", + "stages.2.blocks.13.norm2.weight", + "stages.2.blocks.14.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.14.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.14.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.14.attn.w_msa.logit_scale", + "stages.2.blocks.14.attn.w_msa.proj.bias", + "stages.2.blocks.14.attn.w_msa.proj.weight", + "stages.2.blocks.14.attn.w_msa.q_bias", + "stages.2.blocks.14.attn.w_msa.qkv.weight", + "stages.2.blocks.14.attn.w_msa.v_bias", + "stages.2.blocks.14.ffn.layers.0.bias", + "stages.2.blocks.14.ffn.layers.0.weight", + "stages.2.blocks.14.ffn.layers.3.bias", + "stages.2.blocks.14.ffn.layers.3.weight", + "stages.2.blocks.14.norm1.bias", + "stages.2.blocks.14.norm1.weight", + "stages.2.blocks.14.norm2.bias", + "stages.2.blocks.14.norm2.weight", + "stages.2.blocks.15.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.15.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.15.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.15.attn.w_msa.logit_scale", + "stages.2.blocks.15.attn.w_msa.proj.bias", + "stages.2.blocks.15.attn.w_msa.proj.weight", + "stages.2.blocks.15.attn.w_msa.q_bias", + "stages.2.blocks.15.attn.w_msa.qkv.weight", + "stages.2.blocks.15.attn.w_msa.v_bias", + "stages.2.blocks.15.ffn.layers.0.bias", + "stages.2.blocks.15.ffn.layers.0.weight", + "stages.2.blocks.15.ffn.layers.3.bias", + "stages.2.blocks.15.ffn.layers.3.weight", + "stages.2.blocks.15.norm1.bias", + "stages.2.blocks.15.norm1.weight", + "stages.2.blocks.15.norm2.bias", + "stages.2.blocks.15.norm2.weight", + "stages.2.blocks.16.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.16.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.16.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.16.attn.w_msa.logit_scale", + "stages.2.blocks.16.attn.w_msa.proj.bias", + "stages.2.blocks.16.attn.w_msa.proj.weight", + "stages.2.blocks.16.attn.w_msa.q_bias", + "stages.2.blocks.16.attn.w_msa.qkv.weight", + "stages.2.blocks.16.attn.w_msa.v_bias", + "stages.2.blocks.16.ffn.layers.0.bias", + "stages.2.blocks.16.ffn.layers.0.weight", + "stages.2.blocks.16.ffn.layers.3.bias", + "stages.2.blocks.16.ffn.layers.3.weight", + "stages.2.blocks.16.norm1.bias", + "stages.2.blocks.16.norm1.weight", + "stages.2.blocks.16.norm2.bias", + "stages.2.blocks.16.norm2.weight", + "stages.2.blocks.17.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.17.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.17.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.17.attn.w_msa.logit_scale", + "stages.2.blocks.17.attn.w_msa.proj.bias", + "stages.2.blocks.17.attn.w_msa.proj.weight", + "stages.2.blocks.17.attn.w_msa.q_bias", + "stages.2.blocks.17.attn.w_msa.qkv.weight", + "stages.2.blocks.17.attn.w_msa.v_bias", + "stages.2.blocks.17.ffn.layers.0.bias", + "stages.2.blocks.17.ffn.layers.0.weight", + "stages.2.blocks.17.ffn.layers.3.bias", + "stages.2.blocks.17.ffn.layers.3.weight", + "stages.2.blocks.17.norm1.bias", + "stages.2.blocks.17.norm1.weight", + "stages.2.blocks.17.norm2.bias", + "stages.2.blocks.17.norm2.weight", + "stages.2.blocks.17.norm3.bias", + "stages.2.blocks.17.norm3.weight", + "stages.2.blocks.2.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.2.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.2.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.2.attn.w_msa.logit_scale", + "stages.2.blocks.2.attn.w_msa.proj.bias", + "stages.2.blocks.2.attn.w_msa.proj.weight", + "stages.2.blocks.2.attn.w_msa.q_bias", + "stages.2.blocks.2.attn.w_msa.qkv.weight", + "stages.2.blocks.2.attn.w_msa.v_bias", + "stages.2.blocks.2.ffn.layers.0.bias", + "stages.2.blocks.2.ffn.layers.0.weight", + "stages.2.blocks.2.ffn.layers.3.bias", + "stages.2.blocks.2.ffn.layers.3.weight", + "stages.2.blocks.2.norm1.bias", + "stages.2.blocks.2.norm1.weight", + "stages.2.blocks.2.norm2.bias", + "stages.2.blocks.2.norm2.weight", + "stages.2.blocks.3.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.3.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.3.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.3.attn.w_msa.logit_scale", + "stages.2.blocks.3.attn.w_msa.proj.bias", + "stages.2.blocks.3.attn.w_msa.proj.weight", + "stages.2.blocks.3.attn.w_msa.q_bias", + "stages.2.blocks.3.attn.w_msa.qkv.weight", + "stages.2.blocks.3.attn.w_msa.v_bias", + "stages.2.blocks.3.ffn.layers.0.bias", + "stages.2.blocks.3.ffn.layers.0.weight", + "stages.2.blocks.3.ffn.layers.3.bias", + "stages.2.blocks.3.ffn.layers.3.weight", + "stages.2.blocks.3.norm1.bias", + "stages.2.blocks.3.norm1.weight", + "stages.2.blocks.3.norm2.bias", + "stages.2.blocks.3.norm2.weight", + "stages.2.blocks.4.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.4.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.4.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.4.attn.w_msa.logit_scale", + "stages.2.blocks.4.attn.w_msa.proj.bias", + "stages.2.blocks.4.attn.w_msa.proj.weight", + "stages.2.blocks.4.attn.w_msa.q_bias", + "stages.2.blocks.4.attn.w_msa.qkv.weight", + "stages.2.blocks.4.attn.w_msa.v_bias", + "stages.2.blocks.4.ffn.layers.0.bias", + "stages.2.blocks.4.ffn.layers.0.weight", + "stages.2.blocks.4.ffn.layers.3.bias", + "stages.2.blocks.4.ffn.layers.3.weight", + "stages.2.blocks.4.norm1.bias", + "stages.2.blocks.4.norm1.weight", + "stages.2.blocks.4.norm2.bias", + "stages.2.blocks.4.norm2.weight", + "stages.2.blocks.5.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.5.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.5.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.5.attn.w_msa.logit_scale", + "stages.2.blocks.5.attn.w_msa.proj.bias", + "stages.2.blocks.5.attn.w_msa.proj.weight", + "stages.2.blocks.5.attn.w_msa.q_bias", + "stages.2.blocks.5.attn.w_msa.qkv.weight", + "stages.2.blocks.5.attn.w_msa.v_bias", + "stages.2.blocks.5.ffn.layers.0.bias", + "stages.2.blocks.5.ffn.layers.0.weight", + "stages.2.blocks.5.ffn.layers.3.bias", + "stages.2.blocks.5.ffn.layers.3.weight", + "stages.2.blocks.5.norm1.bias", + "stages.2.blocks.5.norm1.weight", + "stages.2.blocks.5.norm2.bias", + "stages.2.blocks.5.norm2.weight", + "stages.2.blocks.5.norm3.bias", + "stages.2.blocks.5.norm3.weight", + "stages.2.blocks.6.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.6.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.6.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.6.attn.w_msa.logit_scale", + "stages.2.blocks.6.attn.w_msa.proj.bias", + "stages.2.blocks.6.attn.w_msa.proj.weight", + "stages.2.blocks.6.attn.w_msa.q_bias", + "stages.2.blocks.6.attn.w_msa.qkv.weight", + "stages.2.blocks.6.attn.w_msa.v_bias", + "stages.2.blocks.6.ffn.layers.0.bias", + "stages.2.blocks.6.ffn.layers.0.weight", + "stages.2.blocks.6.ffn.layers.3.bias", + "stages.2.blocks.6.ffn.layers.3.weight", + "stages.2.blocks.6.norm1.bias", + "stages.2.blocks.6.norm1.weight", + "stages.2.blocks.6.norm2.bias", + "stages.2.blocks.6.norm2.weight", + "stages.2.blocks.7.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.7.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.7.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.7.attn.w_msa.logit_scale", + "stages.2.blocks.7.attn.w_msa.proj.bias", + "stages.2.blocks.7.attn.w_msa.proj.weight", + "stages.2.blocks.7.attn.w_msa.q_bias", + "stages.2.blocks.7.attn.w_msa.qkv.weight", + "stages.2.blocks.7.attn.w_msa.v_bias", + "stages.2.blocks.7.ffn.layers.0.bias", + "stages.2.blocks.7.ffn.layers.0.weight", + "stages.2.blocks.7.ffn.layers.3.bias", + "stages.2.blocks.7.ffn.layers.3.weight", + "stages.2.blocks.7.norm1.bias", + "stages.2.blocks.7.norm1.weight", + "stages.2.blocks.7.norm2.bias", + "stages.2.blocks.7.norm2.weight", + "stages.2.blocks.8.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.8.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.8.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.8.attn.w_msa.logit_scale", + "stages.2.blocks.8.attn.w_msa.proj.bias", + "stages.2.blocks.8.attn.w_msa.proj.weight", + "stages.2.blocks.8.attn.w_msa.q_bias", + "stages.2.blocks.8.attn.w_msa.qkv.weight", + "stages.2.blocks.8.attn.w_msa.v_bias", + "stages.2.blocks.8.ffn.layers.0.bias", + "stages.2.blocks.8.ffn.layers.0.weight", + "stages.2.blocks.8.ffn.layers.3.bias", + "stages.2.blocks.8.ffn.layers.3.weight", + "stages.2.blocks.8.norm1.bias", + "stages.2.blocks.8.norm1.weight", + "stages.2.blocks.8.norm2.bias", + "stages.2.blocks.8.norm2.weight", + "stages.2.blocks.9.attn.w_msa.cpb_mlp.0.bias", + "stages.2.blocks.9.attn.w_msa.cpb_mlp.0.weight", + "stages.2.blocks.9.attn.w_msa.cpb_mlp.2.weight", + "stages.2.blocks.9.attn.w_msa.logit_scale", + "stages.2.blocks.9.attn.w_msa.proj.bias", + "stages.2.blocks.9.attn.w_msa.proj.weight", + "stages.2.blocks.9.attn.w_msa.q_bias", + "stages.2.blocks.9.attn.w_msa.qkv.weight", + "stages.2.blocks.9.attn.w_msa.v_bias", + "stages.2.blocks.9.ffn.layers.0.bias", + "stages.2.blocks.9.ffn.layers.0.weight", + "stages.2.blocks.9.ffn.layers.3.bias", + "stages.2.blocks.9.ffn.layers.3.weight", + "stages.2.blocks.9.norm1.bias", + "stages.2.blocks.9.norm1.weight", + "stages.2.blocks.9.norm2.bias", + "stages.2.blocks.9.norm2.weight", + "stages.2.downsample.norm.bias", + "stages.2.downsample.norm.weight", + "stages.2.downsample.reduction.weight", + "stages.3.blocks.0.attn.w_msa.cpb_mlp.0.bias", + "stages.3.blocks.0.attn.w_msa.cpb_mlp.0.weight", + "stages.3.blocks.0.attn.w_msa.cpb_mlp.2.weight", + "stages.3.blocks.0.attn.w_msa.logit_scale", + "stages.3.blocks.0.attn.w_msa.proj.bias", + "stages.3.blocks.0.attn.w_msa.proj.weight", + "stages.3.blocks.0.attn.w_msa.q_bias", + "stages.3.blocks.0.attn.w_msa.qkv.weight", + "stages.3.blocks.0.attn.w_msa.v_bias", + "stages.3.blocks.0.ffn.layers.0.bias", + "stages.3.blocks.0.ffn.layers.0.weight", + "stages.3.blocks.0.ffn.layers.3.bias", + "stages.3.blocks.0.ffn.layers.3.weight", + "stages.3.blocks.0.norm1.bias", + "stages.3.blocks.0.norm1.weight", + "stages.3.blocks.0.norm2.bias", + "stages.3.blocks.0.norm2.weight", + "stages.3.blocks.1.attn.w_msa.cpb_mlp.0.bias", + "stages.3.blocks.1.attn.w_msa.cpb_mlp.0.weight", + "stages.3.blocks.1.attn.w_msa.cpb_mlp.2.weight", + "stages.3.blocks.1.attn.w_msa.logit_scale", + "stages.3.blocks.1.attn.w_msa.proj.bias", + "stages.3.blocks.1.attn.w_msa.proj.weight", + "stages.3.blocks.1.attn.w_msa.q_bias", + "stages.3.blocks.1.attn.w_msa.qkv.weight", + "stages.3.blocks.1.attn.w_msa.v_bias", + "stages.3.blocks.1.ffn.layers.0.bias", + "stages.3.blocks.1.ffn.layers.0.weight", + "stages.3.blocks.1.ffn.layers.3.bias", + "stages.3.blocks.1.ffn.layers.3.weight", + "stages.3.blocks.1.norm1.bias", + "stages.3.blocks.1.norm1.weight", + "stages.3.blocks.1.norm2.bias", + "stages.3.blocks.1.norm2.weight", + "stages.3.downsample.norm.bias", + "stages.3.downsample.norm.weight", + "stages.3.downsample.reduction.weight", + "vocabulary_token", + "vocabulary_weight" + ] +} diff --git a/skysensepp-swinv2-msl-hr/model.safetensors b/skysensepp-swinv2-msl-hr/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..57514019d72cab599a6488e2b4ca70b044c931ba --- /dev/null +++ b/skysensepp-swinv2-msl-hr/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2727a8049e5cc0fe0e16a1c7c4ccb6f4d9cb1df2ed3eae2dba857b507dd49d9d +size 2658512808 diff --git a/skysensepp-swinv2-msl-hr/modeling_skysensepp_swinv2_msl.py b/skysensepp-swinv2-msl-hr/modeling_skysensepp_swinv2_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0ab382892dfc61d579821983bc8f6b68b84d17 --- /dev/null +++ b/skysensepp-swinv2-msl-hr/modeling_skysensepp_swinv2_msl.py @@ -0,0 +1,343 @@ +"""SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace).""" + +from copy import deepcopy +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig +from .modeling_utils import ( + DropPath, + FFN, + PatchEmbed, + PatchMerging, + ShiftWindowMSA, + to_2tuple, +) + + +class SwinBlockV2(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int = 8, + shift: bool = False, + extra_norm: bool = False, + ffn_ratio: float = 4.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + with_cp: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.with_cp = with_cp + self.extra_norm = extra_norm + self.attn = ShiftWindowMSA( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=window_size, + shift_size=window_size // 2 if shift else 0, + drop_path=drop_path, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + self.norm1 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=int(embed_dims * ffn_ratio), + num_fcs=2, + drop_path=drop_path, + act_layer=nn.GELU, + add_identity=False, + ) + self.norm2 = nn.LayerNorm(embed_dims) + if self.extra_norm: + self.norm3 = nn.LayerNorm(embed_dims) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + def _inner_forward(x): + identity = x + x = self.attn(x, hw_shape) + x = self.norm1(x) + x = x + identity + + identity = x + x = self.ffn(x) + x = self.norm2(x) + x = x + identity + + if self.extra_norm: + x = self.norm3(x) + return x + + if self.with_cp and x.requires_grad: + x = cp.checkpoint(_inner_forward, x, use_reentrant=False) + else: + x = _inner_forward(x) + return x + + +class SwinBlockV2Sequence(nn.Module): + def __init__( + self, + embed_dims: int, + depth: int, + num_heads: int, + window_size: int = 8, + downsample: bool = False, + drop_paths: Union[Sequence[float], float] = 0.0, + with_cp: bool = False, + pad_small_map: bool = False, + extra_norm_every_n_blocks: int = 0, + pretrained_window_size: int = 0, + is_post_norm_downsample: bool = True, + ): + super().__init__() + if not isinstance(drop_paths, Sequence): + drop_paths = [drop_paths] * depth + + if downsample: + self.out_channels = 2 * embed_dims + self.downsample = PatchMerging( + in_channels=embed_dims, + out_channels=self.out_channels, + is_post_norm=is_post_norm_downsample, + ) + else: + self.out_channels = embed_dims + self.downsample = None + + self.blocks = nn.ModuleList() + for i in range(depth): + extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0 + self.blocks.append( + SwinBlockV2( + embed_dims=self.out_channels, + num_heads=num_heads, + window_size=window_size, + shift=(i % 2 == 1), + extra_norm=extra_norm, + drop_path=drop_paths[i], + with_cp=with_cp, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + ) + + def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + if self.downsample is not None: + x, out_shape = self.downsample(x, in_shape) + else: + out_shape = in_shape + + for block in self.blocks: + x = block(x, out_shape) + return x, out_shape + + +class ProjMHSA(nn.Module): + """Projected multi-head self-attention used in SkySense++ HR backbone.""" + + def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True): + super().__init__() + self.proj_in = nn.Linear(embed_dims, proj_dims) + self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias) + self.proj_out = nn.Linear(proj_dims, embed_dims) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj_in(x) + x, _ = self.attn(x, x, x) + return self.proj_out(x) + + +class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusSwinV2MSLConfig + base_model_prefix = "skysensepp_swinv2_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel): + """SkySense++ HR backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig): + super().__init__(config) + + self.num_layers = len(config.depths) + self.out_indices = config.out_indices + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.patch_size = config.patch_size + + if isinstance(config.window_size, int): + window_sizes = [config.window_size] * self.num_layers + else: + window_sizes = list(config.window_size) + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm, + input_size=config.img_size, + ) + + self.use_abs_pos_embed = config.use_abs_pos_embed + if self.use_abs_pos_embed: + patch_resolution = self.patch_embed.init_out_size + num_patches = patch_resolution[0] * patch_resolution[1] + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + total_depth = sum(config.depths) + if total_depth > 1: + dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)] + else: + dpr = [0.0] + + self.stages = nn.ModuleList() + embed_dims_list = [config.embed_dims] + for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)): + stage = SwinBlockV2Sequence( + embed_dims=embed_dims_list[-1], + depth=depth, + num_heads=num_heads, + window_size=window_sizes[i], + downsample=(i > 0), + drop_paths=dpr[:depth], + with_cp=config.with_cp, + pad_small_map=config.pad_small_map, + extra_norm_every_n_blocks=config.extra_norm_every_n_blocks, + pretrained_window_size=config.pretrained_window_sizes[i], + is_post_norm_downsample=config.is_post_norm_downsample, + ) + self.stages.append(stage) + dpr = dpr[depth:] + embed_dims_list.append(stage.out_channels) + + for i in self.out_indices: + self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1])) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = ProjMHSA(352, 256, num_heads=16) + self.attn2 = ProjMHSA(704, 512, num_heads=16) + self.attn3 = ProjMHSA(1408, 1024, num_heads=16) + self.norm_attn = nn.LayerNorm(1408) + + self.post_init() + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + if self.use_abs_pos_embed: + x = x + self.absolute_pos_embed + x = self.drop_after_pos(x) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, stage in enumerate(self.stages): + x, hw_shape = stage(x, hw_shape) + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + x = x + attention_blocks[i](x) + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + out = norm_layer(x) + out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous() + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-swinv2-msl-hr/modeling_skysensepp_vit_msl.py b/skysensepp-swinv2-msl-hr/modeling_skysensepp_vit_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..6e56ac39c612c2eb87baf00b1e5afc071dc4d060 --- /dev/null +++ b/skysensepp-swinv2-msl-hr/modeling_skysensepp_vit_msl.py @@ -0,0 +1,265 @@ +"""SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace).""" + +import math +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig +from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + feedforward_channels: int, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + num_fcs: int = 2, + qkv_bias: bool = True, + with_cp: bool = False, + ): + super().__init__() + self.with_cp = with_cp + self.norm1 = nn.LayerNorm(embed_dims) + self.attn = nn.MultiheadAttention( + embed_dim=embed_dims, + num_heads=num_heads, + dropout=attn_drop_rate, + bias=qkv_bias, + batch_first=True, + ) + self.proj_drop = nn.Dropout(drop_rate) + self.norm2 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=num_fcs, + ffn_drop=drop_rate, + drop_path=drop_path_rate, + act_layer=nn.GELU, + add_identity=True, + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + def _inner_forward(x): + residual = x + x_norm = self.norm1(x) + attn_out, _ = self.attn(x_norm, x_norm, x_norm) + attn_out = self.proj_drop(attn_out) + x = residual + self.drop_path(attn_out) + return self.ffn(self.norm2(x), identity=x) + + if self.with_cp and x.requires_grad: + return cp.checkpoint(_inner_forward, x, use_reentrant=False) + return _inner_forward(x) + + +class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusViTMSLConfig + base_model_prefix = "skysensepp_vit_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel): + """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusViTMSLConfig): + super().__init__(config) + + img_size = to_2tuple(config.img_size) + self.img_size = img_size + self.patch_size = config.patch_size + self.with_cls_token = config.with_cls_token + self.output_cls_token = config.output_cls_token + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.interpolate_mode = "bicubic" + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm if config.patch_norm else None, + ) + + num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size) + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + out_indices = list(config.out_indices) + self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices] + + num_layers = config.num_layers + if num_layers > 1: + dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)] + else: + dpr = [0.0] + + self.layers = nn.ModuleList() + for i in range(config.num_layers): + self.layers.append( + TransformerEncoderLayer( + embed_dims=config.embed_dims, + num_heads=config.num_heads, + feedforward_channels=config.mlp_ratio * config.embed_dims, + attn_drop_rate=config.attn_drop_rate, + drop_rate=config.drop_rate, + drop_path_rate=dpr[i], + num_fcs=2, + qkv_bias=config.qkv_bias, + with_cp=config.with_cp, + ) + ) + + self.final_norm = config.final_norm + if config.final_norm: + self.norm = nn.LayerNorm(config.embed_dims) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.norm_attn = nn.LayerNorm(config.embed_dims) + + self.post_init() + + @staticmethod + def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"): + pos_h, pos_w = pos_shape + pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :] + pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) + pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode) + return torch.flatten(pos_embed_weight, 2).transpose(1, 2) + + def _pos_embedding(self, patched_img, hw_shape, pos_embed): + x_len, pos_len = patched_img.shape[1], pos_embed.shape[1] + if x_len != pos_len: + pos_h = self.img_size[0] // self.patch_size + pos_w = self.img_size[1] // self.patch_size + pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode) + return self.drop_after_pos(patched_img + pos_embed) + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + x = self._pos_embedding(x, hw_shape, self.pos_embed) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, layer in enumerate(self.layers): + x = layer(x) + + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + attn_out, _ = attention_blocks[i](x, x, x) + x = x + attn_out + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm: + x = self.norm(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + out = x + out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous() + if self.output_cls_token: + out = [out, x[:, 0]] + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-swinv2-msl-hr/modeling_utils.py b/skysensepp-swinv2-msl-hr/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..93feae77f3a3c46e65167f2ed312a25b7cd4ad3a --- /dev/null +++ b/skysensepp-swinv2-msl-hr/modeling_utils.py @@ -0,0 +1,557 @@ +"""SkySense: Pure PyTorch + HuggingFace Transformers implementation. + +Shared utility modules used across SkySense model implementations. +""" + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def to_2tuple(x): + """Convert to a 2-tuple.""" + if isinstance(x, (list, tuple)): + return tuple(x) + return (x, x) + + +class DropPath(nn.Module): + """Drop paths (stochastic depth) per sample. + + Args: + drop_prob (float): Probability of dropping a path. Default: 0.0. + """ + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor = torch.floor(random_tensor + keep_prob) + output = x / keep_prob * random_tensor + return output + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding using Conv2d. + + Args: + in_channels (int): Number of input channels. Default: 3. + embed_dims (int): Embedding dimension. Default: 96. + kernel_size (int): Kernel size of the projection. Default: 4. + stride (int): Stride of the projection. Default: 4. + padding (int): Padding of the projection. Default: 0. + norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm. + input_size (int or tuple or None): Input resolution for calculating output size. + """ + + def __init__( + self, + in_channels: int = 3, + embed_dims: int = 96, + kernel_size: int = 4, + stride: int = 4, + padding: int = 0, + norm_layer: Optional[type] = nn.LayerNorm, + input_size: Optional[int] = None, + ): + super().__init__() + self.projection = nn.Conv2d( + in_channels, embed_dims, + kernel_size=kernel_size, stride=stride, padding=padding, + ) + self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity() + + # Compute init output size if input_size is given + if input_size is not None: + input_size = to_2tuple(input_size) + self.init_out_size = ( + (input_size[0] - kernel_size + 2 * padding) // stride + 1, + (input_size[1] - kernel_size + 2 * padding) // stride + 1, + ) + else: + self.init_out_size = None + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]: + x = self.projection(x) # (B, C, H, W) + out_size = (x.shape[2], x.shape[3]) + x = x.flatten(2).transpose(1, 2) # (B, H*W, C) + x = self.norm(x) + return x, out_size + + +class FFN(nn.Module): + """Feed-Forward Network. + + Args: + embed_dims (int): Input dimension. + feedforward_channels (int): Hidden dimension. + num_fcs (int): Number of FC layers. Default: 2. + ffn_drop (float): Dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + act_layer (nn.Module): Activation layer class. Default: nn.GELU. + add_identity (bool): Whether to add identity connection. Default: True. + """ + + def __init__( + self, + embed_dims: int, + feedforward_channels: int, + num_fcs: int = 2, + ffn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: type = nn.GELU, + add_identity: bool = True, + ): + super().__init__() + assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}" + self.embed_dims = embed_dims + self.feedforward_channels = feedforward_channels + self.add_identity = add_identity + + layers = [] + in_channels = embed_dims + for i in range(num_fcs - 1): + layers.append(nn.Linear(in_channels, feedforward_channels)) + layers.append(act_layer()) + layers.append(nn.Dropout(ffn_drop)) + in_channels = feedforward_channels + layers.append(nn.Linear(feedforward_channels, embed_dims)) + layers.append(nn.Dropout(ffn_drop)) + self.layers = nn.Sequential(*layers) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor: + out = self.layers(x) + out = self.drop_path(out) + if self.add_identity: + if identity is None: + identity = x + out = out + identity + return out + + +class WindowMSAV2(nn.Module): + """Window-based Multi-head Self-Attention for Swin Transformer V2. + + Uses cosine attention and log-spaced continuous position bias (log-CPB). + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): Window size (Wh, Ww). + pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0). + qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Output projection dropout rate. Default: 0.0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: Tuple[int, int], + pretrained_window_size: Tuple[int, int] = (0, 0), + qkv_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ): + super().__init__() + self.embed_dims = embed_dims + self.num_heads = num_heads + self.window_size = window_size + self.pretrained_window_size = pretrained_window_size + + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1)))) + + # MLP for continuous relative position bias (log-CPB) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) + + # Build relative coords table + self._build_relative_coords_table() + # Build relative position index + self._build_relative_position_index() + + self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False) + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(embed_dims)) + self.v_bias = nn.Parameter(torch.zeros(embed_dims)) + else: + self.q_bias = None + self.v_bias = None + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(embed_dims, embed_dims) + self.proj_drop = nn.Dropout(proj_drop) + self.softmax = nn.Softmax(dim=-1) + + def _build_relative_coords_table(self): + """Build the relative coordinates table for log-CPB.""" + Wh, Ww = self.window_size + # Table of relative coordinates + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2) + + # Normalize to [-1, 1] and apply log-scale + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 # normalize to -8, 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + self.register_buffer("relative_coords_table", coords_table) + + def _build_relative_position_index(self): + """Build the pairwise relative position index for each window token.""" + Wh, Ww = self.window_size + coords_h = torch.arange(Wh) + coords_w = torch.arange(Ww) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij')) + coords_flatten = coords.view(2, -1) + + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # (2, Wh*Ww, Wh*Ww) + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += Wh - 1 + relative_coords[:, :, 1] += Ww - 1 + relative_coords[:, :, 0] *= 2 * Ww - 1 + relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww) + self.register_buffer("relative_position_index", relative_position_index) + + def _compute_position_bias(self, N): + """Compute relative position bias, supporting dynamic window sizes. + + The log-CPB (Continuous Position Bias) MLP can generalize to any window + size by computing bias from normalized relative coordinates. + """ + init_N = self.window_size[0] * self.window_size[1] + if N == init_N: + # Use pre-built tables + relative_position_bias_table = self.cpb_mlp( + self.relative_coords_table + ).view(-1, self.num_heads) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view(N, N, -1) + else: + # Dynamic: compute for actual window size on-the-fly + Wh = Ww = int(math.sqrt(N)) + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + # Build position index for actual window size + ch = torch.arange(Wh, device=self.logit_scale.device) + cw = torch.arange(Ww, device=self.logit_scale.device) + coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij')) + coords_flat = coords.view(2, -1) + rel = coords_flat[:, :, None] - coords_flat[:, None, :] + rel = rel.permute(1, 2, 0).contiguous() + rel[:, :, 0] += Wh - 1 + rel[:, :, 1] += Ww - 1 + rel[:, :, 0] *= 2 * Ww - 1 + pos_index = rel.sum(-1) + + bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads) + relative_position_bias = bias_table[ + pos_index.view(-1) + ].view(N, N, -1) + + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + return relative_position_bias + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: (num_windows*B, N, C) where N = Wh*Ww + mask: (num_windows, N, N) or None + """ + B_, N, C = x.shape + + # Compute QKV with bias + if self.q_bias is not None: + qkv_bias = torch.cat( + (self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias)) + qkv = F.linear(x, self.qkv.weight, qkv_bias) + else: + qkv = self.qkv(x) + + qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + + # Cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=math.log(1.0 / 0.01) + ).exp() + attn = attn * logit_scale + + # Log-CPB relative position bias (supports dynamic window sizes) + relative_position_bias = self._compute_position_bias(N) + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + attn = attn + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + + attn = self.softmax(attn) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class ShiftWindowMSA(nn.Module): + """Shifted Window Multi-head Self-Attention. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. Default: 0. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Projection dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + pad_small_map (bool): Pad small feature maps to window size. Default: False. + pretrained_window_size (int): Pretrained window size. Default: 0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int, + shift_size: int = 0, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.window_size = window_size + self.shift_size = shift_size + self.pad_small_map = pad_small_map + + self.w_msa = WindowMSAV2( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=to_2tuple(window_size), + pretrained_window_size=to_2tuple(pretrained_window_size), + attn_drop=attn_drop, + proj_drop=proj_drop, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W, f"Input length {L} != H*W ({H}*{W})" + + x = x.view(B, H, W, C) + + window_size = self.window_size + shift_size = self.shift_size + + # Pad or shrink window + if self.pad_small_map: + pad_r = (window_size - W % window_size) % window_size + pad_b = (window_size - H % window_size) % window_size + x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b)) + _, Hp, Wp, _ = x.shape + else: + Hp, Wp = H, W + if window_size > Hp: + window_size = Hp + shift_size = 0 + if window_size > Wp: + window_size = Wp + shift_size = 0 + + # Compute attention mask for SW-MSA + attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device) + + # Cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) + + # Partition windows + x_windows = self._window_partition(x, window_size) + # (num_windows*B, window_size*window_size, C) + + # W-MSA/SW-MSA + attn_windows = self.w_msa(x_windows, mask=attn_mask) + + # Merge windows + x = self._window_reverse(attn_windows, window_size, Hp, Wp) + + # Reverse cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) + + if self.pad_small_map and (pad_r > 0 or pad_b > 0): + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + x = self.drop_path(x) + return x + + @staticmethod + def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor: + """Partition into non-overlapping windows.""" + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous() + windows = windows.view(-1, window_size * window_size, C) + return windows + + @staticmethod + def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor: + """Reverse window partition.""" + B_nW = windows.shape[0] + nH = H // window_size + nW = W // window_size + B = B_nW // (nH * nW) + x = windows.view(B, nH, nW, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous() + x = x.view(B, H, W, -1) + return x + + @staticmethod + def _compute_attn_mask(H, W, window_size, shift_size, device): + """Compute attention mask for shifted window attention.""" + if shift_size <= 0: + return None + img_mask = torch.zeros((1, H, W, 1), device=device) + h_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + w_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + # Partition mask + mask_windows = img_mask.view( + 1, H // window_size, window_size, W // window_size, window_size, 1 + ) + mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous() + mask_windows = mask_windows.view(-1, window_size * window_size) + + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0) + attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0) + return attn_mask + + +class PatchMerging(nn.Module): + """Patch Merging Layer for downsampling (2x). + + Args: + in_channels (int): Input channels. + out_channels (int): Output channels. + norm_layer (type): Normalization layer. Default: nn.LayerNorm. + is_post_norm (bool): Apply norm after linear. Default: True. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + norm_layer: type = nn.LayerNorm, + is_post_norm: bool = True, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.is_post_norm = is_post_norm + self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False) + if is_post_norm: + self.norm = norm_layer(out_channels) + else: + self.norm = norm_layer(4 * in_channels) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W + + x = x.view(B, H, W, C) + + # Pad if needed + pad_h = H % 2 + pad_w = W % 2 + if pad_h or pad_w: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], dim=-1) + + out_h = (H + pad_h) // 2 + out_w = (W + pad_w) // 2 + x = x.view(B, out_h * out_w, 4 * C) + + if self.is_post_norm: + x = self.reduction(x) + x = self.norm(x) + else: + x = self.norm(x) + x = self.reduction(x) + + return x, (out_h, out_w) diff --git a/skysensepp-swinv2-msl-hr/pipeline_skysensepp.py b/skysensepp-swinv2-msl-hr/pipeline_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9a5354361c9a08abd1a82b3df8d4ec678a209b --- /dev/null +++ b/skysensepp-swinv2-msl-hr/pipeline_skysensepp.py @@ -0,0 +1,86 @@ +"""Custom HuggingFace pipeline for SkySense++ MSL feature extraction.""" + +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline): + """Pipeline for SkySense++ MSL backbones. + + Expects image tensors plus semantic annotation maps (class indices). + """ + + def _sanitize_parameters( + self, + annotation=None, + mask=None, + output_hidden_states=None, + **kwargs, + ): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + + if annotation is not None: + preprocess_params["annotation"] = annotation + if mask is not None: + forward_params["mask"] = mask + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + + return preprocess_params, forward_params, postprocess_params + + def preprocess( + self, + pixel_values: Any, + annotation: Optional[Any] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if isinstance(pixel_values, dict): + annotation = pixel_values.get("annotation", annotation) + pixel_values = pixel_values.get("pixel_values", pixel_values) + + if isinstance(pixel_values, np.ndarray): + pixel_values = torch.from_numpy(pixel_values).float() + elif not isinstance(pixel_values, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}" + ) + + if annotation is None: + raise ValueError("SkySense++ MSL models require an `annotation` semantic map.") + + if isinstance(annotation, np.ndarray): + annotation = torch.from_numpy(annotation).long() + elif not isinstance(annotation, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for annotation, got {type(annotation)}" + ) + + if pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + if annotation.ndim == 2: + annotation = annotation.unsqueeze(0) + + return {"pixel_values": pixel_values, "annotation": annotation} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + pixel_values=model_inputs["pixel_values"], + annotation=model_inputs["annotation"], + mask=kwargs.get("mask"), + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"last_hidden_state": outputs.last_hidden_state} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-swinv2-msl-hr/pipeline_skysensepp_fusion.py b/skysensepp-swinv2-msl-hr/pipeline_skysensepp_fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..80f070fa9fd3f18a053caa1afef9dc142bce2597 --- /dev/null +++ b/skysensepp-swinv2-msl-hr/pipeline_skysensepp_fusion.py @@ -0,0 +1,53 @@ +"""Optional pipeline for SkySense++ fusion neck.""" + +from typing import Any, Dict + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusFusionNeckPipeline(Pipeline): + """Pipeline for the optional SkySense++ fusion neck module. + + Expects concatenated multi-modal tokens per spatial location: + ``(batch, num_modalities, input_dims)``. + """ + + def _sanitize_parameters(self, output_hidden_states=None, **kwargs): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + return preprocess_params, forward_params, postprocess_params + + def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]: + if isinstance(hidden_states, dict): + hidden_states = hidden_states["hidden_states"] + + if isinstance(hidden_states, np.ndarray): + hidden_states = torch.from_numpy(hidden_states).float() + elif not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}" + ) + if hidden_states.ndim == 2: + hidden_states = hidden_states.unsqueeze(0) + return {"hidden_states": hidden_states} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + hidden_states=model_inputs["hidden_states"], + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"pooler_output": outputs.pooler_output} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-vit-msl-s1/__init__.py b/skysensepp-vit-msl-s1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d49c51cc97fc8d74fa2f98f3c4677a5b59f8bc2 --- /dev/null +++ b/skysensepp-vit-msl-s1/__init__.py @@ -0,0 +1,25 @@ +"""SkySense++: Multi-Modal Remote Sensing Foundation Model (HuggingFace).""" + +from .configuration_skysensepp import ( + SkySensePlusPlusSwinV2MSLConfig, + SkySensePlusPlusViTMSLConfig, +) +from .modeling_skysensepp_swinv2_msl import ( + SkySensePlusPlusSwinV2MSLModel, + SkySensePlusPlusSwinV2MSLPreTrainedModel, +) +from .modeling_skysensepp_vit_msl import ( + SkySensePlusPlusViTMSLModel, + SkySensePlusPlusViTMSLPreTrainedModel, +) +from .pipeline_skysensepp import SkySensePlusPlusMSLFeatureExtractionPipeline + +__all__ = [ + "SkySensePlusPlusSwinV2MSLConfig", + "SkySensePlusPlusViTMSLConfig", + "SkySensePlusPlusSwinV2MSLModel", + "SkySensePlusPlusSwinV2MSLPreTrainedModel", + "SkySensePlusPlusViTMSLModel", + "SkySensePlusPlusViTMSLPreTrainedModel", + "SkySensePlusPlusMSLFeatureExtractionPipeline", +] diff --git a/skysensepp-vit-msl-s1/config.json b/skysensepp-vit-msl-s1/config.json new file mode 100644 index 0000000000000000000000000000000000000000..195bf3287ce3d6208e84ed17703ce0d4c58f3da5 --- /dev/null +++ b/skysensepp-vit-msl-s1/config.json @@ -0,0 +1,68 @@ +{ + "return_dict": true, + "output_hidden_states": false, + "dtype": "float32", + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": [ + "SkySensePlusPlusViTMSLModel" + ], + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "transformers_version": "5.0.0", + "img_size": 16, + "patch_size": 4, + "in_channels": 2, + "embed_dims": 1024, + "num_layers": 24, + "num_heads": 16, + "mlp_ratio": 4, + "out_indices": [ + 5, + 11, + 17, + 23 + ], + "qkv_bias": true, + "drop_rate": 0.0, + "attn_drop_rate": 0.0, + "drop_path_rate": 0.3, + "with_cls_token": false, + "output_cls_token": false, + "patch_norm": false, + "final_norm": false, + "with_cp": false, + "vocabulary_size": 64, + "num_vocabulary_tokens": 65, + "merge_stage": 4, + "use_attn": false, + "modality": "s1", + "model_type": "skysensepp_vit_msl", + "output_attentions": false, + "auto_map": { + "AutoConfig": "configuration_skysensepp.SkySensePlusPlusViTMSLConfig", + "AutoModel": "modeling_skysensepp_vit_msl.SkySensePlusPlusViTMSLModel" + }, + "custom_pipelines": { + "skysensepp-feature-extraction": { + "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + }, + "image-feature-extraction": { + "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + } +} diff --git a/skysensepp-vit-msl-s1/configuration_skysensepp.py b/skysensepp-vit-msl-s1/configuration_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..e086dd1de7d378e81cc124a89a36e33f45bc4b3f --- /dev/null +++ b/skysensepp-vit-msl-s1/configuration_skysensepp.py @@ -0,0 +1,124 @@ +"""Configuration classes for SkySense++ MSL backbones.""" + +from transformers import PretrainedConfig + + +class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig): + """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical).""" + + model_type = "skysensepp_swinv2_msl" + + arch_zoo = { + "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0}, + "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0}, + "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6}, + "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6}, + } + + def __init__( + self, + arch="huge", + img_size=224, + patch_size=4, + in_channels=3, + window_size=8, + drop_rate=0.0, + drop_path_rate=0.2, + out_indices=(0, 1, 2, 3), + use_abs_pos_embed=False, + with_cp=False, + pad_small_map=False, + pretrained_window_sizes=(0, 0, 0, 0), + is_post_norm_downsample=True, + vocabulary_size=64, + merge_stage=2, + use_attn=True, + **kwargs, + ): + super().__init__(**kwargs) + + arch = arch.lower() + if arch not in self.arch_zoo: + raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}") + arch_settings = self.arch_zoo[arch] + + self.arch = arch + self.embed_dims = arch_settings["embed_dims"] + self.depths = arch_settings["depths"] + self.num_heads = arch_settings["num_heads"] + self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"] + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.window_size = window_size + self.drop_rate = drop_rate + self.drop_path_rate = drop_path_rate + self.out_indices = list(out_indices) + self.use_abs_pos_embed = use_abs_pos_embed + self.with_cp = with_cp + self.pad_small_map = pad_small_map + self.pretrained_window_sizes = list(pretrained_window_sizes) + self.is_post_norm_downsample = is_post_norm_downsample + + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + + +class SkySensePlusPlusViTMSLConfig(PretrainedConfig): + """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1).""" + + model_type = "skysensepp_vit_msl" + + def __init__( + self, + img_size=16, + patch_size=4, + in_channels=10, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + out_indices=(5, 11, 17, 23), + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=False, + output_cls_token=False, + patch_norm=False, + final_norm=False, + with_cp=False, + vocabulary_size=64, + merge_stage=4, + use_attn=False, + modality="s2", + **kwargs, + ): + super().__init__(**kwargs) + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.out_indices = list(out_indices) + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.patch_norm = patch_norm + self.final_norm = final_norm + self.with_cp = with_cp + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + self.modality = modality diff --git a/skysensepp-vit-msl-s1/conversion_manifest.json b/skysensepp-vit-msl-s1/conversion_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..007f976f4694a344503667546ac14b21197caab2 --- /dev/null +++ b/skysensepp-vit-msl-s1/conversion_manifest.json @@ -0,0 +1,305 @@ +{ + "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release_s1.pth", + "modality": "s1", + "model_class": "SkySensePlusPlusViTMSLModel", + "num_tensors": 295, + "missing_keys": [], + "unexpected_keys": [], + "tensor_names": [ + "cls_token", + "layers.0.attn.in_proj_bias", + "layers.0.attn.in_proj_weight", + "layers.0.attn.out_proj.bias", + "layers.0.attn.out_proj.weight", + "layers.0.ffn.layers.0.bias", + "layers.0.ffn.layers.0.weight", + "layers.0.ffn.layers.3.bias", + "layers.0.ffn.layers.3.weight", + "layers.0.norm1.bias", + "layers.0.norm1.weight", + "layers.0.norm2.bias", + "layers.0.norm2.weight", + "layers.1.attn.in_proj_bias", + "layers.1.attn.in_proj_weight", + "layers.1.attn.out_proj.bias", + "layers.1.attn.out_proj.weight", + "layers.1.ffn.layers.0.bias", + "layers.1.ffn.layers.0.weight", + "layers.1.ffn.layers.3.bias", + "layers.1.ffn.layers.3.weight", + "layers.1.norm1.bias", + "layers.1.norm1.weight", + "layers.1.norm2.bias", + "layers.1.norm2.weight", + "layers.10.attn.in_proj_bias", + "layers.10.attn.in_proj_weight", + "layers.10.attn.out_proj.bias", + "layers.10.attn.out_proj.weight", + "layers.10.ffn.layers.0.bias", + "layers.10.ffn.layers.0.weight", + "layers.10.ffn.layers.3.bias", + "layers.10.ffn.layers.3.weight", + "layers.10.norm1.bias", + "layers.10.norm1.weight", + "layers.10.norm2.bias", + "layers.10.norm2.weight", + "layers.11.attn.in_proj_bias", + "layers.11.attn.in_proj_weight", + "layers.11.attn.out_proj.bias", + "layers.11.attn.out_proj.weight", + "layers.11.ffn.layers.0.bias", + "layers.11.ffn.layers.0.weight", + "layers.11.ffn.layers.3.bias", + "layers.11.ffn.layers.3.weight", + "layers.11.norm1.bias", + "layers.11.norm1.weight", + "layers.11.norm2.bias", + "layers.11.norm2.weight", + "layers.12.attn.in_proj_bias", + "layers.12.attn.in_proj_weight", + "layers.12.attn.out_proj.bias", + "layers.12.attn.out_proj.weight", + "layers.12.ffn.layers.0.bias", + "layers.12.ffn.layers.0.weight", + "layers.12.ffn.layers.3.bias", + "layers.12.ffn.layers.3.weight", + "layers.12.norm1.bias", + "layers.12.norm1.weight", + "layers.12.norm2.bias", + "layers.12.norm2.weight", + "layers.13.attn.in_proj_bias", + "layers.13.attn.in_proj_weight", + "layers.13.attn.out_proj.bias", + "layers.13.attn.out_proj.weight", + "layers.13.ffn.layers.0.bias", + "layers.13.ffn.layers.0.weight", + "layers.13.ffn.layers.3.bias", + "layers.13.ffn.layers.3.weight", + "layers.13.norm1.bias", + "layers.13.norm1.weight", + "layers.13.norm2.bias", + "layers.13.norm2.weight", + "layers.14.attn.in_proj_bias", + "layers.14.attn.in_proj_weight", + "layers.14.attn.out_proj.bias", + "layers.14.attn.out_proj.weight", + "layers.14.ffn.layers.0.bias", + "layers.14.ffn.layers.0.weight", + "layers.14.ffn.layers.3.bias", + "layers.14.ffn.layers.3.weight", + "layers.14.norm1.bias", + "layers.14.norm1.weight", + "layers.14.norm2.bias", + "layers.14.norm2.weight", + "layers.15.attn.in_proj_bias", + "layers.15.attn.in_proj_weight", + "layers.15.attn.out_proj.bias", + "layers.15.attn.out_proj.weight", + "layers.15.ffn.layers.0.bias", + "layers.15.ffn.layers.0.weight", + "layers.15.ffn.layers.3.bias", + "layers.15.ffn.layers.3.weight", + "layers.15.norm1.bias", + "layers.15.norm1.weight", + "layers.15.norm2.bias", + "layers.15.norm2.weight", + "layers.16.attn.in_proj_bias", + "layers.16.attn.in_proj_weight", + "layers.16.attn.out_proj.bias", + "layers.16.attn.out_proj.weight", + "layers.16.ffn.layers.0.bias", + "layers.16.ffn.layers.0.weight", + "layers.16.ffn.layers.3.bias", + "layers.16.ffn.layers.3.weight", + "layers.16.norm1.bias", + "layers.16.norm1.weight", + "layers.16.norm2.bias", + "layers.16.norm2.weight", + "layers.17.attn.in_proj_bias", + "layers.17.attn.in_proj_weight", + "layers.17.attn.out_proj.bias", + "layers.17.attn.out_proj.weight", + "layers.17.ffn.layers.0.bias", + "layers.17.ffn.layers.0.weight", + "layers.17.ffn.layers.3.bias", + "layers.17.ffn.layers.3.weight", + "layers.17.norm1.bias", + "layers.17.norm1.weight", + "layers.17.norm2.bias", + "layers.17.norm2.weight", + "layers.18.attn.in_proj_bias", + "layers.18.attn.in_proj_weight", + "layers.18.attn.out_proj.bias", + "layers.18.attn.out_proj.weight", + "layers.18.ffn.layers.0.bias", + "layers.18.ffn.layers.0.weight", + "layers.18.ffn.layers.3.bias", + "layers.18.ffn.layers.3.weight", + "layers.18.norm1.bias", + "layers.18.norm1.weight", + "layers.18.norm2.bias", + "layers.18.norm2.weight", + "layers.19.attn.in_proj_bias", + "layers.19.attn.in_proj_weight", + "layers.19.attn.out_proj.bias", + "layers.19.attn.out_proj.weight", + "layers.19.ffn.layers.0.bias", + "layers.19.ffn.layers.0.weight", + "layers.19.ffn.layers.3.bias", + "layers.19.ffn.layers.3.weight", + "layers.19.norm1.bias", + "layers.19.norm1.weight", + "layers.19.norm2.bias", + "layers.19.norm2.weight", + "layers.2.attn.in_proj_bias", + "layers.2.attn.in_proj_weight", + "layers.2.attn.out_proj.bias", + "layers.2.attn.out_proj.weight", + "layers.2.ffn.layers.0.bias", + "layers.2.ffn.layers.0.weight", + "layers.2.ffn.layers.3.bias", + "layers.2.ffn.layers.3.weight", + "layers.2.norm1.bias", + "layers.2.norm1.weight", + "layers.2.norm2.bias", + "layers.2.norm2.weight", + "layers.20.attn.in_proj_bias", + "layers.20.attn.in_proj_weight", + "layers.20.attn.out_proj.bias", + "layers.20.attn.out_proj.weight", + "layers.20.ffn.layers.0.bias", + "layers.20.ffn.layers.0.weight", + "layers.20.ffn.layers.3.bias", + "layers.20.ffn.layers.3.weight", + "layers.20.norm1.bias", + "layers.20.norm1.weight", + "layers.20.norm2.bias", + "layers.20.norm2.weight", + "layers.21.attn.in_proj_bias", + "layers.21.attn.in_proj_weight", + "layers.21.attn.out_proj.bias", + "layers.21.attn.out_proj.weight", + "layers.21.ffn.layers.0.bias", + "layers.21.ffn.layers.0.weight", + "layers.21.ffn.layers.3.bias", + "layers.21.ffn.layers.3.weight", + "layers.21.norm1.bias", + "layers.21.norm1.weight", + "layers.21.norm2.bias", + "layers.21.norm2.weight", + "layers.22.attn.in_proj_bias", + "layers.22.attn.in_proj_weight", + "layers.22.attn.out_proj.bias", + "layers.22.attn.out_proj.weight", + "layers.22.ffn.layers.0.bias", + "layers.22.ffn.layers.0.weight", + "layers.22.ffn.layers.3.bias", + "layers.22.ffn.layers.3.weight", + "layers.22.norm1.bias", + "layers.22.norm1.weight", + "layers.22.norm2.bias", + "layers.22.norm2.weight", + "layers.23.attn.in_proj_bias", + "layers.23.attn.in_proj_weight", + "layers.23.attn.out_proj.bias", + "layers.23.attn.out_proj.weight", + "layers.23.ffn.layers.0.bias", + "layers.23.ffn.layers.0.weight", + "layers.23.ffn.layers.3.bias", + "layers.23.ffn.layers.3.weight", + "layers.23.norm1.bias", + "layers.23.norm1.weight", + "layers.23.norm2.bias", + "layers.23.norm2.weight", + "layers.3.attn.in_proj_bias", + "layers.3.attn.in_proj_weight", + "layers.3.attn.out_proj.bias", + "layers.3.attn.out_proj.weight", + "layers.3.ffn.layers.0.bias", + "layers.3.ffn.layers.0.weight", + "layers.3.ffn.layers.3.bias", + "layers.3.ffn.layers.3.weight", + "layers.3.norm1.bias", + "layers.3.norm1.weight", + "layers.3.norm2.bias", + "layers.3.norm2.weight", + "layers.4.attn.in_proj_bias", + "layers.4.attn.in_proj_weight", + "layers.4.attn.out_proj.bias", + "layers.4.attn.out_proj.weight", + "layers.4.ffn.layers.0.bias", + "layers.4.ffn.layers.0.weight", + "layers.4.ffn.layers.3.bias", + "layers.4.ffn.layers.3.weight", + "layers.4.norm1.bias", + "layers.4.norm1.weight", + "layers.4.norm2.bias", + "layers.4.norm2.weight", + "layers.5.attn.in_proj_bias", + "layers.5.attn.in_proj_weight", + "layers.5.attn.out_proj.bias", + "layers.5.attn.out_proj.weight", + "layers.5.ffn.layers.0.bias", + "layers.5.ffn.layers.0.weight", + "layers.5.ffn.layers.3.bias", + "layers.5.ffn.layers.3.weight", + "layers.5.norm1.bias", + "layers.5.norm1.weight", + "layers.5.norm2.bias", + "layers.5.norm2.weight", + "layers.6.attn.in_proj_bias", + "layers.6.attn.in_proj_weight", + "layers.6.attn.out_proj.bias", + "layers.6.attn.out_proj.weight", + "layers.6.ffn.layers.0.bias", + "layers.6.ffn.layers.0.weight", + "layers.6.ffn.layers.3.bias", + "layers.6.ffn.layers.3.weight", + "layers.6.norm1.bias", + "layers.6.norm1.weight", + "layers.6.norm2.bias", + "layers.6.norm2.weight", + "layers.7.attn.in_proj_bias", + "layers.7.attn.in_proj_weight", + "layers.7.attn.out_proj.bias", + "layers.7.attn.out_proj.weight", + "layers.7.ffn.layers.0.bias", + "layers.7.ffn.layers.0.weight", + "layers.7.ffn.layers.3.bias", + "layers.7.ffn.layers.3.weight", + "layers.7.norm1.bias", + "layers.7.norm1.weight", + "layers.7.norm2.bias", + "layers.7.norm2.weight", + "layers.8.attn.in_proj_bias", + "layers.8.attn.in_proj_weight", + "layers.8.attn.out_proj.bias", + "layers.8.attn.out_proj.weight", + "layers.8.ffn.layers.0.bias", + "layers.8.ffn.layers.0.weight", + "layers.8.ffn.layers.3.bias", + "layers.8.ffn.layers.3.weight", + "layers.8.norm1.bias", + "layers.8.norm1.weight", + "layers.8.norm2.bias", + "layers.8.norm2.weight", + "layers.9.attn.in_proj_bias", + "layers.9.attn.in_proj_weight", + "layers.9.attn.out_proj.bias", + "layers.9.attn.out_proj.weight", + "layers.9.ffn.layers.0.bias", + "layers.9.ffn.layers.0.weight", + "layers.9.ffn.layers.3.bias", + "layers.9.ffn.layers.3.weight", + "layers.9.norm1.bias", + "layers.9.norm1.weight", + "layers.9.norm2.bias", + "layers.9.norm2.weight", + "mask_token", + "patch_embed.projection.bias", + "patch_embed.projection.weight", + "pos_embed", + "vocabulary_token", + "vocabulary_weight" + ] +} diff --git a/skysensepp-vit-msl-s1/model.safetensors b/skysensepp-vit-msl-s1/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..e660bb171284cc7a63a00261b4f012590d15f3a5 --- /dev/null +++ b/skysensepp-vit-msl-s1/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9a9a658f9efd8f0ebb65d92fe075bec23c1b9edf1fd653b2d7457f2065acaf1f +size 1209741688 diff --git a/skysensepp-vit-msl-s1/modeling_skysensepp_swinv2_msl.py b/skysensepp-vit-msl-s1/modeling_skysensepp_swinv2_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0ab382892dfc61d579821983bc8f6b68b84d17 --- /dev/null +++ b/skysensepp-vit-msl-s1/modeling_skysensepp_swinv2_msl.py @@ -0,0 +1,343 @@ +"""SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace).""" + +from copy import deepcopy +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig +from .modeling_utils import ( + DropPath, + FFN, + PatchEmbed, + PatchMerging, + ShiftWindowMSA, + to_2tuple, +) + + +class SwinBlockV2(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int = 8, + shift: bool = False, + extra_norm: bool = False, + ffn_ratio: float = 4.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + with_cp: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.with_cp = with_cp + self.extra_norm = extra_norm + self.attn = ShiftWindowMSA( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=window_size, + shift_size=window_size // 2 if shift else 0, + drop_path=drop_path, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + self.norm1 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=int(embed_dims * ffn_ratio), + num_fcs=2, + drop_path=drop_path, + act_layer=nn.GELU, + add_identity=False, + ) + self.norm2 = nn.LayerNorm(embed_dims) + if self.extra_norm: + self.norm3 = nn.LayerNorm(embed_dims) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + def _inner_forward(x): + identity = x + x = self.attn(x, hw_shape) + x = self.norm1(x) + x = x + identity + + identity = x + x = self.ffn(x) + x = self.norm2(x) + x = x + identity + + if self.extra_norm: + x = self.norm3(x) + return x + + if self.with_cp and x.requires_grad: + x = cp.checkpoint(_inner_forward, x, use_reentrant=False) + else: + x = _inner_forward(x) + return x + + +class SwinBlockV2Sequence(nn.Module): + def __init__( + self, + embed_dims: int, + depth: int, + num_heads: int, + window_size: int = 8, + downsample: bool = False, + drop_paths: Union[Sequence[float], float] = 0.0, + with_cp: bool = False, + pad_small_map: bool = False, + extra_norm_every_n_blocks: int = 0, + pretrained_window_size: int = 0, + is_post_norm_downsample: bool = True, + ): + super().__init__() + if not isinstance(drop_paths, Sequence): + drop_paths = [drop_paths] * depth + + if downsample: + self.out_channels = 2 * embed_dims + self.downsample = PatchMerging( + in_channels=embed_dims, + out_channels=self.out_channels, + is_post_norm=is_post_norm_downsample, + ) + else: + self.out_channels = embed_dims + self.downsample = None + + self.blocks = nn.ModuleList() + for i in range(depth): + extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0 + self.blocks.append( + SwinBlockV2( + embed_dims=self.out_channels, + num_heads=num_heads, + window_size=window_size, + shift=(i % 2 == 1), + extra_norm=extra_norm, + drop_path=drop_paths[i], + with_cp=with_cp, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + ) + + def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + if self.downsample is not None: + x, out_shape = self.downsample(x, in_shape) + else: + out_shape = in_shape + + for block in self.blocks: + x = block(x, out_shape) + return x, out_shape + + +class ProjMHSA(nn.Module): + """Projected multi-head self-attention used in SkySense++ HR backbone.""" + + def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True): + super().__init__() + self.proj_in = nn.Linear(embed_dims, proj_dims) + self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias) + self.proj_out = nn.Linear(proj_dims, embed_dims) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj_in(x) + x, _ = self.attn(x, x, x) + return self.proj_out(x) + + +class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusSwinV2MSLConfig + base_model_prefix = "skysensepp_swinv2_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel): + """SkySense++ HR backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig): + super().__init__(config) + + self.num_layers = len(config.depths) + self.out_indices = config.out_indices + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.patch_size = config.patch_size + + if isinstance(config.window_size, int): + window_sizes = [config.window_size] * self.num_layers + else: + window_sizes = list(config.window_size) + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm, + input_size=config.img_size, + ) + + self.use_abs_pos_embed = config.use_abs_pos_embed + if self.use_abs_pos_embed: + patch_resolution = self.patch_embed.init_out_size + num_patches = patch_resolution[0] * patch_resolution[1] + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + total_depth = sum(config.depths) + if total_depth > 1: + dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)] + else: + dpr = [0.0] + + self.stages = nn.ModuleList() + embed_dims_list = [config.embed_dims] + for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)): + stage = SwinBlockV2Sequence( + embed_dims=embed_dims_list[-1], + depth=depth, + num_heads=num_heads, + window_size=window_sizes[i], + downsample=(i > 0), + drop_paths=dpr[:depth], + with_cp=config.with_cp, + pad_small_map=config.pad_small_map, + extra_norm_every_n_blocks=config.extra_norm_every_n_blocks, + pretrained_window_size=config.pretrained_window_sizes[i], + is_post_norm_downsample=config.is_post_norm_downsample, + ) + self.stages.append(stage) + dpr = dpr[depth:] + embed_dims_list.append(stage.out_channels) + + for i in self.out_indices: + self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1])) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = ProjMHSA(352, 256, num_heads=16) + self.attn2 = ProjMHSA(704, 512, num_heads=16) + self.attn3 = ProjMHSA(1408, 1024, num_heads=16) + self.norm_attn = nn.LayerNorm(1408) + + self.post_init() + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + if self.use_abs_pos_embed: + x = x + self.absolute_pos_embed + x = self.drop_after_pos(x) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, stage in enumerate(self.stages): + x, hw_shape = stage(x, hw_shape) + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + x = x + attention_blocks[i](x) + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + out = norm_layer(x) + out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous() + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-vit-msl-s1/modeling_skysensepp_vit_msl.py b/skysensepp-vit-msl-s1/modeling_skysensepp_vit_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..6e56ac39c612c2eb87baf00b1e5afc071dc4d060 --- /dev/null +++ b/skysensepp-vit-msl-s1/modeling_skysensepp_vit_msl.py @@ -0,0 +1,265 @@ +"""SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace).""" + +import math +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig +from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + feedforward_channels: int, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + num_fcs: int = 2, + qkv_bias: bool = True, + with_cp: bool = False, + ): + super().__init__() + self.with_cp = with_cp + self.norm1 = nn.LayerNorm(embed_dims) + self.attn = nn.MultiheadAttention( + embed_dim=embed_dims, + num_heads=num_heads, + dropout=attn_drop_rate, + bias=qkv_bias, + batch_first=True, + ) + self.proj_drop = nn.Dropout(drop_rate) + self.norm2 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=num_fcs, + ffn_drop=drop_rate, + drop_path=drop_path_rate, + act_layer=nn.GELU, + add_identity=True, + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + def _inner_forward(x): + residual = x + x_norm = self.norm1(x) + attn_out, _ = self.attn(x_norm, x_norm, x_norm) + attn_out = self.proj_drop(attn_out) + x = residual + self.drop_path(attn_out) + return self.ffn(self.norm2(x), identity=x) + + if self.with_cp and x.requires_grad: + return cp.checkpoint(_inner_forward, x, use_reentrant=False) + return _inner_forward(x) + + +class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusViTMSLConfig + base_model_prefix = "skysensepp_vit_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel): + """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusViTMSLConfig): + super().__init__(config) + + img_size = to_2tuple(config.img_size) + self.img_size = img_size + self.patch_size = config.patch_size + self.with_cls_token = config.with_cls_token + self.output_cls_token = config.output_cls_token + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.interpolate_mode = "bicubic" + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm if config.patch_norm else None, + ) + + num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size) + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + out_indices = list(config.out_indices) + self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices] + + num_layers = config.num_layers + if num_layers > 1: + dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)] + else: + dpr = [0.0] + + self.layers = nn.ModuleList() + for i in range(config.num_layers): + self.layers.append( + TransformerEncoderLayer( + embed_dims=config.embed_dims, + num_heads=config.num_heads, + feedforward_channels=config.mlp_ratio * config.embed_dims, + attn_drop_rate=config.attn_drop_rate, + drop_rate=config.drop_rate, + drop_path_rate=dpr[i], + num_fcs=2, + qkv_bias=config.qkv_bias, + with_cp=config.with_cp, + ) + ) + + self.final_norm = config.final_norm + if config.final_norm: + self.norm = nn.LayerNorm(config.embed_dims) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.norm_attn = nn.LayerNorm(config.embed_dims) + + self.post_init() + + @staticmethod + def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"): + pos_h, pos_w = pos_shape + pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :] + pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) + pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode) + return torch.flatten(pos_embed_weight, 2).transpose(1, 2) + + def _pos_embedding(self, patched_img, hw_shape, pos_embed): + x_len, pos_len = patched_img.shape[1], pos_embed.shape[1] + if x_len != pos_len: + pos_h = self.img_size[0] // self.patch_size + pos_w = self.img_size[1] // self.patch_size + pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode) + return self.drop_after_pos(patched_img + pos_embed) + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + x = self._pos_embedding(x, hw_shape, self.pos_embed) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, layer in enumerate(self.layers): + x = layer(x) + + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + attn_out, _ = attention_blocks[i](x, x, x) + x = x + attn_out + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm: + x = self.norm(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + out = x + out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous() + if self.output_cls_token: + out = [out, x[:, 0]] + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-vit-msl-s1/modeling_utils.py b/skysensepp-vit-msl-s1/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..93feae77f3a3c46e65167f2ed312a25b7cd4ad3a --- /dev/null +++ b/skysensepp-vit-msl-s1/modeling_utils.py @@ -0,0 +1,557 @@ +"""SkySense: Pure PyTorch + HuggingFace Transformers implementation. + +Shared utility modules used across SkySense model implementations. +""" + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def to_2tuple(x): + """Convert to a 2-tuple.""" + if isinstance(x, (list, tuple)): + return tuple(x) + return (x, x) + + +class DropPath(nn.Module): + """Drop paths (stochastic depth) per sample. + + Args: + drop_prob (float): Probability of dropping a path. Default: 0.0. + """ + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor = torch.floor(random_tensor + keep_prob) + output = x / keep_prob * random_tensor + return output + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding using Conv2d. + + Args: + in_channels (int): Number of input channels. Default: 3. + embed_dims (int): Embedding dimension. Default: 96. + kernel_size (int): Kernel size of the projection. Default: 4. + stride (int): Stride of the projection. Default: 4. + padding (int): Padding of the projection. Default: 0. + norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm. + input_size (int or tuple or None): Input resolution for calculating output size. + """ + + def __init__( + self, + in_channels: int = 3, + embed_dims: int = 96, + kernel_size: int = 4, + stride: int = 4, + padding: int = 0, + norm_layer: Optional[type] = nn.LayerNorm, + input_size: Optional[int] = None, + ): + super().__init__() + self.projection = nn.Conv2d( + in_channels, embed_dims, + kernel_size=kernel_size, stride=stride, padding=padding, + ) + self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity() + + # Compute init output size if input_size is given + if input_size is not None: + input_size = to_2tuple(input_size) + self.init_out_size = ( + (input_size[0] - kernel_size + 2 * padding) // stride + 1, + (input_size[1] - kernel_size + 2 * padding) // stride + 1, + ) + else: + self.init_out_size = None + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]: + x = self.projection(x) # (B, C, H, W) + out_size = (x.shape[2], x.shape[3]) + x = x.flatten(2).transpose(1, 2) # (B, H*W, C) + x = self.norm(x) + return x, out_size + + +class FFN(nn.Module): + """Feed-Forward Network. + + Args: + embed_dims (int): Input dimension. + feedforward_channels (int): Hidden dimension. + num_fcs (int): Number of FC layers. Default: 2. + ffn_drop (float): Dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + act_layer (nn.Module): Activation layer class. Default: nn.GELU. + add_identity (bool): Whether to add identity connection. Default: True. + """ + + def __init__( + self, + embed_dims: int, + feedforward_channels: int, + num_fcs: int = 2, + ffn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: type = nn.GELU, + add_identity: bool = True, + ): + super().__init__() + assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}" + self.embed_dims = embed_dims + self.feedforward_channels = feedforward_channels + self.add_identity = add_identity + + layers = [] + in_channels = embed_dims + for i in range(num_fcs - 1): + layers.append(nn.Linear(in_channels, feedforward_channels)) + layers.append(act_layer()) + layers.append(nn.Dropout(ffn_drop)) + in_channels = feedforward_channels + layers.append(nn.Linear(feedforward_channels, embed_dims)) + layers.append(nn.Dropout(ffn_drop)) + self.layers = nn.Sequential(*layers) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor: + out = self.layers(x) + out = self.drop_path(out) + if self.add_identity: + if identity is None: + identity = x + out = out + identity + return out + + +class WindowMSAV2(nn.Module): + """Window-based Multi-head Self-Attention for Swin Transformer V2. + + Uses cosine attention and log-spaced continuous position bias (log-CPB). + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): Window size (Wh, Ww). + pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0). + qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Output projection dropout rate. Default: 0.0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: Tuple[int, int], + pretrained_window_size: Tuple[int, int] = (0, 0), + qkv_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ): + super().__init__() + self.embed_dims = embed_dims + self.num_heads = num_heads + self.window_size = window_size + self.pretrained_window_size = pretrained_window_size + + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1)))) + + # MLP for continuous relative position bias (log-CPB) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) + + # Build relative coords table + self._build_relative_coords_table() + # Build relative position index + self._build_relative_position_index() + + self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False) + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(embed_dims)) + self.v_bias = nn.Parameter(torch.zeros(embed_dims)) + else: + self.q_bias = None + self.v_bias = None + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(embed_dims, embed_dims) + self.proj_drop = nn.Dropout(proj_drop) + self.softmax = nn.Softmax(dim=-1) + + def _build_relative_coords_table(self): + """Build the relative coordinates table for log-CPB.""" + Wh, Ww = self.window_size + # Table of relative coordinates + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2) + + # Normalize to [-1, 1] and apply log-scale + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 # normalize to -8, 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + self.register_buffer("relative_coords_table", coords_table) + + def _build_relative_position_index(self): + """Build the pairwise relative position index for each window token.""" + Wh, Ww = self.window_size + coords_h = torch.arange(Wh) + coords_w = torch.arange(Ww) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij')) + coords_flatten = coords.view(2, -1) + + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # (2, Wh*Ww, Wh*Ww) + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += Wh - 1 + relative_coords[:, :, 1] += Ww - 1 + relative_coords[:, :, 0] *= 2 * Ww - 1 + relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww) + self.register_buffer("relative_position_index", relative_position_index) + + def _compute_position_bias(self, N): + """Compute relative position bias, supporting dynamic window sizes. + + The log-CPB (Continuous Position Bias) MLP can generalize to any window + size by computing bias from normalized relative coordinates. + """ + init_N = self.window_size[0] * self.window_size[1] + if N == init_N: + # Use pre-built tables + relative_position_bias_table = self.cpb_mlp( + self.relative_coords_table + ).view(-1, self.num_heads) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view(N, N, -1) + else: + # Dynamic: compute for actual window size on-the-fly + Wh = Ww = int(math.sqrt(N)) + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + # Build position index for actual window size + ch = torch.arange(Wh, device=self.logit_scale.device) + cw = torch.arange(Ww, device=self.logit_scale.device) + coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij')) + coords_flat = coords.view(2, -1) + rel = coords_flat[:, :, None] - coords_flat[:, None, :] + rel = rel.permute(1, 2, 0).contiguous() + rel[:, :, 0] += Wh - 1 + rel[:, :, 1] += Ww - 1 + rel[:, :, 0] *= 2 * Ww - 1 + pos_index = rel.sum(-1) + + bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads) + relative_position_bias = bias_table[ + pos_index.view(-1) + ].view(N, N, -1) + + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + return relative_position_bias + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: (num_windows*B, N, C) where N = Wh*Ww + mask: (num_windows, N, N) or None + """ + B_, N, C = x.shape + + # Compute QKV with bias + if self.q_bias is not None: + qkv_bias = torch.cat( + (self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias)) + qkv = F.linear(x, self.qkv.weight, qkv_bias) + else: + qkv = self.qkv(x) + + qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + + # Cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=math.log(1.0 / 0.01) + ).exp() + attn = attn * logit_scale + + # Log-CPB relative position bias (supports dynamic window sizes) + relative_position_bias = self._compute_position_bias(N) + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + attn = attn + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + + attn = self.softmax(attn) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class ShiftWindowMSA(nn.Module): + """Shifted Window Multi-head Self-Attention. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. Default: 0. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Projection dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + pad_small_map (bool): Pad small feature maps to window size. Default: False. + pretrained_window_size (int): Pretrained window size. Default: 0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int, + shift_size: int = 0, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.window_size = window_size + self.shift_size = shift_size + self.pad_small_map = pad_small_map + + self.w_msa = WindowMSAV2( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=to_2tuple(window_size), + pretrained_window_size=to_2tuple(pretrained_window_size), + attn_drop=attn_drop, + proj_drop=proj_drop, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W, f"Input length {L} != H*W ({H}*{W})" + + x = x.view(B, H, W, C) + + window_size = self.window_size + shift_size = self.shift_size + + # Pad or shrink window + if self.pad_small_map: + pad_r = (window_size - W % window_size) % window_size + pad_b = (window_size - H % window_size) % window_size + x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b)) + _, Hp, Wp, _ = x.shape + else: + Hp, Wp = H, W + if window_size > Hp: + window_size = Hp + shift_size = 0 + if window_size > Wp: + window_size = Wp + shift_size = 0 + + # Compute attention mask for SW-MSA + attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device) + + # Cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) + + # Partition windows + x_windows = self._window_partition(x, window_size) + # (num_windows*B, window_size*window_size, C) + + # W-MSA/SW-MSA + attn_windows = self.w_msa(x_windows, mask=attn_mask) + + # Merge windows + x = self._window_reverse(attn_windows, window_size, Hp, Wp) + + # Reverse cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) + + if self.pad_small_map and (pad_r > 0 or pad_b > 0): + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + x = self.drop_path(x) + return x + + @staticmethod + def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor: + """Partition into non-overlapping windows.""" + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous() + windows = windows.view(-1, window_size * window_size, C) + return windows + + @staticmethod + def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor: + """Reverse window partition.""" + B_nW = windows.shape[0] + nH = H // window_size + nW = W // window_size + B = B_nW // (nH * nW) + x = windows.view(B, nH, nW, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous() + x = x.view(B, H, W, -1) + return x + + @staticmethod + def _compute_attn_mask(H, W, window_size, shift_size, device): + """Compute attention mask for shifted window attention.""" + if shift_size <= 0: + return None + img_mask = torch.zeros((1, H, W, 1), device=device) + h_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + w_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + # Partition mask + mask_windows = img_mask.view( + 1, H // window_size, window_size, W // window_size, window_size, 1 + ) + mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous() + mask_windows = mask_windows.view(-1, window_size * window_size) + + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0) + attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0) + return attn_mask + + +class PatchMerging(nn.Module): + """Patch Merging Layer for downsampling (2x). + + Args: + in_channels (int): Input channels. + out_channels (int): Output channels. + norm_layer (type): Normalization layer. Default: nn.LayerNorm. + is_post_norm (bool): Apply norm after linear. Default: True. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + norm_layer: type = nn.LayerNorm, + is_post_norm: bool = True, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.is_post_norm = is_post_norm + self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False) + if is_post_norm: + self.norm = norm_layer(out_channels) + else: + self.norm = norm_layer(4 * in_channels) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W + + x = x.view(B, H, W, C) + + # Pad if needed + pad_h = H % 2 + pad_w = W % 2 + if pad_h or pad_w: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], dim=-1) + + out_h = (H + pad_h) // 2 + out_w = (W + pad_w) // 2 + x = x.view(B, out_h * out_w, 4 * C) + + if self.is_post_norm: + x = self.reduction(x) + x = self.norm(x) + else: + x = self.norm(x) + x = self.reduction(x) + + return x, (out_h, out_w) diff --git a/skysensepp-vit-msl-s1/pipeline_skysensepp.py b/skysensepp-vit-msl-s1/pipeline_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9a5354361c9a08abd1a82b3df8d4ec678a209b --- /dev/null +++ b/skysensepp-vit-msl-s1/pipeline_skysensepp.py @@ -0,0 +1,86 @@ +"""Custom HuggingFace pipeline for SkySense++ MSL feature extraction.""" + +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline): + """Pipeline for SkySense++ MSL backbones. + + Expects image tensors plus semantic annotation maps (class indices). + """ + + def _sanitize_parameters( + self, + annotation=None, + mask=None, + output_hidden_states=None, + **kwargs, + ): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + + if annotation is not None: + preprocess_params["annotation"] = annotation + if mask is not None: + forward_params["mask"] = mask + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + + return preprocess_params, forward_params, postprocess_params + + def preprocess( + self, + pixel_values: Any, + annotation: Optional[Any] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if isinstance(pixel_values, dict): + annotation = pixel_values.get("annotation", annotation) + pixel_values = pixel_values.get("pixel_values", pixel_values) + + if isinstance(pixel_values, np.ndarray): + pixel_values = torch.from_numpy(pixel_values).float() + elif not isinstance(pixel_values, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}" + ) + + if annotation is None: + raise ValueError("SkySense++ MSL models require an `annotation` semantic map.") + + if isinstance(annotation, np.ndarray): + annotation = torch.from_numpy(annotation).long() + elif not isinstance(annotation, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for annotation, got {type(annotation)}" + ) + + if pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + if annotation.ndim == 2: + annotation = annotation.unsqueeze(0) + + return {"pixel_values": pixel_values, "annotation": annotation} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + pixel_values=model_inputs["pixel_values"], + annotation=model_inputs["annotation"], + mask=kwargs.get("mask"), + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"last_hidden_state": outputs.last_hidden_state} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-vit-msl-s1/pipeline_skysensepp_fusion.py b/skysensepp-vit-msl-s1/pipeline_skysensepp_fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..80f070fa9fd3f18a053caa1afef9dc142bce2597 --- /dev/null +++ b/skysensepp-vit-msl-s1/pipeline_skysensepp_fusion.py @@ -0,0 +1,53 @@ +"""Optional pipeline for SkySense++ fusion neck.""" + +from typing import Any, Dict + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusFusionNeckPipeline(Pipeline): + """Pipeline for the optional SkySense++ fusion neck module. + + Expects concatenated multi-modal tokens per spatial location: + ``(batch, num_modalities, input_dims)``. + """ + + def _sanitize_parameters(self, output_hidden_states=None, **kwargs): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + return preprocess_params, forward_params, postprocess_params + + def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]: + if isinstance(hidden_states, dict): + hidden_states = hidden_states["hidden_states"] + + if isinstance(hidden_states, np.ndarray): + hidden_states = torch.from_numpy(hidden_states).float() + elif not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}" + ) + if hidden_states.ndim == 2: + hidden_states = hidden_states.unsqueeze(0) + return {"hidden_states": hidden_states} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + hidden_states=model_inputs["hidden_states"], + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"pooler_output": outputs.pooler_output} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-vit-msl-s2/__init__.py b/skysensepp-vit-msl-s2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9d49c51cc97fc8d74fa2f98f3c4677a5b59f8bc2 --- /dev/null +++ b/skysensepp-vit-msl-s2/__init__.py @@ -0,0 +1,25 @@ +"""SkySense++: Multi-Modal Remote Sensing Foundation Model (HuggingFace).""" + +from .configuration_skysensepp import ( + SkySensePlusPlusSwinV2MSLConfig, + SkySensePlusPlusViTMSLConfig, +) +from .modeling_skysensepp_swinv2_msl import ( + SkySensePlusPlusSwinV2MSLModel, + SkySensePlusPlusSwinV2MSLPreTrainedModel, +) +from .modeling_skysensepp_vit_msl import ( + SkySensePlusPlusViTMSLModel, + SkySensePlusPlusViTMSLPreTrainedModel, +) +from .pipeline_skysensepp import SkySensePlusPlusMSLFeatureExtractionPipeline + +__all__ = [ + "SkySensePlusPlusSwinV2MSLConfig", + "SkySensePlusPlusViTMSLConfig", + "SkySensePlusPlusSwinV2MSLModel", + "SkySensePlusPlusSwinV2MSLPreTrainedModel", + "SkySensePlusPlusViTMSLModel", + "SkySensePlusPlusViTMSLPreTrainedModel", + "SkySensePlusPlusMSLFeatureExtractionPipeline", +] diff --git a/skysensepp-vit-msl-s2/config.json b/skysensepp-vit-msl-s2/config.json new file mode 100644 index 0000000000000000000000000000000000000000..2bdc97fa29535d2c178d651921c9ce6972d45dcf --- /dev/null +++ b/skysensepp-vit-msl-s2/config.json @@ -0,0 +1,68 @@ +{ + "return_dict": true, + "output_hidden_states": false, + "dtype": "float32", + "chunk_size_feed_forward": 0, + "is_encoder_decoder": false, + "architectures": [ + "SkySensePlusPlusViTMSLModel" + ], + "id2label": { + "0": "LABEL_0", + "1": "LABEL_1" + }, + "label2id": { + "LABEL_0": 0, + "LABEL_1": 1 + }, + "problem_type": null, + "_name_or_path": "", + "transformers_version": "5.0.0", + "img_size": 16, + "patch_size": 4, + "in_channels": 10, + "embed_dims": 1024, + "num_layers": 24, + "num_heads": 16, + "mlp_ratio": 4, + "out_indices": [ + 5, + 11, + 17, + 23 + ], + "qkv_bias": true, + "drop_rate": 0.0, + "attn_drop_rate": 0.0, + "drop_path_rate": 0.3, + "with_cls_token": false, + "output_cls_token": false, + "patch_norm": false, + "final_norm": false, + "with_cp": false, + "vocabulary_size": 64, + "num_vocabulary_tokens": 65, + "merge_stage": 4, + "use_attn": false, + "modality": "s2", + "model_type": "skysensepp_vit_msl", + "output_attentions": false, + "auto_map": { + "AutoConfig": "configuration_skysensepp.SkySensePlusPlusViTMSLConfig", + "AutoModel": "modeling_skysensepp_vit_msl.SkySensePlusPlusViTMSLModel" + }, + "custom_pipelines": { + "skysensepp-feature-extraction": { + "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + }, + "image-feature-extraction": { + "impl": "pipeline_skysensepp.SkySensePlusPlusMSLFeatureExtractionPipeline", + "pt": [ + "AutoModel" + ] + } + } +} diff --git a/skysensepp-vit-msl-s2/configuration_skysensepp.py b/skysensepp-vit-msl-s2/configuration_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..e086dd1de7d378e81cc124a89a36e33f45bc4b3f --- /dev/null +++ b/skysensepp-vit-msl-s2/configuration_skysensepp.py @@ -0,0 +1,124 @@ +"""Configuration classes for SkySense++ MSL backbones.""" + +from transformers import PretrainedConfig + + +class SkySensePlusPlusSwinV2MSLConfig(PretrainedConfig): + """Configuration for SkySense++ Swin Transformer V2 MSL backbone (HR optical).""" + + model_type = "skysensepp_swinv2_msl" + + arch_zoo = { + "tiny": {"embed_dims": 96, "depths": [2, 2, 6, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "small": {"embed_dims": 96, "depths": [2, 2, 18, 2], "num_heads": [3, 6, 12, 24], "extra_norm_every_n_blocks": 0}, + "base": {"embed_dims": 128, "depths": [2, 2, 18, 2], "num_heads": [4, 8, 16, 32], "extra_norm_every_n_blocks": 0}, + "large": {"embed_dims": 192, "depths": [2, 2, 18, 2], "num_heads": [6, 12, 24, 48], "extra_norm_every_n_blocks": 0}, + "huge": {"embed_dims": 352, "depths": [2, 2, 18, 2], "num_heads": [8, 16, 32, 64], "extra_norm_every_n_blocks": 6}, + "giant": {"embed_dims": 512, "depths": [2, 2, 42, 4], "num_heads": [16, 32, 64, 128], "extra_norm_every_n_blocks": 6}, + } + + def __init__( + self, + arch="huge", + img_size=224, + patch_size=4, + in_channels=3, + window_size=8, + drop_rate=0.0, + drop_path_rate=0.2, + out_indices=(0, 1, 2, 3), + use_abs_pos_embed=False, + with_cp=False, + pad_small_map=False, + pretrained_window_sizes=(0, 0, 0, 0), + is_post_norm_downsample=True, + vocabulary_size=64, + merge_stage=2, + use_attn=True, + **kwargs, + ): + super().__init__(**kwargs) + + arch = arch.lower() + if arch not in self.arch_zoo: + raise ValueError(f"Unknown arch '{arch}'. Choose from {list(self.arch_zoo.keys())}") + arch_settings = self.arch_zoo[arch] + + self.arch = arch + self.embed_dims = arch_settings["embed_dims"] + self.depths = arch_settings["depths"] + self.num_heads = arch_settings["num_heads"] + self.extra_norm_every_n_blocks = arch_settings["extra_norm_every_n_blocks"] + + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.window_size = window_size + self.drop_rate = drop_rate + self.drop_path_rate = drop_path_rate + self.out_indices = list(out_indices) + self.use_abs_pos_embed = use_abs_pos_embed + self.with_cp = with_cp + self.pad_small_map = pad_small_map + self.pretrained_window_sizes = list(pretrained_window_sizes) + self.is_post_norm_downsample = is_post_norm_downsample + + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + + +class SkySensePlusPlusViTMSLConfig(PretrainedConfig): + """Configuration for SkySense++ Vision Transformer MSL backbone (S2/S1).""" + + model_type = "skysensepp_vit_msl" + + def __init__( + self, + img_size=16, + patch_size=4, + in_channels=10, + embed_dims=1024, + num_layers=24, + num_heads=16, + mlp_ratio=4, + out_indices=(5, 11, 17, 23), + qkv_bias=True, + drop_rate=0.0, + attn_drop_rate=0.0, + drop_path_rate=0.3, + with_cls_token=False, + output_cls_token=False, + patch_norm=False, + final_norm=False, + with_cp=False, + vocabulary_size=64, + merge_stage=4, + use_attn=False, + modality="s2", + **kwargs, + ): + super().__init__(**kwargs) + self.img_size = img_size + self.patch_size = patch_size + self.in_channels = in_channels + self.embed_dims = embed_dims + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_ratio = mlp_ratio + self.out_indices = list(out_indices) + self.qkv_bias = qkv_bias + self.drop_rate = drop_rate + self.attn_drop_rate = attn_drop_rate + self.drop_path_rate = drop_path_rate + self.with_cls_token = with_cls_token + self.output_cls_token = output_cls_token + self.patch_norm = patch_norm + self.final_norm = final_norm + self.with_cp = with_cp + self.vocabulary_size = vocabulary_size + self.num_vocabulary_tokens = vocabulary_size + 1 + self.merge_stage = merge_stage + self.use_attn = use_attn + self.modality = modality diff --git a/skysensepp-vit-msl-s2/conversion_manifest.json b/skysensepp-vit-msl-s2/conversion_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..cd2c6752cbcdcd2b625a07b6fb71ec355143539a --- /dev/null +++ b/skysensepp-vit-msl-s2/conversion_manifest.json @@ -0,0 +1,305 @@ +{ + "source_checkpoint": "/exstorage/czy/models/raw/skysensepp_release_s2.pth", + "modality": "s2", + "model_class": "SkySensePlusPlusViTMSLModel", + "num_tensors": 295, + "missing_keys": [], + "unexpected_keys": [], + "tensor_names": [ + "cls_token", + "layers.0.attn.in_proj_bias", + "layers.0.attn.in_proj_weight", + "layers.0.attn.out_proj.bias", + "layers.0.attn.out_proj.weight", + "layers.0.ffn.layers.0.bias", + "layers.0.ffn.layers.0.weight", + "layers.0.ffn.layers.3.bias", + "layers.0.ffn.layers.3.weight", + "layers.0.norm1.bias", + "layers.0.norm1.weight", + "layers.0.norm2.bias", + "layers.0.norm2.weight", + "layers.1.attn.in_proj_bias", + "layers.1.attn.in_proj_weight", + "layers.1.attn.out_proj.bias", + "layers.1.attn.out_proj.weight", + "layers.1.ffn.layers.0.bias", + "layers.1.ffn.layers.0.weight", + "layers.1.ffn.layers.3.bias", + "layers.1.ffn.layers.3.weight", + "layers.1.norm1.bias", + "layers.1.norm1.weight", + "layers.1.norm2.bias", + "layers.1.norm2.weight", + "layers.10.attn.in_proj_bias", + "layers.10.attn.in_proj_weight", + "layers.10.attn.out_proj.bias", + "layers.10.attn.out_proj.weight", + "layers.10.ffn.layers.0.bias", + "layers.10.ffn.layers.0.weight", + "layers.10.ffn.layers.3.bias", + "layers.10.ffn.layers.3.weight", + "layers.10.norm1.bias", + "layers.10.norm1.weight", + "layers.10.norm2.bias", + "layers.10.norm2.weight", + "layers.11.attn.in_proj_bias", + "layers.11.attn.in_proj_weight", + "layers.11.attn.out_proj.bias", + "layers.11.attn.out_proj.weight", + "layers.11.ffn.layers.0.bias", + "layers.11.ffn.layers.0.weight", + "layers.11.ffn.layers.3.bias", + "layers.11.ffn.layers.3.weight", + "layers.11.norm1.bias", + "layers.11.norm1.weight", + "layers.11.norm2.bias", + "layers.11.norm2.weight", + "layers.12.attn.in_proj_bias", + "layers.12.attn.in_proj_weight", + "layers.12.attn.out_proj.bias", + "layers.12.attn.out_proj.weight", + "layers.12.ffn.layers.0.bias", + "layers.12.ffn.layers.0.weight", + "layers.12.ffn.layers.3.bias", + "layers.12.ffn.layers.3.weight", + "layers.12.norm1.bias", + "layers.12.norm1.weight", + "layers.12.norm2.bias", + "layers.12.norm2.weight", + "layers.13.attn.in_proj_bias", + "layers.13.attn.in_proj_weight", + "layers.13.attn.out_proj.bias", + "layers.13.attn.out_proj.weight", + "layers.13.ffn.layers.0.bias", + "layers.13.ffn.layers.0.weight", + "layers.13.ffn.layers.3.bias", + "layers.13.ffn.layers.3.weight", + "layers.13.norm1.bias", + "layers.13.norm1.weight", + "layers.13.norm2.bias", + "layers.13.norm2.weight", + "layers.14.attn.in_proj_bias", + "layers.14.attn.in_proj_weight", + "layers.14.attn.out_proj.bias", + "layers.14.attn.out_proj.weight", + "layers.14.ffn.layers.0.bias", + "layers.14.ffn.layers.0.weight", + "layers.14.ffn.layers.3.bias", + "layers.14.ffn.layers.3.weight", + "layers.14.norm1.bias", + "layers.14.norm1.weight", + "layers.14.norm2.bias", + "layers.14.norm2.weight", + "layers.15.attn.in_proj_bias", + "layers.15.attn.in_proj_weight", + "layers.15.attn.out_proj.bias", + "layers.15.attn.out_proj.weight", + "layers.15.ffn.layers.0.bias", + "layers.15.ffn.layers.0.weight", + "layers.15.ffn.layers.3.bias", + "layers.15.ffn.layers.3.weight", + "layers.15.norm1.bias", + "layers.15.norm1.weight", + "layers.15.norm2.bias", + "layers.15.norm2.weight", + "layers.16.attn.in_proj_bias", + "layers.16.attn.in_proj_weight", + "layers.16.attn.out_proj.bias", + "layers.16.attn.out_proj.weight", + "layers.16.ffn.layers.0.bias", + "layers.16.ffn.layers.0.weight", + "layers.16.ffn.layers.3.bias", + "layers.16.ffn.layers.3.weight", + "layers.16.norm1.bias", + "layers.16.norm1.weight", + "layers.16.norm2.bias", + "layers.16.norm2.weight", + "layers.17.attn.in_proj_bias", + "layers.17.attn.in_proj_weight", + "layers.17.attn.out_proj.bias", + "layers.17.attn.out_proj.weight", + "layers.17.ffn.layers.0.bias", + "layers.17.ffn.layers.0.weight", + "layers.17.ffn.layers.3.bias", + "layers.17.ffn.layers.3.weight", + "layers.17.norm1.bias", + "layers.17.norm1.weight", + "layers.17.norm2.bias", + "layers.17.norm2.weight", + "layers.18.attn.in_proj_bias", + "layers.18.attn.in_proj_weight", + "layers.18.attn.out_proj.bias", + "layers.18.attn.out_proj.weight", + "layers.18.ffn.layers.0.bias", + "layers.18.ffn.layers.0.weight", + "layers.18.ffn.layers.3.bias", + "layers.18.ffn.layers.3.weight", + "layers.18.norm1.bias", + "layers.18.norm1.weight", + "layers.18.norm2.bias", + "layers.18.norm2.weight", + "layers.19.attn.in_proj_bias", + "layers.19.attn.in_proj_weight", + "layers.19.attn.out_proj.bias", + "layers.19.attn.out_proj.weight", + "layers.19.ffn.layers.0.bias", + "layers.19.ffn.layers.0.weight", + "layers.19.ffn.layers.3.bias", + "layers.19.ffn.layers.3.weight", + "layers.19.norm1.bias", + "layers.19.norm1.weight", + "layers.19.norm2.bias", + "layers.19.norm2.weight", + "layers.2.attn.in_proj_bias", + "layers.2.attn.in_proj_weight", + "layers.2.attn.out_proj.bias", + "layers.2.attn.out_proj.weight", + "layers.2.ffn.layers.0.bias", + "layers.2.ffn.layers.0.weight", + "layers.2.ffn.layers.3.bias", + "layers.2.ffn.layers.3.weight", + "layers.2.norm1.bias", + "layers.2.norm1.weight", + "layers.2.norm2.bias", + "layers.2.norm2.weight", + "layers.20.attn.in_proj_bias", + "layers.20.attn.in_proj_weight", + "layers.20.attn.out_proj.bias", + "layers.20.attn.out_proj.weight", + "layers.20.ffn.layers.0.bias", + "layers.20.ffn.layers.0.weight", + "layers.20.ffn.layers.3.bias", + "layers.20.ffn.layers.3.weight", + "layers.20.norm1.bias", + "layers.20.norm1.weight", + "layers.20.norm2.bias", + "layers.20.norm2.weight", + "layers.21.attn.in_proj_bias", + "layers.21.attn.in_proj_weight", + "layers.21.attn.out_proj.bias", + "layers.21.attn.out_proj.weight", + "layers.21.ffn.layers.0.bias", + "layers.21.ffn.layers.0.weight", + "layers.21.ffn.layers.3.bias", + "layers.21.ffn.layers.3.weight", + "layers.21.norm1.bias", + "layers.21.norm1.weight", + "layers.21.norm2.bias", + "layers.21.norm2.weight", + "layers.22.attn.in_proj_bias", + "layers.22.attn.in_proj_weight", + "layers.22.attn.out_proj.bias", + "layers.22.attn.out_proj.weight", + "layers.22.ffn.layers.0.bias", + "layers.22.ffn.layers.0.weight", + "layers.22.ffn.layers.3.bias", + "layers.22.ffn.layers.3.weight", + "layers.22.norm1.bias", + "layers.22.norm1.weight", + "layers.22.norm2.bias", + "layers.22.norm2.weight", + "layers.23.attn.in_proj_bias", + "layers.23.attn.in_proj_weight", + "layers.23.attn.out_proj.bias", + "layers.23.attn.out_proj.weight", + "layers.23.ffn.layers.0.bias", + "layers.23.ffn.layers.0.weight", + "layers.23.ffn.layers.3.bias", + "layers.23.ffn.layers.3.weight", + "layers.23.norm1.bias", + "layers.23.norm1.weight", + "layers.23.norm2.bias", + "layers.23.norm2.weight", + "layers.3.attn.in_proj_bias", + "layers.3.attn.in_proj_weight", + "layers.3.attn.out_proj.bias", + "layers.3.attn.out_proj.weight", + "layers.3.ffn.layers.0.bias", + "layers.3.ffn.layers.0.weight", + "layers.3.ffn.layers.3.bias", + "layers.3.ffn.layers.3.weight", + "layers.3.norm1.bias", + "layers.3.norm1.weight", + "layers.3.norm2.bias", + "layers.3.norm2.weight", + "layers.4.attn.in_proj_bias", + "layers.4.attn.in_proj_weight", + "layers.4.attn.out_proj.bias", + "layers.4.attn.out_proj.weight", + "layers.4.ffn.layers.0.bias", + "layers.4.ffn.layers.0.weight", + "layers.4.ffn.layers.3.bias", + "layers.4.ffn.layers.3.weight", + "layers.4.norm1.bias", + "layers.4.norm1.weight", + "layers.4.norm2.bias", + "layers.4.norm2.weight", + "layers.5.attn.in_proj_bias", + "layers.5.attn.in_proj_weight", + "layers.5.attn.out_proj.bias", + "layers.5.attn.out_proj.weight", + "layers.5.ffn.layers.0.bias", + "layers.5.ffn.layers.0.weight", + "layers.5.ffn.layers.3.bias", + "layers.5.ffn.layers.3.weight", + "layers.5.norm1.bias", + "layers.5.norm1.weight", + "layers.5.norm2.bias", + "layers.5.norm2.weight", + "layers.6.attn.in_proj_bias", + "layers.6.attn.in_proj_weight", + "layers.6.attn.out_proj.bias", + "layers.6.attn.out_proj.weight", + "layers.6.ffn.layers.0.bias", + "layers.6.ffn.layers.0.weight", + "layers.6.ffn.layers.3.bias", + "layers.6.ffn.layers.3.weight", + "layers.6.norm1.bias", + "layers.6.norm1.weight", + "layers.6.norm2.bias", + "layers.6.norm2.weight", + "layers.7.attn.in_proj_bias", + "layers.7.attn.in_proj_weight", + "layers.7.attn.out_proj.bias", + "layers.7.attn.out_proj.weight", + "layers.7.ffn.layers.0.bias", + "layers.7.ffn.layers.0.weight", + "layers.7.ffn.layers.3.bias", + "layers.7.ffn.layers.3.weight", + "layers.7.norm1.bias", + "layers.7.norm1.weight", + "layers.7.norm2.bias", + "layers.7.norm2.weight", + "layers.8.attn.in_proj_bias", + "layers.8.attn.in_proj_weight", + "layers.8.attn.out_proj.bias", + "layers.8.attn.out_proj.weight", + "layers.8.ffn.layers.0.bias", + "layers.8.ffn.layers.0.weight", + "layers.8.ffn.layers.3.bias", + "layers.8.ffn.layers.3.weight", + "layers.8.norm1.bias", + "layers.8.norm1.weight", + "layers.8.norm2.bias", + "layers.8.norm2.weight", + "layers.9.attn.in_proj_bias", + "layers.9.attn.in_proj_weight", + "layers.9.attn.out_proj.bias", + "layers.9.attn.out_proj.weight", + "layers.9.ffn.layers.0.bias", + "layers.9.ffn.layers.0.weight", + "layers.9.ffn.layers.3.bias", + "layers.9.ffn.layers.3.weight", + "layers.9.norm1.bias", + "layers.9.norm1.weight", + "layers.9.norm2.bias", + "layers.9.norm2.weight", + "mask_token", + "patch_embed.projection.bias", + "patch_embed.projection.weight", + "pos_embed", + "vocabulary_token", + "vocabulary_weight" + ] +} diff --git a/skysensepp-vit-msl-s2/model.safetensors b/skysensepp-vit-msl-s2/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..b8e87ca7a84d5a8de1af53424d8fc6e2815d003b --- /dev/null +++ b/skysensepp-vit-msl-s2/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:27f2a40bdad5ffc10598d808ec7b5d481f9848c6d6a63155fb4f4f113480c3e5 +size 1210265976 diff --git a/skysensepp-vit-msl-s2/modeling_skysensepp_swinv2_msl.py b/skysensepp-vit-msl-s2/modeling_skysensepp_swinv2_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..2c0ab382892dfc61d579821983bc8f6b68b84d17 --- /dev/null +++ b/skysensepp-vit-msl-s2/modeling_skysensepp_swinv2_msl.py @@ -0,0 +1,343 @@ +"""SkySense++ Swin Transformer V2 MSL backbone (pure PyTorch + HuggingFace).""" + +from copy import deepcopy +from typing import Optional, Sequence, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusSwinV2MSLConfig +from .modeling_utils import ( + DropPath, + FFN, + PatchEmbed, + PatchMerging, + ShiftWindowMSA, + to_2tuple, +) + + +class SwinBlockV2(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int = 8, + shift: bool = False, + extra_norm: bool = False, + ffn_ratio: float = 4.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + with_cp: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.with_cp = with_cp + self.extra_norm = extra_norm + self.attn = ShiftWindowMSA( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=window_size, + shift_size=window_size // 2 if shift else 0, + drop_path=drop_path, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + self.norm1 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=int(embed_dims * ffn_ratio), + num_fcs=2, + drop_path=drop_path, + act_layer=nn.GELU, + add_identity=False, + ) + self.norm2 = nn.LayerNorm(embed_dims) + if self.extra_norm: + self.norm3 = nn.LayerNorm(embed_dims) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + def _inner_forward(x): + identity = x + x = self.attn(x, hw_shape) + x = self.norm1(x) + x = x + identity + + identity = x + x = self.ffn(x) + x = self.norm2(x) + x = x + identity + + if self.extra_norm: + x = self.norm3(x) + return x + + if self.with_cp and x.requires_grad: + x = cp.checkpoint(_inner_forward, x, use_reentrant=False) + else: + x = _inner_forward(x) + return x + + +class SwinBlockV2Sequence(nn.Module): + def __init__( + self, + embed_dims: int, + depth: int, + num_heads: int, + window_size: int = 8, + downsample: bool = False, + drop_paths: Union[Sequence[float], float] = 0.0, + with_cp: bool = False, + pad_small_map: bool = False, + extra_norm_every_n_blocks: int = 0, + pretrained_window_size: int = 0, + is_post_norm_downsample: bool = True, + ): + super().__init__() + if not isinstance(drop_paths, Sequence): + drop_paths = [drop_paths] * depth + + if downsample: + self.out_channels = 2 * embed_dims + self.downsample = PatchMerging( + in_channels=embed_dims, + out_channels=self.out_channels, + is_post_norm=is_post_norm_downsample, + ) + else: + self.out_channels = embed_dims + self.downsample = None + + self.blocks = nn.ModuleList() + for i in range(depth): + extra_norm = extra_norm_every_n_blocks > 0 and (i + 1) % extra_norm_every_n_blocks == 0 + self.blocks.append( + SwinBlockV2( + embed_dims=self.out_channels, + num_heads=num_heads, + window_size=window_size, + shift=(i % 2 == 1), + extra_norm=extra_norm, + drop_path=drop_paths[i], + with_cp=with_cp, + pad_small_map=pad_small_map, + pretrained_window_size=pretrained_window_size, + ) + ) + + def forward(self, x: torch.Tensor, in_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + if self.downsample is not None: + x, out_shape = self.downsample(x, in_shape) + else: + out_shape = in_shape + + for block in self.blocks: + x = block(x, out_shape) + return x, out_shape + + +class ProjMHSA(nn.Module): + """Projected multi-head self-attention used in SkySense++ HR backbone.""" + + def __init__(self, embed_dims: int, proj_dims: int, num_heads: int = 16, bias: bool = True): + super().__init__() + self.proj_in = nn.Linear(embed_dims, proj_dims) + self.attn = nn.MultiheadAttention(proj_dims, num_heads, batch_first=True, bias=bias) + self.proj_out = nn.Linear(proj_dims, embed_dims) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.proj_in(x) + x, _ = self.attn(x, x, x) + return self.proj_out(x) + + +class SkySensePlusPlusSwinV2MSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusSwinV2MSLConfig + base_model_prefix = "skysensepp_swinv2_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.LayerNorm): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusSwinV2MSLModel(SkySensePlusPlusSwinV2MSLPreTrainedModel): + """SkySense++ HR backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusSwinV2MSLConfig): + super().__init__(config) + + self.num_layers = len(config.depths) + self.out_indices = config.out_indices + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.patch_size = config.patch_size + + if isinstance(config.window_size, int): + window_sizes = [config.window_size] * self.num_layers + else: + window_sizes = list(config.window_size) + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm, + input_size=config.img_size, + ) + + self.use_abs_pos_embed = config.use_abs_pos_embed + if self.use_abs_pos_embed: + patch_resolution = self.patch_embed.init_out_size + num_patches = patch_resolution[0] * patch_resolution[1] + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + total_depth = sum(config.depths) + if total_depth > 1: + dpr = [config.drop_path_rate * i / (total_depth - 1) for i in range(total_depth)] + else: + dpr = [0.0] + + self.stages = nn.ModuleList() + embed_dims_list = [config.embed_dims] + for i, (depth, num_heads) in enumerate(zip(config.depths, config.num_heads)): + stage = SwinBlockV2Sequence( + embed_dims=embed_dims_list[-1], + depth=depth, + num_heads=num_heads, + window_size=window_sizes[i], + downsample=(i > 0), + drop_paths=dpr[:depth], + with_cp=config.with_cp, + pad_small_map=config.pad_small_map, + extra_norm_every_n_blocks=config.extra_norm_every_n_blocks, + pretrained_window_size=config.pretrained_window_sizes[i], + is_post_norm_downsample=config.is_post_norm_downsample, + ) + self.stages.append(stage) + dpr = dpr[depth:] + embed_dims_list.append(stage.out_channels) + + for i in self.out_indices: + self.add_module(f"norm{i}", nn.LayerNorm(embed_dims_list[i + 1])) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = ProjMHSA(352, 256, num_heads=16) + self.attn2 = ProjMHSA(704, 512, num_heads=16) + self.attn3 = ProjMHSA(1408, 1024, num_heads=16) + self.norm_attn = nn.LayerNorm(1408) + + self.post_init() + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + if self.use_abs_pos_embed: + x = x + self.absolute_pos_embed + x = self.drop_after_pos(x) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, stage in enumerate(self.stages): + x, hw_shape = stage(x, hw_shape) + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + x = x + attention_blocks[i](x) + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + norm_layer = getattr(self, f"norm{i}") + out = norm_layer(x) + out = out.view(-1, *hw_shape, stage.out_channels).permute(0, 3, 1, 2).contiguous() + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-vit-msl-s2/modeling_skysensepp_vit_msl.py b/skysensepp-vit-msl-s2/modeling_skysensepp_vit_msl.py new file mode 100644 index 0000000000000000000000000000000000000000..6e56ac39c612c2eb87baf00b1e5afc071dc4d060 --- /dev/null +++ b/skysensepp-vit-msl-s2/modeling_skysensepp_vit_msl.py @@ -0,0 +1,265 @@ +"""SkySense++ Vision Transformer MSL backbone (pure PyTorch + HuggingFace).""" + +import math +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from transformers import PreTrainedModel +from transformers.modeling_outputs import BaseModelOutput + +from .configuration_skysensepp import SkySensePlusPlusViTMSLConfig +from .modeling_utils import DropPath, FFN, PatchEmbed, to_2tuple + + +class TransformerEncoderLayer(nn.Module): + def __init__( + self, + embed_dims: int, + num_heads: int, + feedforward_channels: int, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + num_fcs: int = 2, + qkv_bias: bool = True, + with_cp: bool = False, + ): + super().__init__() + self.with_cp = with_cp + self.norm1 = nn.LayerNorm(embed_dims) + self.attn = nn.MultiheadAttention( + embed_dim=embed_dims, + num_heads=num_heads, + dropout=attn_drop_rate, + bias=qkv_bias, + batch_first=True, + ) + self.proj_drop = nn.Dropout(drop_rate) + self.norm2 = nn.LayerNorm(embed_dims) + self.ffn = FFN( + embed_dims=embed_dims, + feedforward_channels=feedforward_channels, + num_fcs=num_fcs, + ffn_drop=drop_rate, + drop_path=drop_path_rate, + act_layer=nn.GELU, + add_identity=True, + ) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + def _inner_forward(x): + residual = x + x_norm = self.norm1(x) + attn_out, _ = self.attn(x_norm, x_norm, x_norm) + attn_out = self.proj_drop(attn_out) + x = residual + self.drop_path(attn_out) + return self.ffn(self.norm2(x), identity=x) + + if self.with_cp and x.requires_grad: + return cp.checkpoint(_inner_forward, x, use_reentrant=False) + return _inner_forward(x) + + +class SkySensePlusPlusViTMSLPreTrainedModel(PreTrainedModel): + config_class = SkySensePlusPlusViTMSLConfig + base_model_prefix = "skysensepp_vit_msl" + supports_gradient_checkpointing = True + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + nn.init.trunc_normal_(module.weight, std=0.02) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + nn.init.ones_(module.weight) + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Conv2d): + nn.init.kaiming_normal_(module.weight, mode="fan_in") + if module.bias is not None: + nn.init.zeros_(module.bias) + + +class SkySensePlusPlusViTMSLModel(SkySensePlusPlusViTMSLPreTrainedModel): + """SkySense++ S2/S1 backbone with semantic vocabulary and annotation conditioning.""" + + def __init__(self, config: SkySensePlusPlusViTMSLConfig): + super().__init__(config) + + img_size = to_2tuple(config.img_size) + self.img_size = img_size + self.patch_size = config.patch_size + self.with_cls_token = config.with_cls_token + self.output_cls_token = config.output_cls_token + self.merge_stage = config.merge_stage + self.use_attn = config.use_attn + self.interpolate_mode = "bicubic" + + self.patch_embed = PatchEmbed( + in_channels=config.in_channels, + embed_dims=config.embed_dims, + kernel_size=config.patch_size, + stride=config.patch_size, + norm_layer=nn.LayerNorm if config.patch_norm else None, + ) + + num_patches = (img_size[0] // config.patch_size) * (img_size[1] // config.patch_size) + self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, config.embed_dims)) + self.drop_after_pos = nn.Dropout(p=config.drop_rate) + + out_indices = list(config.out_indices) + self.out_indices = [idx if idx >= 0 else config.num_layers + idx for idx in out_indices] + + num_layers = config.num_layers + if num_layers > 1: + dpr = [config.drop_path_rate * i / (num_layers - 1) for i in range(num_layers)] + else: + dpr = [0.0] + + self.layers = nn.ModuleList() + for i in range(config.num_layers): + self.layers.append( + TransformerEncoderLayer( + embed_dims=config.embed_dims, + num_heads=config.num_heads, + feedforward_channels=config.mlp_ratio * config.embed_dims, + attn_drop_rate=config.attn_drop_rate, + drop_rate=config.drop_rate, + drop_path_rate=dpr[i], + num_fcs=2, + qkv_bias=config.qkv_bias, + with_cp=config.with_cp, + ) + ) + + self.final_norm = config.final_norm + if config.final_norm: + self.norm = nn.LayerNorm(config.embed_dims) + + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dims)) + self.vocabulary_token = nn.Parameter( + torch.zeros(config.num_vocabulary_tokens, config.embed_dims) + ) + self.vocabulary_weight = nn.Parameter(torch.zeros(1, config.patch_size * config.patch_size)) + + if self.use_attn: + self.attn1 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn2 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.attn3 = nn.MultiheadAttention(config.embed_dims, config.num_heads, batch_first=True, bias=True) + self.norm_attn = nn.LayerNorm(config.embed_dims) + + self.post_init() + + @staticmethod + def resize_pos_embed(pos_embed, input_shape, pos_shape, mode="bicubic"): + pos_h, pos_w = pos_shape + pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w) :] + pos_embed_weight = pos_embed_weight.reshape(1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2) + pos_embed_weight = F.interpolate(pos_embed_weight, size=input_shape, align_corners=False, mode=mode) + return torch.flatten(pos_embed_weight, 2).transpose(1, 2) + + def _pos_embedding(self, patched_img, hw_shape, pos_embed): + x_len, pos_len = patched_img.shape[1], pos_embed.shape[1] + if x_len != pos_len: + pos_h = self.img_size[0] // self.patch_size + pos_w = self.img_size[1] // self.patch_size + pos_embed = self.resize_pos_embed(pos_embed, hw_shape, (pos_h, pos_w), self.interpolate_mode) + return self.drop_after_pos(patched_img + pos_embed) + + def create_ann_token(self, anno_img: torch.Tensor) -> torch.Tensor: + batch_size, height, width = anno_img.shape + ann_token = torch.index_select( + self.vocabulary_token, 0, anno_img.reshape(-1) + ).reshape(batch_size, height, width, -1) + + num_patch_h = height // self.patch_size + num_patch_w = width // self.patch_size + weight = F.softmax(self.vocabulary_weight, dim=1) * self.patch_size * self.patch_size + weight = ( + weight.reshape(1, 1, self.patch_size, 1, self.patch_size) + .repeat(1, num_patch_h, 1, num_patch_w, 1) + .reshape(1, height, width, 1) + ) + ann_token = ann_token * weight + ann_token = F.avg_pool2d( + torch.einsum("bhwc->bchw", ann_token), self.patch_size, self.patch_size + ) + return torch.einsum("bchw->bhwc", ann_token).reshape( + batch_size, num_patch_h * num_patch_w, self.config.embed_dims + ) + + def forward( + self, + pixel_values: torch.Tensor, + annotation: torch.Tensor, + mask: Optional[torch.Tensor] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutput]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + x, hw_shape = self.patch_embed(pixel_values) + y = self.create_ann_token(annotation) + batch_size, num_tokens, channels = y.shape + + if mask is not None: + mask_tokens = self.mask_token.expand(batch_size, num_tokens, -1) + weight = mask.flatten(1).unsqueeze(-1).type_as(mask_tokens) + y = y * (1.0 - weight) + mask_tokens * weight + + if self.merge_stage == 0: + x = (x + y) * 0.5 + else: + x = x.reshape(batch_size, *hw_shape, channels) + y = y.reshape(batch_size, *hw_shape, channels) + x = torch.cat((x, y), dim=2) + hw_shape = (hw_shape[0], hw_shape[1] * 2) + x = x.reshape(batch_size, -1, channels) + + x = self._pos_embedding(x, hw_shape, self.pos_embed) + + all_hidden_states = () if output_hidden_states else None + feature_maps = [] + merge_idx = self.merge_stage - 1 + + for i, layer in enumerate(self.layers): + x = layer(x) + + if i == merge_idx: + x = x.reshape(batch_size, *hw_shape, x.shape[-1]) + x = (x[:, :, : x.shape[2] // 2] + x[:, :, x.shape[2] // 2 :]) * 0.5 + x = x.reshape(batch_size, -1, x.shape[-1]) + hw_shape = (hw_shape[0], hw_shape[1] // 2) + + if self.use_attn: + attention_blocks = [self.attn1, self.attn2, self.attn3] + if i <= len(attention_blocks) - 1: + attn_out, _ = attention_blocks[i](x, x, x) + x = x + attn_out + if i == len(attention_blocks) - 1: + x = self.norm_attn(x) + + if (not self.use_attn) and (i == len(self.layers) - 1) and self.final_norm: + x = self.norm(x) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (x,) + + if i in self.out_indices: + out = x + out = out.reshape(batch_size, hw_shape[0], hw_shape[1], channels).permute(0, 3, 1, 2).contiguous() + if self.output_cls_token: + out = [out, x[:, 0]] + feature_maps.append(out) + + if not return_dict: + return tuple(feature_maps) + + return BaseModelOutput( + last_hidden_state=feature_maps[-1] if feature_maps else x, + hidden_states=all_hidden_states, + ) diff --git a/skysensepp-vit-msl-s2/modeling_utils.py b/skysensepp-vit-msl-s2/modeling_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..93feae77f3a3c46e65167f2ed312a25b7cd4ad3a --- /dev/null +++ b/skysensepp-vit-msl-s2/modeling_utils.py @@ -0,0 +1,557 @@ +"""SkySense: Pure PyTorch + HuggingFace Transformers implementation. + +Shared utility modules used across SkySense model implementations. +""" + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def to_2tuple(x): + """Convert to a 2-tuple.""" + if isinstance(x, (list, tuple)): + return tuple(x) + return (x, x) + + +class DropPath(nn.Module): + """Drop paths (stochastic depth) per sample. + + Args: + drop_prob (float): Probability of dropping a path. Default: 0.0. + """ + + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.drop_prob == 0.0 or not self.training: + return x + keep_prob = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor = torch.floor(random_tensor + keep_prob) + output = x / keep_prob * random_tensor + return output + + +class PatchEmbed(nn.Module): + """Image to Patch Embedding using Conv2d. + + Args: + in_channels (int): Number of input channels. Default: 3. + embed_dims (int): Embedding dimension. Default: 96. + kernel_size (int): Kernel size of the projection. Default: 4. + stride (int): Stride of the projection. Default: 4. + padding (int): Padding of the projection. Default: 0. + norm_layer (nn.Module or None): Normalization layer. Default: nn.LayerNorm. + input_size (int or tuple or None): Input resolution for calculating output size. + """ + + def __init__( + self, + in_channels: int = 3, + embed_dims: int = 96, + kernel_size: int = 4, + stride: int = 4, + padding: int = 0, + norm_layer: Optional[type] = nn.LayerNorm, + input_size: Optional[int] = None, + ): + super().__init__() + self.projection = nn.Conv2d( + in_channels, embed_dims, + kernel_size=kernel_size, stride=stride, padding=padding, + ) + self.norm = norm_layer(embed_dims) if norm_layer else nn.Identity() + + # Compute init output size if input_size is given + if input_size is not None: + input_size = to_2tuple(input_size) + self.init_out_size = ( + (input_size[0] - kernel_size + 2 * padding) // stride + 1, + (input_size[1] - kernel_size + 2 * padding) // stride + 1, + ) + else: + self.init_out_size = None + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]: + x = self.projection(x) # (B, C, H, W) + out_size = (x.shape[2], x.shape[3]) + x = x.flatten(2).transpose(1, 2) # (B, H*W, C) + x = self.norm(x) + return x, out_size + + +class FFN(nn.Module): + """Feed-Forward Network. + + Args: + embed_dims (int): Input dimension. + feedforward_channels (int): Hidden dimension. + num_fcs (int): Number of FC layers. Default: 2. + ffn_drop (float): Dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + act_layer (nn.Module): Activation layer class. Default: nn.GELU. + add_identity (bool): Whether to add identity connection. Default: True. + """ + + def __init__( + self, + embed_dims: int, + feedforward_channels: int, + num_fcs: int = 2, + ffn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: type = nn.GELU, + add_identity: bool = True, + ): + super().__init__() + assert num_fcs >= 2, f"num_fcs must be >= 2, got {num_fcs}" + self.embed_dims = embed_dims + self.feedforward_channels = feedforward_channels + self.add_identity = add_identity + + layers = [] + in_channels = embed_dims + for i in range(num_fcs - 1): + layers.append(nn.Linear(in_channels, feedforward_channels)) + layers.append(act_layer()) + layers.append(nn.Dropout(ffn_drop)) + in_channels = feedforward_channels + layers.append(nn.Linear(feedforward_channels, embed_dims)) + layers.append(nn.Dropout(ffn_drop)) + self.layers = nn.Sequential(*layers) + + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, identity: Optional[torch.Tensor] = None) -> torch.Tensor: + out = self.layers(x) + out = self.drop_path(out) + if self.add_identity: + if identity is None: + identity = x + out = out + identity + return out + + +class WindowMSAV2(nn.Module): + """Window-based Multi-head Self-Attention for Swin Transformer V2. + + Uses cosine attention and log-spaced continuous position bias (log-CPB). + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (tuple[int]): Window size (Wh, Ww). + pretrained_window_size (tuple[int]): Pretrained window size for CPB. Default: (0, 0). + qkv_bias (bool): If True, add learnable bias to q, k, v. Default: True. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Output projection dropout rate. Default: 0.0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: Tuple[int, int], + pretrained_window_size: Tuple[int, int] = (0, 0), + qkv_bias: bool = True, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ): + super().__init__() + self.embed_dims = embed_dims + self.num_heads = num_heads + self.window_size = window_size + self.pretrained_window_size = pretrained_window_size + + self.logit_scale = nn.Parameter( + torch.log(10 * torch.ones((num_heads, 1, 1)))) + + # MLP for continuous relative position bias (log-CPB) + self.cpb_mlp = nn.Sequential( + nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False), + ) + + # Build relative coords table + self._build_relative_coords_table() + # Build relative position index + self._build_relative_position_index() + + self.qkv = nn.Linear(embed_dims, embed_dims * 3, bias=False) + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(embed_dims)) + self.v_bias = nn.Parameter(torch.zeros(embed_dims)) + else: + self.q_bias = None + self.v_bias = None + + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(embed_dims, embed_dims) + self.proj_drop = nn.Dropout(proj_drop) + self.softmax = nn.Softmax(dim=-1) + + def _build_relative_coords_table(self): + """Build the relative coordinates table for log-CPB.""" + Wh, Ww = self.window_size + # Table of relative coordinates + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) # (1, (2Wh-1)*(2Ww-1), 2) + + # Normalize to [-1, 1] and apply log-scale + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 # normalize to -8, 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + self.register_buffer("relative_coords_table", coords_table) + + def _build_relative_position_index(self): + """Build the pairwise relative position index for each window token.""" + Wh, Ww = self.window_size + coords_h = torch.arange(Wh) + coords_w = torch.arange(Ww) + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing='ij')) + coords_flatten = coords.view(2, -1) + + relative_coords = ( + coords_flatten[:, :, None] - coords_flatten[:, None, :] + ) # (2, Wh*Ww, Wh*Ww) + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += Wh - 1 + relative_coords[:, :, 1] += Ww - 1 + relative_coords[:, :, 0] *= 2 * Ww - 1 + relative_position_index = relative_coords.sum(-1) # (Wh*Ww, Wh*Ww) + self.register_buffer("relative_position_index", relative_position_index) + + def _compute_position_bias(self, N): + """Compute relative position bias, supporting dynamic window sizes. + + The log-CPB (Continuous Position Bias) MLP can generalize to any window + size by computing bias from normalized relative coordinates. + """ + init_N = self.window_size[0] * self.window_size[1] + if N == init_N: + # Use pre-built tables + relative_position_bias_table = self.cpb_mlp( + self.relative_coords_table + ).view(-1, self.num_heads) + relative_position_bias = relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view(N, N, -1) + else: + # Dynamic: compute for actual window size on-the-fly + Wh = Ww = int(math.sqrt(N)) + coords_h = torch.arange(-(Wh - 1), Wh, dtype=torch.float32, device=self.logit_scale.device) + coords_w = torch.arange(-(Ww - 1), Ww, dtype=torch.float32, device=self.logit_scale.device) + coords_table = torch.stack( + torch.meshgrid(coords_h, coords_w, indexing='ij') + ).flatten(1).transpose(0, 1).unsqueeze(0) + if self.pretrained_window_size[0] > 0: + coords_table[:, :, 0] /= (self.pretrained_window_size[0] - 1) + coords_table[:, :, 1] /= (self.pretrained_window_size[1] - 1) + else: + coords_table[:, :, 0] /= max(Wh - 1, 1) + coords_table[:, :, 1] /= max(Ww - 1, 1) + coords_table *= 8 + coords_table = ( + torch.sign(coords_table) + * torch.log2(torch.abs(coords_table) + 1.0) + / math.log2(8) + ) + # Build position index for actual window size + ch = torch.arange(Wh, device=self.logit_scale.device) + cw = torch.arange(Ww, device=self.logit_scale.device) + coords = torch.stack(torch.meshgrid(ch, cw, indexing='ij')) + coords_flat = coords.view(2, -1) + rel = coords_flat[:, :, None] - coords_flat[:, None, :] + rel = rel.permute(1, 2, 0).contiguous() + rel[:, :, 0] += Wh - 1 + rel[:, :, 1] += Ww - 1 + rel[:, :, 0] *= 2 * Ww - 1 + pos_index = rel.sum(-1) + + bias_table = self.cpb_mlp(coords_table).view(-1, self.num_heads) + relative_position_bias = bias_table[ + pos_index.view(-1) + ].view(N, N, -1) + + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + return relative_position_bias + + def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: (num_windows*B, N, C) where N = Wh*Ww + mask: (num_windows, N, N) or None + """ + B_, N, C = x.shape + + # Compute QKV with bias + if self.q_bias is not None: + qkv_bias = torch.cat( + (self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias)) + qkv = F.linear(x, self.qkv.weight, qkv_bias) + else: + qkv = self.qkv(x) + + qkv = qkv.reshape(B_, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.permute(2, 0, 3, 1, 4) + q, k, v = qkv.unbind(0) + + # Cosine attention + attn = F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1) + logit_scale = torch.clamp( + self.logit_scale, max=math.log(1.0 / 0.01) + ).exp() + attn = attn * logit_scale + + # Log-CPB relative position bias (supports dynamic window sizes) + relative_position_bias = self._compute_position_bias(N) + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + attn = attn + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + + attn = self.softmax(attn) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class ShiftWindowMSA(nn.Module): + """Shifted Window Multi-head Self-Attention. + + Args: + embed_dims (int): Number of input channels. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. Default: 0. + attn_drop (float): Attention dropout rate. Default: 0.0. + proj_drop (float): Projection dropout rate. Default: 0.0. + drop_path (float): Drop path rate. Default: 0.0. + pad_small_map (bool): Pad small feature maps to window size. Default: False. + pretrained_window_size (int): Pretrained window size. Default: 0. + """ + + def __init__( + self, + embed_dims: int, + num_heads: int, + window_size: int, + shift_size: int = 0, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + drop_path: float = 0.0, + pad_small_map: bool = False, + pretrained_window_size: int = 0, + ): + super().__init__() + self.window_size = window_size + self.shift_size = shift_size + self.pad_small_map = pad_small_map + + self.w_msa = WindowMSAV2( + embed_dims=embed_dims, + num_heads=num_heads, + window_size=to_2tuple(window_size), + pretrained_window_size=to_2tuple(pretrained_window_size), + attn_drop=attn_drop, + proj_drop=proj_drop, + ) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> torch.Tensor: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W, f"Input length {L} != H*W ({H}*{W})" + + x = x.view(B, H, W, C) + + window_size = self.window_size + shift_size = self.shift_size + + # Pad or shrink window + if self.pad_small_map: + pad_r = (window_size - W % window_size) % window_size + pad_b = (window_size - H % window_size) % window_size + x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b)) + _, Hp, Wp, _ = x.shape + else: + Hp, Wp = H, W + if window_size > Hp: + window_size = Hp + shift_size = 0 + if window_size > Wp: + window_size = Wp + shift_size = 0 + + # Compute attention mask for SW-MSA + attn_mask = self._compute_attn_mask(Hp, Wp, window_size, shift_size, x.device) + + # Cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) + + # Partition windows + x_windows = self._window_partition(x, window_size) + # (num_windows*B, window_size*window_size, C) + + # W-MSA/SW-MSA + attn_windows = self.w_msa(x_windows, mask=attn_mask) + + # Merge windows + x = self._window_reverse(attn_windows, window_size, Hp, Wp) + + # Reverse cyclic shift + if shift_size > 0: + x = torch.roll(x, shifts=(shift_size, shift_size), dims=(1, 2)) + + if self.pad_small_map and (pad_r > 0 or pad_b > 0): + x = x[:, :H, :W, :].contiguous() + + x = x.view(B, H * W, C) + x = self.drop_path(x) + return x + + @staticmethod + def _window_partition(x: torch.Tensor, window_size: int) -> torch.Tensor: + """Partition into non-overlapping windows.""" + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous() + windows = windows.view(-1, window_size * window_size, C) + return windows + + @staticmethod + def _window_reverse(windows: torch.Tensor, window_size: int, H: int, W: int) -> torch.Tensor: + """Reverse window partition.""" + B_nW = windows.shape[0] + nH = H // window_size + nW = W // window_size + B = B_nW // (nH * nW) + x = windows.view(B, nH, nW, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous() + x = x.view(B, H, W, -1) + return x + + @staticmethod + def _compute_attn_mask(H, W, window_size, shift_size, device): + """Compute attention mask for shifted window attention.""" + if shift_size <= 0: + return None + img_mask = torch.zeros((1, H, W, 1), device=device) + h_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + w_slices = ( + slice(0, -window_size), + slice(-window_size, -shift_size), + slice(-shift_size, None), + ) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + # Partition mask + mask_windows = img_mask.view( + 1, H // window_size, window_size, W // window_size, window_size, 1 + ) + mask_windows = mask_windows.permute(0, 1, 3, 2, 4, 5).contiguous() + mask_windows = mask_windows.view(-1, window_size * window_size) + + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0) + attn_mask = attn_mask.masked_fill(attn_mask == 0, 0.0) + return attn_mask + + +class PatchMerging(nn.Module): + """Patch Merging Layer for downsampling (2x). + + Args: + in_channels (int): Input channels. + out_channels (int): Output channels. + norm_layer (type): Normalization layer. Default: nn.LayerNorm. + is_post_norm (bool): Apply norm after linear. Default: True. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + norm_layer: type = nn.LayerNorm, + is_post_norm: bool = True, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.is_post_norm = is_post_norm + self.reduction = nn.Linear(4 * in_channels, out_channels, bias=False) + if is_post_norm: + self.norm = norm_layer(out_channels) + else: + self.norm = norm_layer(4 * in_channels) + + def forward(self, x: torch.Tensor, hw_shape: Tuple[int, int]) -> Tuple[torch.Tensor, Tuple[int, int]]: + B, L, C = x.shape + H, W = hw_shape + assert L == H * W + + x = x.view(B, H, W, C) + + # Pad if needed + pad_h = H % 2 + pad_w = W % 2 + if pad_h or pad_w: + x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h)) + + x0 = x[:, 0::2, 0::2, :] + x1 = x[:, 1::2, 0::2, :] + x2 = x[:, 0::2, 1::2, :] + x3 = x[:, 1::2, 1::2, :] + x = torch.cat([x0, x1, x2, x3], dim=-1) + + out_h = (H + pad_h) // 2 + out_w = (W + pad_w) // 2 + x = x.view(B, out_h * out_w, 4 * C) + + if self.is_post_norm: + x = self.reduction(x) + x = self.norm(x) + else: + x = self.norm(x) + x = self.reduction(x) + + return x, (out_h, out_w) diff --git a/skysensepp-vit-msl-s2/pipeline_skysensepp.py b/skysensepp-vit-msl-s2/pipeline_skysensepp.py new file mode 100644 index 0000000000000000000000000000000000000000..0d9a5354361c9a08abd1a82b3df8d4ec678a209b --- /dev/null +++ b/skysensepp-vit-msl-s2/pipeline_skysensepp.py @@ -0,0 +1,86 @@ +"""Custom HuggingFace pipeline for SkySense++ MSL feature extraction.""" + +from typing import Any, Dict, Optional, Union + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusMSLFeatureExtractionPipeline(Pipeline): + """Pipeline for SkySense++ MSL backbones. + + Expects image tensors plus semantic annotation maps (class indices). + """ + + def _sanitize_parameters( + self, + annotation=None, + mask=None, + output_hidden_states=None, + **kwargs, + ): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + + if annotation is not None: + preprocess_params["annotation"] = annotation + if mask is not None: + forward_params["mask"] = mask + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + + return preprocess_params, forward_params, postprocess_params + + def preprocess( + self, + pixel_values: Any, + annotation: Optional[Any] = None, + **kwargs, + ) -> Dict[str, torch.Tensor]: + if isinstance(pixel_values, dict): + annotation = pixel_values.get("annotation", annotation) + pixel_values = pixel_values.get("pixel_values", pixel_values) + + if isinstance(pixel_values, np.ndarray): + pixel_values = torch.from_numpy(pixel_values).float() + elif not isinstance(pixel_values, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for pixel_values, got {type(pixel_values)}" + ) + + if annotation is None: + raise ValueError("SkySense++ MSL models require an `annotation` semantic map.") + + if isinstance(annotation, np.ndarray): + annotation = torch.from_numpy(annotation).long() + elif not isinstance(annotation, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for annotation, got {type(annotation)}" + ) + + if pixel_values.ndim == 3: + pixel_values = pixel_values.unsqueeze(0) + if annotation.ndim == 2: + annotation = annotation.unsqueeze(0) + + return {"pixel_values": pixel_values, "annotation": annotation} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + pixel_values=model_inputs["pixel_values"], + annotation=model_inputs["annotation"], + mask=kwargs.get("mask"), + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"last_hidden_state": outputs.last_hidden_state} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result diff --git a/skysensepp-vit-msl-s2/pipeline_skysensepp_fusion.py b/skysensepp-vit-msl-s2/pipeline_skysensepp_fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..80f070fa9fd3f18a053caa1afef9dc142bce2597 --- /dev/null +++ b/skysensepp-vit-msl-s2/pipeline_skysensepp_fusion.py @@ -0,0 +1,53 @@ +"""Optional pipeline for SkySense++ fusion neck.""" + +from typing import Any, Dict + +import numpy as np +import torch +from transformers import Pipeline + + +class SkySensePlusPlusFusionNeckPipeline(Pipeline): + """Pipeline for the optional SkySense++ fusion neck module. + + Expects concatenated multi-modal tokens per spatial location: + ``(batch, num_modalities, input_dims)``. + """ + + def _sanitize_parameters(self, output_hidden_states=None, **kwargs): + preprocess_params = {} + forward_params = {} + postprocess_params = {} + if output_hidden_states is not None: + forward_params["output_hidden_states"] = output_hidden_states + return preprocess_params, forward_params, postprocess_params + + def preprocess(self, hidden_states: Any, **kwargs) -> Dict[str, torch.Tensor]: + if isinstance(hidden_states, dict): + hidden_states = hidden_states["hidden_states"] + + if isinstance(hidden_states, np.ndarray): + hidden_states = torch.from_numpy(hidden_states).float() + elif not isinstance(hidden_states, torch.Tensor): + raise TypeError( + f"Expected tensor or ndarray for hidden_states, got {type(hidden_states)}" + ) + if hidden_states.ndim == 2: + hidden_states = hidden_states.unsqueeze(0) + return {"hidden_states": hidden_states} + + def _forward(self, model_inputs: Dict[str, torch.Tensor], **kwargs) -> Dict[str, Any]: + with torch.no_grad(): + outputs = self.model( + hidden_states=model_inputs["hidden_states"], + output_hidden_states=kwargs.get("output_hidden_states", False), + return_dict=True, + ) + return {"outputs": outputs} + + def postprocess(self, model_outputs: Dict[str, Any], **kwargs) -> Dict[str, Any]: + outputs = model_outputs["outputs"] + result = {"pooler_output": outputs.pooler_output} + if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: + result["hidden_states"] = outputs.hidden_states + return result