I use the custom decoder I made but I think there's better or you can learn from those in the PyTorch_toolbelt decoders the most related two versions are

unet_v2.py

from typing import Tuple, List

import torch
import torch.nn.functional as F
from torch import nn

from .common import DecoderModule
from ..activations import ABN

__all__ = ["UNetDecoderV2", "UnetCentralBlockV2", "UnetDecoderBlockV2"]


class UnetCentralBlockV2(nn.Module):
    def __init__(self, in_dec_filters, out_filters, mask_channels, abn_block=ABN):
        super().__init__()
        self.bottleneck = nn.Conv2d(in_dec_filters, out_filters, kernel_size=1)

        self.conv1 = nn.Conv2d(out_filters, out_filters, kernel_size=3, padding=1, stride=2, bias=False)
        self.abn1 = abn_block(out_filters)
        self.conv2 = nn.Conv2d(out_filters, out_filters, kernel_size=3, padding=1, bias=False)
        self.abn2 = abn_block(out_filters)
        self.dsv = nn.Conv2d(out_filters, mask_channels, kernel_size=1)

    def forward(self, x):
        x = self.bottleneck(x)

        x = self.conv1(x)
        x = self.abn1(x)
        x = self.conv2(x)
        x = self.abn2(x)

        dsv = self.dsv(x)

        return x, dsv


class UnetDecoderBlockV2(nn.Module):
    """"""

    def __init__(
        self,
        in_dec_filters: int,
        in_enc_filters: int,
        out_filters: int,
        mask_channels: int,
        abn_block=ABN,
        pre_dropout_rate=0.0,
        post_dropout_rate=0.0,
        scale_factor=None,
        scale_mode="nearest",
        align_corners=None,
    ):
        super(UnetDecoderBlockV2, self).__init__()

        self.scale_factor = scale_factor
        self.scale_mode = scale_mode
        self.align_corners = align_corners

        self.bottleneck = nn.Conv2d(in_dec_filters + in_enc_filters, out_filters, kernel_size=1)

        self.conv1 = nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=1, padding=1, bias=False)
        self.abn1 = abn_block(out_filters)
        self.conv2 = nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=1, padding=1, bias=False)
        self.abn2 = abn_block(out_filters)

        self.pre_drop = nn.Dropout2d(pre_dropout_rate, inplace=True)

        self.post_drop = nn.Dropout2d(post_dropout_rate)

        self.dsv = nn.Conv2d(out_filters, mask_channels, kernel_size=1)

    def forward(self, x: torch.Tensor, enc: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        if self.scale_factor is not None:
            x = F.interpolate(
                x, scale_factor=self.scale_factor, mode=self.scale_mode, align_corners=self.align_corners
            )
        else:
            lat_size = enc.size()[2:]
            x = F.interpolate(x, size=lat_size, mode=self.scale_mode, align_corners=self.align_corners)

        x = torch.cat([x, enc], 1)

        x = self.bottleneck(x)
        x = self.pre_drop(x)

        x = self.conv1(x)
        x = self.abn1(x)

        x = self.conv2(x)
        x = self.abn2(x)

        x = self.post_drop(x)

        dsv = self.dsv(x)
        return x, dsv


class UNetDecoderV2(DecoderModule):
    def __init__(self, feature_maps: List[int], decoder_features: int, mask_channels: int, dropout=0.0, abn_block=ABN):
        super().__init__()

        if not isinstance(decoder_features, list):
            decoder_features = [decoder_features * (2 ** i) for i in range(len(feature_maps))]

        blocks = []
        for block_index, in_enc_features in enumerate(feature_maps[:-1]):
            blocks.append(
                UnetDecoderBlockV2(
                    decoder_features[block_index + 1],
                    in_enc_features,
                    decoder_features[block_index],
                    mask_channels,
                    abn_block=abn_block,
                    post_dropout_rate=dropout,
                )
            )

        self.center = UnetCentralBlockV2(feature_maps[-1], decoder_features[-1], mask_channels, abn_block=abn_block)
        self.blocks = nn.ModuleList(blocks)
        self.output_filters = decoder_features

        self.final = nn.Conv2d(decoder_features[0], mask_channels, kernel_size=1)

    def forward(self, feature_maps):

        output, dsv = self.center(feature_maps[-1])
        dsv_list = [dsv]

        for decoder_block, encoder_output in zip(reversed(self.blocks), reversed(feature_maps[:-1])):
            output, dsv = decoder_block(output, encoder_output)
            dsv_list.append(dsv)

        dsv_list = list(reversed(dsv_list))

        output = self.final(output)
        return output, dsv_list


#the normal one which has issues unet.py
from typing import List, Union, Type

import torch
from torch import nn

from .common import DecoderModule
from .. import DeconvolutionUpsample2d
from ..unet import UnetBlock

__all__ = ["UNetDecoder"]


class UNetDecoder(DecoderModule):
    def __init__(
        self,
        feature_maps: List[int],
        decoder_features: Union[int, List[int]] = None,
        unet_block=UnetBlock,
        upsample_block: Union[nn.Upsample, nn.ConvTranspose2d, Type[nn.PixelShuffle]] = None,
    ):
        super().__init__()

        # if not isinstance(decoder_features, list):
        #     decoder_features = [decoder_features * (2 ** i) for i in range(len(feature_maps))]
        # else:
        #     assert len(decoder_features) == len(
        #         feature_maps
        #     ), f"Incorrect number of decoder features: {decoder_features}, {feature_maps}"

        if upsample_block is None:
            upsample_block = nn.ConvTranspose2d

        blocks = []
        upsamples = []

        num_blocks = len(feature_maps) - 1  # Number of outputs is one less than encoder layers

        if decoder_features is None:
            decoder_features = [None] * num_blocks
        else:
            if len(decoder_features) != num_blocks:
                raise ValueError(f"decoder_features must have length of {num_blocks}")
        in_channels_for_upsample_block = feature_maps[-1]

        for block_index in reversed(range(num_blocks)):
            features_from_encoder = feature_maps[block_index]

            if isinstance(upsample_block, nn.Upsample):
                upsamples.append(upsample_block)
                out_channels_from_upsample_block = in_channels_for_upsample_block
            elif issubclass(upsample_block, nn.Upsample):
                upsamples.append(upsample_block(scale_factor=2))
                out_channels_from_upsample_block = in_channels_for_upsample_block
            elif issubclass(upsample_block, nn.PixelShuffle):
                upsamples.append(upsample_block(upscale_factor=2))
                out_channels_from_upsample_block = in_channels_for_upsample_block // 4
            elif issubclass(upsample_block, nn.ConvTranspose2d):
                up = upsample_block(
                    in_channels_for_upsample_block,
                    in_channels_for_upsample_block // 2,
                    kernel_size=3,
                    stride=2,
                    padding=1,
                )
                upsamples.append(up)
                out_channels_from_upsample_block = up.out_channels
            else:
                up = upsample_block(in_channels_for_upsample_block)
                upsamples.append(up)
                out_channels_from_upsample_block = up.out_channels

            in_channels = features_from_encoder + out_channels_from_upsample_block
            out_channels = decoder_features[block_index] or in_channels // 2
            blocks.append(unet_block(in_channels, out_channels))

            in_channels_for_upsample_block = out_channels
            decoder_features[block_index] = out_channels

        self.blocks = nn.ModuleList(blocks)
        self.upsamples = nn.ModuleList(upsamples)
        self.output_filters = decoder_features

    @property
    @torch.jit.unused
    def channels(self) -> List[int]:
        return self.output_filters

    def forward(self, feature_maps: List[torch.Tensor]) -> List[torch.Tensor]:
        x = feature_maps[-1]
        outputs = []
        num_feature_maps = len(feature_maps)
        for index, (upsample_block, decoder_block) in enumerate(zip(self.upsamples, self.blocks)):
            encoder_input = feature_maps[num_feature_maps - index - 2]

            if isinstance(upsample_block, (nn.ConvTranspose2d, DeconvolutionUpsample2d)):
                x = upsample_block(x, output_size=encoder_input.size())
            else:
                x = upsample_block(x)

            x = torch.cat([x, encoder_input], dim=1)
            x = decoder_block(x)
            outputs.append(x)

        # Returns list of tensors in same order as input (fine-to-coarse)
        return outputs[::-1]


#here's my custom decoder which give the next error

import timm
import torch.nn as nn
import torch
import torch.nn.functional as F
from pytorch_toolbelt.modules.decoders import UNetDecoder

class CustomTimmEncoder(nn.Module):
    def __init__(self, model_name, num_channels=2, features_only=True, pretrained=False):
        super(CustomTimmEncoder, self).__init__()
        self.encoder = timm.create_model(
            model_name,
            features_only=features_only,
            in_chans=num_channels,
            pretrained=pretrained
        )
        # Extract channels and strides
        self.channels = [f['num_chs'] for f in self.encoder.feature_info]
        self.strides = [f['reduction'] for f in self.encoder.feature_info]
        self.feature_info = self.encoder.feature_info
        
    def forward(self, x):
        features = self.encoder(x)
        for idx, feature in enumerate(features):
            print(f"[DEBUG] Feature {idx} shape: {feature.shape}")
        return features
    
class DebugUNetDecoder(UNetDecoder):
    def __init__(self, encoder_channels, decoder_channels, unet_block, upsample_block, center=False):
        super(DebugUNetDecoder, self).__init__(
            encoder_channels=encoder_channels,
            decoder_channels=decoder_channels,
            unet_block=unet_block,
            upsample_block=upsample_block,
            center=center  # Explicitly pass 'center' argument
        )
        if not hasattr(self, 'center'):
            self.center = None
        if not hasattr(self, 'up_blocks'):
            self.up_blocks = nn.ModuleList()  # Initialize an empty ModuleList or handle as needed

    def forward(self, features):
        x = features[-1]
        x = self.center(x) if self.center is not None else x  # Apply center block if it exists
        for i in reversed(range(len(features) - 1)):
            encoder_input = features[i]
            decoder_block = self.up_blocks[i]
            x = decoder_block(x, encoder_input)
            print(f"[DEBUG] Decoder stage {i}:")
            print(f"    x shape after up_block: {x.shape}")
            print(f"    encoder_input shape:       {encoder_input.shape}")
        return x
    

#finally the error

python train_multilabel.py +config=configs/multilabel_circle_net.yaml
DEBUG:hydra.core.utils:Setting JobRuntime:name=UNKNOWN_NAME
DEBUG:hydra.core.utils:Setting JobRuntime:name=train_multilabel
Dataset valid_sar_cv
        ignore_low_confidence_detections False
        ignore_low_confidence_labels     False
        filter_low_confidence_objects    False
Train dataset
        Unknown vessels 90
        Unknown fishing 11242
Valid dataset
        Unknown vessels 20
        Unknown fishing 2778
Preparing training dataset: 100%|███████████████████| 35/35 [00:00<00:00, 709.92it/s]
[DEBUG] self.box_coder in dataset.__init__: MultilabelCircleNetCoder(image_size=(512, 512), stride=4, max_objects=512, labels_encoding=circle, heatmap_encoding=umich, fixed_radius=None, ignore_index=-1, ignore_low_confidence_detections=False)
Preparing validation dataset: 100%|███████████████████| 7/7 [00:00<00:00, 589.16it/s]
Train instances 13933 dataset 35 sampler 8192
Valid instances 3655 dataset 1039
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 528, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 119, in run
    ret = run_job(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 524, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 415, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 202, in build_model
    model = super().build_model(config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/pipeline.py", line 395, in build_model
    model = self.get_model(config["model"])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 379, in get_model
    model: nn.Module = get_detection_model(model_config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/factory.py", line 27, in get_detection_model
    return instantiate(model_config, _recursive_=False)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 226, in instantiate
    return instantiate_node(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 347, in instantiate_node
    return _call_target(_target_, partial, args, kwargs, full_key)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 92, in _call_target
    return _target_(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 202, in from_config
    encoder_module = CustomTimmEncoder(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/custom_encoder.py", line 10, in __init__
    self.encoder = timm.create_model(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/_factory.py", line 117, in create_model
    model = create_fn(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/_registry.py", line 143, in _fn
    warnings.warn(f'Mapping deprecated model name {deprecated_name} to current {current_name}.', stacklevel=2)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/warnings.py", line 109, in _showwarnmsg
    sw(msg.message, msg.category, msg.filename, msg.lineno,
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 53, in warn_with_traceback
    traceback.print_stack(file=log)
/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/timm/models/_factory.py:117: UserWarning:

Mapping deprecated model name tf_efficientnet_b5_ns to current tf_efficientnet_b5.ns_jft_in1k.

[DEBUG] After adjustment, encoder.conv_stem.in_channels: 2
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 92, in _call_target
    return _target_(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 247, in from_config
    return cls(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 132, in __init__
    self.decoder = DebugUNetDecoder(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/custom_encoder.py", line 29, in __init__
    super(DebugUNetDecoder, self).__init__(
TypeError: __init__() got an unexpected keyword argument 'encoder_channels'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 528, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 524, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 415, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 202, in build_model
    model = super().build_model(config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/pipeline.py", line 395, in build_model
    model = self.get_model(config["model"])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 379, in get_model
    model: nn.Module = get_detection_model(model_config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/factory.py", line 27, in get_detection_model
    return instantiate(model_config, _recursive_=False)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 226, in instantiate
    return instantiate_node(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 347, in instantiate_node
    return _call_target(_target_, partial, args, kwargs, full_key)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/instantiate/_instantiate2.py", line 97, in _call_target
    raise InstantiationException(msg) from e
hydra.errors.InstantiationException: Error in call to target 'xview3.centernet.models.centernet_unet.CenterNetUNetModel.from_config':
TypeError("__init__() got an unexpected keyword argument 'encoder_channels'")
full_key: model
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers %