Anpassung an der modeling für b5 nutzung

modeling_my_segformer.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
 
-from segformer_plusplus.model.backbone.mit import MixVisionTransformer  # Backbone
+from segformer_plusplus.model.backbone.mit import MixVisionTransformer   # Backbone
 from mix_vision_transformer_config import MySegformerConfig             # Config
 from segformer_plusplus.model.head.segformer_head import SegformerHead # <-- dein Head
 
@@ -16,7 +16,8 @@ class MySegformerForSemanticSegmentation(PreTrainedModel):
 
         # Backbone (MixVisionTransformer)
         self.backbone = MixVisionTransformer(
-            embed_dims=config.embed_dims,       # z.B. [64, 128, 320, 512]
+            # Pass only the first element of embed_dims for the initial patch embedding
+            embed_dims=config.embed_dims[0],   # <--- KORRIGIERTE ZEILE
             num_stages=config.num_stages,
             num_layers=config.num_layers,
             num_heads=config.num_heads,
@@ -32,14 +33,17 @@ class MySegformerForSemanticSegmentation(PreTrainedModel):
         )
 
         # Head direkt importieren
-        in_channels = [64, 128, 320, 512]
+        # Use config.decode_head.in_channels directly, as it's defined in the config.
+        # This ensures consistency with the backbone's output channels for the head.
+        in_channels = config.decode_head["in_channels"]
 
         self.segmentation_head = SegformerHead(
-            in_channels=in_channels,                # Liste der Embeddings aus Backbone
-            in_index=list(config.out_indices),      # welche Feature Maps genutzt werden
-            out_channels=getattr(config, "num_classes", 19),  # Anzahl Klassen
-            dropout_ratio=0.1,
-            align_corners=False
+            in_channels=in_channels,                  # Liste der Embeddings aus Backbone
+            in_index=list(config.out_indices),        # welche Feature Maps genutzt werden
+            channels=config.decode_head["channels"],   # channels parameter for SegformerHead itself
+            dropout_ratio=config.decode_head["dropout_ratio"],
+            num_classes=getattr(config, "num_classes", 19), # Ensure num_classes is passed if not directly in decode_head config
+            align_corners=config.decode_head["align_corners"]
         )
 
         self.post_init()
@@ -48,11 +52,11 @@ class MySegformerForSemanticSegmentation(PreTrainedModel):
         # Backbone → Features (Liste von Tensors)
         features = self.backbone(x)
 
-                # Debug: Ausgabe der Shapes der Backbone-Features
+        # Debug: Ausgabe der Shapes der Backbone-Features
         for i, f in enumerate(features):
             print(f"Feature {i}: shape = {f.shape}")
 
         # Head → logits
         logits = self.segmentation_head(features)
 
-        return logits
+        return logits