Anpassungen der COnfigs für b5 backbone

config.json

@@ -1,9 +1,9 @@
 {
   "model_type": "my_segformer",
   "in_channels": 3,
-  "embed_dims": 64,
+  "embed_dims": [64, 128, 320, 512],
   "num_stages": 4,
-  "num_layers": [2, 2, 2, 2],
+  "num_layers": [3, 6, 40, 3],
   "num_heads": [1, 2, 5, 8],
   "patch_sizes": [7, 3, 3, 3],
   "strides": [4, 2, 2, 2],
@@ -24,4 +24,4 @@
     "align_corners": false,
     "interpolate_mode": "bilinear"
   }
-}
+}

mix_vision_transformer_config.py

@@ -6,9 +6,11 @@ class MySegformerConfig(PretrainedConfig):
     def __init__(
         self,
         in_channels=3,
-        embed_dims=[32, 64, 160, 256],
+        # Corrected for SegFormer-B5: list of embedding dimensions for each stage
+        embed_dims=[64, 128, 320, 512],
         num_stages=4,
-        num_layers=[2, 2, 2, 2],
+        # Corrected for SegFormer-B5: number of transformer layers in each stage
+        num_layers=[3, 6, 40, 3],
         num_heads=[1, 2, 5, 8],
         patch_sizes=[7, 3, 3, 3],
         strides=[4, 2, 2, 2],
@@ -24,8 +26,8 @@ class MySegformerConfig(PretrainedConfig):
         **kwargs
     ):
         super().__init__(**kwargs)
-        self.embed_dims = embed_dims  # einfach int behalten
         self.in_channels = in_channels
+        self.embed_dims = embed_dims # This will now be a list, which is correct for SegFormer
         self.num_stages = num_stages
         self.num_layers = num_layers
         self.num_heads = num_heads
@@ -40,11 +42,12 @@ class MySegformerConfig(PretrainedConfig):
         self.out_indices = out_indices
         self.num_classes = num_classes
 
-        # optionaler Block für Head-Config (falls decode_head nicht übergeben wurde)
+        # Optional block for Head-Config (if decode_head not passed)
         if decode_head is None:
             decode_head = {
-                "in_channels": [64, 128, 320, 512],          # [32, 64, 160, 256]
-                "in_index": list(range(self.num_stages)),# [0, 1, 2, 3]
+                # Corrected for SegFormer-B5: input channels for the decode head from each stage
+                "in_channels": [64, 128, 320, 512],
+                "in_index": list(range(self.num_stages)),
                 "channels": 256,
                 "dropout_ratio": 0.1,
                 "out_channels": self.num_classes,
@@ -52,4 +55,4 @@ class MySegformerConfig(PretrainedConfig):
                 "interpolate_mode": "bilinear"
             }
 
-        self.decode_head = decode_head
+        self.decode_head = decode_head