#here's my necks.py now with all upgrades 

import collections
from typing import List, Dict
#from pytorch_toolbelt.modules import FPNFuse
import torch
from torch.nn import functional as F
from torch import nn, Tensor

__all__ = ["TensorAtIndex", "PickTensors", "HyperColumnNeck"]


class TensorAtIndex(nn.Module):
    def __init__(self, channels: List[int], index: int = 0):
        super().__init__()
        self.index = index
        self.channels = channels[index]

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        return feature_maps[self.index]


class PickTensors(nn.Module):
    def __init__(self, channels: List[int], indexes: Dict[int, str]):
        super().__init__()
        self.indexes = indexes
        self.channels = channels[0]

    def forward(self, feature_maps: List[Tensor]) -> Dict[str, Tensor]:
        return collections.OrderedDict((name, feature_maps[index]) for index, name in self.indexes.items())


class HyperColumnNeck(nn.Module):
    def __init__(self, channels: List[int], dropout_rate=0.0):
        super().__init__()
        self.fuse = FPNFuse(align_corners=True)
        self.channels = sum(channels)
        self.dropout = nn.Dropout2d(p=dropout_rate, inplace=True)

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        return self.dropout(self.fuse(feature_maps))

class FPNFuse(nn.Module):
    def __init__(
        self,
        input_channels: List[int],
        output_channels: int,
        mode: str = 'bilinear',
        align_corners: bool = False,
    ):
        super().__init__()
        self.mode = mode
        self.align_corners = align_corners
        self.lateral_convs = nn.ModuleList([
            nn.Conv2d(in_ch, output_channels, kernel_size=1, bias=False)
            for in_ch in input_channels
        ])
        self.fuse_conv = nn.Conv2d(output_channels, output_channels, kernel_size=3, padding=1)
        
    def forward(self, feature_maps: List[Tensor]):
        print(f"[DEBUG] feature_maps in FPNFuse: type={type(feature_maps)}")
        valid_feature_maps = []
        dst_size = None
        for idx, f in enumerate(feature_maps):
            if f.dim() == 4:
                print(f"[DEBUG] Valid feature_maps[{idx}] shape: {f.shape}")
                # Apply lateral convolution
                processed_feature = self.lateral_convs[idx](f)
                valid_feature_maps.append(processed_feature)
                if dst_size is None:
                    dst_size = processed_feature.size()[-2:]
            else:
                print(f"[WARNING] Skipping feature_maps[{idx}] with shape {f.shape}")
        if dst_size is None:
            raise ValueError("No valid feature maps with spatial dimensions found.")
        layers = []
        for f in valid_feature_maps:
            upsampled_feature = F.interpolate(f, size=dst_size, mode=self.mode, align_corners=self.align_corners)
            layers.append(upsampled_feature)
        x = torch.cat(layers, dim=1)
        x = self.fuse_conv(x)
        return x

class FPNNeck(nn.Module):
    def __init__(self, channels: List[int], out_channels: int):
        super().__init__()
        self.fpn = FPNFuse(channels, out_channels, mode='bilinear', align_corners=False)
        self.out_conv = nn.Conv2d(out_channels * len(channels), out_channels, kernel_size=1)

    def forward(self, feature_maps: List[Tensor]) -> Tensor:
        fpn_output = self.fpn(feature_maps)
        print(f"[DEBUG] fpn_output shape: {fpn_output.shape}")
        return self.out_conv(fpn_output)

#here's my custom_encoder.py

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(nn.Module):
    def __init__(self, encoder_channels, decoder_channels, unet_block, upsample_block, center=False):
        super(DebugUNetDecoder, self).__init__()
        self.center = center
        self.num_stages = len(decoder_channels)
        self.up_samples = nn.ModuleList()
        self.unet_blocks = nn.ModuleList()

        if self.center:
            self.center_block = unet_block(encoder_channels[-1], encoder_channels[-1])
        else:
            self.center_block = nn.Identity()

        for i in range(self.num_stages):
            self.up_samples.append(upsample_block())

            if i == 0:
                in_channels = encoder_channels[-1] + encoder_channels[-2]
            else:
                in_channels = decoder_channels[i - 1] + encoder_channels[-(i + 2)]
            out_channels = decoder_channels[i]

            self.unet_blocks.append(unet_block(in_channels, out_channels))

    def forward(self, features):
        x = features[-1]
        x = self.center_block(x)
        num_features = len(features)
        print(f"[DEBUG] Number of encoder features: {num_features}")

        decoded_features = []
        for i in range(self.num_stages):
            print(f"[DEBUG] Decoder stage {i}")
            x = self.up_samples[i](x)

            encoder_index = num_features - i - 2
            if encoder_index >= 0:
                encoder_feature = features[encoder_index]
                print(f"Using encoder feature at index {encoder_index}")

                if x.size()[2:] != encoder_feature.size()[2:]:
                    encoder_feature = F.interpolate(encoder_feature, size=x.size()[2:], mode='nearest')
                    print(f"encoder_feature.shape after resizing at stage {i}: {encoder_feature.shape}")

                x = torch.cat([x, encoder_feature], dim=1)
                print(f"x.shape after concat at stage {i}: {x.shape}")
            else:
                print(f"No encoder feature for decoder stage {i}, skipping concatenation")

            x = self.unet_blocks[i](x)
            decoded_features.append(x)

        return decoded_features


#here's the error log the important parts with the explanations 

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, 164.16it/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, 547.91it/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 198, 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
[DEBUG] model.box_coder after initialization: MultilabelCircleNetCoder(image_size=[768, 768], stride=1, max_objects=512, labels_encoding=circle, heatmap_encoding=umich, fixed_radius=None, ignore_index=-1, ignore_low_confidence_detections=False)
BoxCoder
MultilabelCircleNetCoder(image_size=[768, 768], stride=1, max_objects=512, labels_encoding=circle, heatmap_encoding=umich, fixed_radius=None, ignore_index=-1, ignore_low_confidence_detections=False)
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.session:Could not import boto3, continuing with reduced functionality.
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x3333322e0>
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x32bf11280>
...

DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x305946220>
[DEBUG] Channel 0 shape: (512, 512)
[DEBUG] Channel 1 shape: (512, 512)
[DEBUG] feature_maps in FPNFuse: type=<class 'list'>
[DEBUG] Valid feature_maps[0] shape: torch.Size([8, 512, 32, 32])
DEBUG:rasterio.env:Entering env context: <rasterio.env.Env object at 0x324ecfa60>
...

DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x3059462e0>
DEBUG:rasterio._base:Sharing flag: 0
[DEBUG] Valid feature_maps[1] shape: torch.Size([8, 384, 64, 64])
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x33803edf0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x33803ee80> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x33803ee80>.
...

DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x33803edf0>
DEBUG:rasterio._base:Sharing flag: 0
[DEBUG] Valid feature_maps[2] shape: torch.Size([8, 256, 128, 128])
[DEBUG] Valid feature_maps[3] shape: torch.Size([8, 128, 256, 256])
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x33142f310>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x33142fb80> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x33142fb80>.
...
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x33fec2400>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=10240, row_off=4096, width=2048, height=2048)
DEBUG:rasterio._io:IO window xoff=10240.0 yoff=4096.0 width=2048.0 height=2048.0
DEBUG:rasterio._base:Nodata success: 1, Nodata value: -32768.000000
DEBUG:rasterio._base:Dataset <open DatasetReader name='/Volumes/Crucial/xview3_dataset/full_dataset/validation/4a97701b4bd81bf7v/VV_dB.tif' mode='r'> is started.
...

Decoded features length: 4
[DEBUG] decoded_features[0] shape: torch.Size([8, 512, 32, 32])
[DEBUG] decoded_features[1] shape: torch.Size([8, 384, 64, 64])
[DEBUG] decoded_features[2] shape: torch.Size([8, 256, 128, 128])
[DEBUG] decoded_features[3] shape: torch.Size([8, 128, 256, 256])
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x305946220>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x309c45e50> options
...

Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  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 206, in build_model
    model = nas.search(config.nas.num_iterations, loaders['train'], loaders['valid'])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 36, in search
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 36, in <listcomp>
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 112, in evaluate_model
    outputs = model(inputs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 169, in forward
    neck_features = self.neck(decoded_features)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/necks.py", line 89, in forward
    fpn_output = self.fpn(feature_maps)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/necks.py", line 79, in forward
    x = self.fuse_conv(x)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 458, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/nn/modules/conv.py", line 454, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Given groups=1, weight of size [256, 256, 3, 3], expected input[8, 1024, 32, 32] to have 256 channels, but got 1024 channels instead

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers %