Improve auxiliary_in_channels default behavior in UperNet (#37540)

Improve auxiliary_in_channels behavior in UperNet

Co-authored-by: Pavel Iakubovskii <[email protected]>

Author: simonreise

src/transformers/models/upernet/configuration_upernet.py

@@ -96,7 +96,7 @@ def __init__(
         pool_scales=[1, 2, 3, 6],
         use_auxiliary_head=True,
         auxiliary_loss_weight=0.4,
-        auxiliary_in_channels=384,
+        auxiliary_in_channels=None,
         auxiliary_channels=256,
         auxiliary_num_convs=1,
         auxiliary_concat_input=False,

src/transformers/models/upernet/modeling_upernet.py

@@ -218,12 +218,14 @@ class UperNetFCNHead(nn.Module):
     """
 
     def __init__(
-        self, config, in_index: int = 2, kernel_size: int = 3, dilation: Union[int, Tuple[int, int]] = 1
+        self, config, in_channels, in_index: int = 2, kernel_size: int = 3, dilation: Union[int, Tuple[int, int]] = 1
     ) -> None:
         super().__init__()
 
         self.config = config
-        self.in_channels = config.auxiliary_in_channels
+        self.in_channels = (
+            in_channels[in_index] if config.auxiliary_in_channels is None else config.auxiliary_in_channels
+        )
         self.channels = config.auxiliary_channels
         self.num_convs = config.auxiliary_num_convs
         self.concat_input = config.auxiliary_concat_input
@@ -292,7 +294,9 @@ def __init__(self, config):
 
         # Semantic segmentation head(s)
         self.decode_head = UperNetHead(config, in_channels=self.backbone.channels)
-        self.auxiliary_head = UperNetFCNHead(config) if config.use_auxiliary_head else None
+        self.auxiliary_head = (
+            UperNetFCNHead(config, in_channels=self.backbone.channels) if config.use_auxiliary_head else None
+        )
 
         # Initialize weights and apply final processing
         self.post_init()