Anpassungen für HF, Checkpoint umgewandelt, config.json angepasst

config.json

@@ -1,16 +1,57 @@
 {
+  "architectures": [
+    "MySegformerForSemanticSegmentation"
+  ],
+  "attn_drop_rate": 0.0,
+  "drop_path_rate": 0.0,
+  "drop_rate": 0.0,
+  "embed_dims": [
+    64,
+    128,
+    320,
+    512
+  ],
+  "mlp_ratio": 4,
   "model_type": "my_segformer",
-  "embed_dims": 64,
+  "num_classes": 19,
+  "num_heads": [
+    1,
+    2,
+    4,
+    8
+  ],
+  "num_layers": [
+    3,
+    4,
+    6,
+    3
+  ],
   "num_stages": 4,
-  "num_layers": [3,4,6,3],
-  "num_heads": [1,2,4,8],
-  "patch_sizes": [7,3,3,3],
-  "strides": [4,2,2,2],
-  "sr_ratios": [8,4,2,1],
-  "mlp_ratio": 4,
+  "out_indices": [
+    0,
+    1,
+    2,
+    3
+  ],
+  "patch_sizes": [
+    7,
+    3,
+    3,
+    3
+  ],
   "qkv_bias": true,
-  "drop_rate": 0.0,
-  "attn_drop_rate": 0.0,
-  "drop_path_rate": 0.0,
-  "out_indices": [0,1,2,3]
+  "sr_ratios": [
+    8,
+    4,
+    2,
+    1
+  ],
+  "strides": [
+    4,
+    2,
+    2,
+    2
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.55.0"
 }

hf_segformer_converted/config.json

@@ -0,0 +1,57 @@
+{
+  "architectures": [
+    "MySegformerForSemanticSegmentation"
+  ],
+  "attn_drop_rate": 0.0,
+  "drop_path_rate": 0.0,
+  "drop_rate": 0.0,
+  "embed_dims": [
+    64,
+    128,
+    320,
+    512
+  ],
+  "mlp_ratio": 4,
+  "model_type": "my_segformer",
+  "num_classes": 19,
+  "num_heads": [
+    1,
+    2,
+    4,
+    8
+  ],
+  "num_layers": [
+    3,
+    4,
+    6,
+    3
+  ],
+  "num_stages": 4,
+  "out_indices": [
+    0,
+    1,
+    2,
+    3
+  ],
+  "patch_sizes": [
+    7,
+    3,
+    3,
+    3
+  ],
+  "qkv_bias": true,
+  "sr_ratios": [
+    8,
+    4,
+    2,
+    1
+  ],
+  "strides": [
+    4,
+    2,
+    2,
+    2
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.55.0"
+}

mix_vision_transformer_config.py

@@ -5,7 +5,7 @@ class MySegformerConfig(PretrainedConfig):
 
     def __init__(
         self,
-        embed_dims=64,
+        embed_dims=[64, 128, 320, 512],
         num_stages=4,
         num_layers=[3, 4, 6, 3],
         num_heads=[1, 2, 4, 8],
@@ -21,6 +21,11 @@ class MySegformerConfig(PretrainedConfig):
         **kwargs
     ):
         super().__init__(**kwargs)
+
+        # Absicherung, falls embed_dims als int übergeben wird
+        if isinstance(embed_dims, int):
+            embed_dims = [embed_dims]
+
         self.embed_dims = embed_dims
         self.num_stages = num_stages
         self.num_layers = num_layers
@@ -34,3 +39,4 @@ class MySegformerConfig(PretrainedConfig):
         self.attn_drop_rate = attn_drop_rate
         self.drop_path_rate = drop_path_rate
         self.out_indices = out_indices
+

mmengineToHFCheckpoint.py

@@ -0,0 +1,79 @@
+import torch
+from mix_vision_transformer_config import MySegformerConfig
+from modeling_my_segformer import MySegformerForSemanticSegmentation
+
+def convert_mmengine_checkpoint_to_hf(mm_checkpoint_path, hf_save_dir):
+    # 1. Lade mmengine checkpoint
+    mm_ckpt = torch.load(mm_checkpoint_path, map_location="cpu")
+    if 'state_dict' in mm_ckpt:
+        mm_state_dict = mm_ckpt['state_dict']
+    else:
+        mm_state_dict = mm_ckpt
+
+    # 2. Erstelle Config & Modell (achte darauf, dass Config-Parameter zum Checkpoint passen)
+    config = MySegformerConfig(
+        embed_dims=[64, 128, 320, 512],  # <--- korrekte Liste mit 4 Werten
+        num_stages=4,
+        num_layers=[3, 4, 6, 3],
+        num_heads=[1, 2, 4, 8],
+        patch_sizes=[7, 3, 3, 3],
+        strides=[4, 2, 2, 2],
+        sr_ratios=[8, 4, 2, 1],
+        mlp_ratio=4,
+        qkv_bias=True,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        out_indices=(0, 1, 2, 3),
+        num_classes=19
+    )
+
+    model = MySegformerForSemanticSegmentation(config)
+
+    # 3. Mappe mmengine Keys auf HF Keys
+    hf_state_dict = {}
+
+    for k, v in mm_state_dict.items():
+        new_k = k
+
+        # Falls "module." als Prefix da ist (DataParallel), entfernen
+        if new_k.startswith("module."):
+            new_k = new_k[len("module."):]
+
+        # Mapping von decode_head.* -> segmentation_head.*
+        if new_k.startswith("decode_head."):
+            new_k = new_k.replace("decode_head.", "segmentation_head.")
+
+        # BatchNorm-Namen vereinheitlichen
+        new_k = new_k.replace(".bn.", ".")
+
+        # Nur Keys übernehmen, die im HF-Modell existieren
+        if new_k not in model.state_dict():
+            print(f"⚠️ Ignoriere {new_k} (nicht im HF-Modell)")
+            continue
+
+        hf_state_dict[new_k] = v
+
+    # 4. Lade die Gewichte ins Modell
+    missing_keys, unexpected_keys = model.load_state_dict(hf_state_dict, strict=False)
+
+    print("Missing keys:", missing_keys)
+    print("Unexpected keys:", unexpected_keys)
+
+    # 5. Speichere das HF-kompatible Modell & Config
+    model.save_pretrained(hf_save_dir)
+    config.save_pretrained(hf_save_dir)
+
+    print(f"✅ Model und Config erfolgreich gespeichert in {hf_save_dir}")
+
+    # 5b. Auch als klassische .pth-Datei speichern
+    pth_path = hf_save_dir.rstrip("/") + ".pth"
+    torch.save(model.state_dict(), pth_path)
+    print(f"✅ Zusätzlich als .pth gespeichert unter {pth_path}")
+
+
+if __name__ == "__main__":
+    mm_checkpoint_path = "./segformer-b5-bsm_hq.pth"
+    hf_save_dir = "hf_segformer_converted"
+
+    convert_mmengine_checkpoint_to_hf(mm_checkpoint_path, hf_save_dir)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:818e1b51093355a5915fd844e68edc7af583b2a397272c84e6e12a670896bc23
+size 98934820

modeling_my_segformer.py

@@ -2,17 +2,17 @@ from transformers import PreTrainedModel
 import torch
 import torch.nn as nn
 from segformer_plusplus.utils import resize
-from segformer_plusplus.model.backbone.mit import MixVisionTransformer  # deine Backbone-Importierung
-from mix_vision_transformer_config import MySegformerConfig  # deine Config-Importierung
+from segformer_plusplus.model.backbone.mit import MixVisionTransformer  # Backbone-Import
+from mix_vision_transformer_config import MySegformerConfig  # Config-Import
 
-# Head-Implementierung (etwas vereinfacht und angepasst)
+# Head-Implementierung (vereinfacht)
 class SegformerHead(nn.Module):
     def __init__(self,
                  in_channels=[64, 128, 256, 512],  # anpassen je nach Backbone-Ausgabe!
                  in_index=[0, 1, 2, 3],
                  channels=256,
                  dropout_ratio=0.1,
-                 out_channels=19,  # Anzahl Klassen, anpassen!
+                 out_channels=19,  # Anzahl Klassen anpassen!
                  norm_cfg=None,
                  align_corners=False,
                  interpolate_mode='bilinear'):
@@ -26,6 +26,11 @@ class SegformerHead(nn.Module):
         self.align_corners = align_corners
         self.interpolate_mode = interpolate_mode
 
+        print(f"in_channels: {self.in_channels}, type: {type(self.in_channels)}")
+        print(f"in_index: {self.in_index}, type: {type(self.in_index)}")
+        print(f"len(in_channels): {len(self.in_channels) if hasattr(self.in_channels, '__len__') else 'no len'}")
+        print(f"len(in_index): {len(self.in_index) if hasattr(self.in_index, '__len__') else 'no len'}")
+
         self.act_cfg = dict(type='ReLU')
         self.conv_seg = nn.Conv2d(channels, out_channels, kernel_size=1)
         self.dropout = nn.Dropout2d(dropout_ratio) if dropout_ratio > 0 else None
@@ -43,6 +48,7 @@ class SegformerHead(nn.Module):
                     out_channels=channels,
                     kernel_size=1,
                     stride=1,
+                    bias=False,
                     norm_cfg=norm_cfg,
                     act_cfg=self.act_cfg))
 
@@ -50,6 +56,7 @@ class SegformerHead(nn.Module):
             in_channels=channels * num_inputs,
             out_channels=channels,
             kernel_size=1,
+            bias=False,
             norm_cfg=norm_cfg)
 
     def cls_seg(self, feat):
@@ -81,9 +88,9 @@ class MySegformerForSemanticSegmentation(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
-        # Backbone initialisieren mit Parametern aus Config
+        # Wichtig: die gesamte Liste übergeben, nicht nur das erste Element
         self.backbone = MixVisionTransformer(
-            embed_dims=config.embed_dims,
+            embed_dims=config.embed_dims,    # GANZE Liste, z.B. [64, 128, 320, 512]
             num_stages=config.num_stages,
             num_layers=config.num_layers,
             num_heads=config.num_heads,
@@ -98,17 +105,18 @@ class MySegformerForSemanticSegmentation(PreTrainedModel):
             out_indices=config.out_indices
         )
 
-        # Head initialisieren, out_channels aus config oder fix setzen
+        # Sicherstellen, dass in_channels eine Liste ist
+        in_channels = config.embed_dims
+        if isinstance(in_channels, int):
+            in_channels = [in_channels]
+
+        print(f"config.embed_dims: {config.embed_dims}, type: {type(config.embed_dims)}")
         self.segmentation_head = SegformerHead(
-            in_channels=[64, 128, 256, 512],  # <- Anpassen, je nachdem wie Backbone ausgibt!
+            in_channels=config.embed_dims,            # z.B. [64, 128, 320, 512]
+            in_index=list(config.out_indices),        # z.B. [0, 1, 2, 3]
             out_channels=config.num_classes if hasattr(config, 'num_classes') else 19,
             dropout_ratio=0.1,
             align_corners=False
         )
 
         self.post_init()
-
-    def forward(self, x):
-        features = self.backbone(x)
-        segmentation_output = self.segmentation_head(features)
-        return segmentation_output

preTrainedTest.py

@@ -0,0 +1,13 @@
+from modeling_my_segformer import MySegformerForSemanticSegmentation
+from mix_vision_transformer_config import MySegformerConfig
+
+# Der Pfad zu deinem HF-Repo (kann auch einfach als String benutzt werden)
+model_name_or_path = "TimM77/SegformerPlusPlus"
+
+# Config laden (automatisch aus config.json im Repo)
+config = MySegformerConfig.from_pretrained(model_name_or_path)
+
+# Modell laden (Gewichte aus pytorch_model.bin + Config)
+model = MySegformerForSemanticSegmentation.from_pretrained(model_name_or_path, config=config)
+
+print(model, config)

segformer_plusplus/model/backbone/mit.py

@@ -415,7 +415,7 @@ class MixVisionTransformer(BaseModule):
         cur = 0
         self.layers = ModuleList()
         for i, num_layer in enumerate(num_layers):
-            embed_dims_i = embed_dims * num_heads[i]
+            embed_dims_i = embed_dims[i]
             patch_embed = PatchEmbed(
                 in_channels=in_channels,
                 embed_dims=embed_dims_i,