Improve auxiliary_in_channels default behavior in UperNet

[simonreise]

Feature request

Now the number of input channels in auxiliary UperNet FCN Head is controlled by auxiliary_in_channels parameter in UperNetConfig, which is set to 384 by default. Not every backbone have 384 channels by default, so sometimes user have to set this parameter manually.

I propose to change the default model behaviour to automatically handle the number of input channels but keep the auxiliary_in_channels parameter with None as a default value so it still can be set manually by user if they need it.

Motivation

Now the number of input channels in auxiliary UperNet FCN Head is controlled by auxiliary_in_channels parameter in UperNetConfig, which is set to 384 by default. It works nice with some backbones (like Swin and Convnext) but crashes with backbones like BiT or ResNet.

So, now, if user wants to use some backbones, they have to set the parameter manually, which is frustrating and it is not always straightforward which value you should use with every specific backbone.

Your contribution

However, in most cases this behaviour can be automated.

I propose to set auxiliary_in_channels default value to None, so by default the model will handle the number of input channels automatically, but the parameter is still can be set by user if they need it.

I propose to add in_channels parameter to self.auxiliary_head and set self.in_channels from config.auxiliary_in_channels if it is not None or from in_channels if it is None.

UperNetForSemanticSegmentation.__init__

# Old
# self.auxiliary_head = UperNetFCNHead(config) if config.use_auxiliary_head else None
# New
self.auxiliary_head = UperNetFCNHead(config, in_channels=self.backbone.channels) if config.use_auxiliary_head else None

class UperNetFCNHead(nn.Module):
  def __init__(
          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
          # Old
          #self.in_channels = config.auxiliary_in_channels
          # New
          self.in_channels = in_channels[in_index] if config.auxiliary_in_channels is None else config.auxiliary_in_channels

I tested it with BEiT, BiT, ConvNeXT, ConvNeXTV2, DINOV2, DINOV2WithRegisters, FocalNet, Hiera, PVTV2, ResNet, Swin, SwinV2 and ViTDet backbones.

[Rocketknight1]

cc @qubvel @NielsRogge

[qubvel]

Hey @simonreise, looks good to me, feel free to open a PR and ping me for review