to be able to fix the project correclty you need to review deeply my config and encoder, decoder, fan , fbn , swin transformer

here are

class CustomTimmEncoder(nn.Module):
def __init__(
self,
model_name: str = "efficientnetv2_rw_s",
num_channels: int = 2,
features_only: bool = True,
pretrained: bool = False,
activation: str = "silu",
drop_rate: float = 0.0,
drop_path_rate: float = 0.0,
image_size: Tuple[int, int] = (224, 224),
):
super(CustomTimmEncoder, self).__init__()
self.model_name = model_name # Store model name for conditional handling
self.encoder = timm.create_model(
model_name,
features_only=features_only,
in_chans=num_channels,
pretrained=pretrained,
act_layer=activation,
drop_rate=drop_rate,
drop_path_rate=drop_path_rate,
)
# Extract channels and strides
# Extract actual feature channels and strides
# In CustomTimmEncoder
self.feature_info = self.encoder.feature_info
self.channels = [f['num_chs'] for f in self.feature_info]
#self.channels = [f['num_chs'] for f in self.encoder.feature_info]
self.strides = [f['reduction'] for f in self.feature_info]
self.feature_info = self.feature_info
self.image_size = image_size
self.embed_dim = self.channels[-1]
self.depths = [2, 2, 6, 2]
self.num_heads = [3, 6, 12, 24]
self.window_size = 7
self.mlp_ratio = 4.
self.qkv_bias = True
self.qk_scale = None
self.drop_rate = drop_rate
self.attn_drop_rate = 0.0
self.drop_path_rate = drop_path_rate
self.norm_layer = nn.LayerNorm
self.use_dropout = True
self.use_custom_pos_embed = False
self.transformer_output_channels = self.channels[-1]
self.output_stride = 32


# Conditional handling based on encoder type
if "swin" in model_name.lower() or "transformer" in model_name.lower():
# Assuming transformer-based encoder has 'patch_embed'
if hasattr(self.encoder, 'patch_embed'):
self.conv_stem = self.encoder.patch_embed.proj
else:
raise AttributeError(f"The encoder '{model_name}' does not have 'patch_embed' attribute.")
else:
# For CNN-based encoders like EfficientNet
if hasattr(self.encoder, 'conv_stem'):
self.conv_stem = self.encoder.conv_stem
elif hasattr(self.encoder, 'conv1'):
self.conv_stem = self.encoder.conv1 # Example for some other CNNs
else:
raise AttributeError(f"The encoder '{model_name}' does not have a recognized stem convolution layer.")


def forward(self, x):
encoder_features = self.encoder(x)
if not isinstance(encoder_features, list):
encoder_features = [encoder_features]
logger.debug(f"Encoder features shapes: {[f.shape for f in encoder_features]}")


for idx, encoder_feature in enumerate(encoder_features):
if encoder_feature.dim() == 3:
# Missing batch dimension, add it
encoder_feature = encoder_feature.unsqueeze(0)
encoder_features[idx] = encoder_feature
print(f"[WARNING] Feature {idx} missing batch dimension, adding it.")
elif encoder_feature.dim() < 3:
raise ValueError(f"Feature {idx} has unexpected number of dimensions: {encoder_feature.dim()}")
print(f"[DEBUG] Feature {idx} shape: {encoder_feature.shape}") # Existing debug


# Additional debug to ensure all features have batch dimension
for idx, encoder_feature in enumerate(encoder_features):
if encoder_feature.dim() != 4:
print(f"[ERROR] Feature {idx} does not have 4 dimensions after processing: {encoder_feature.dim()}")
raise ValueError(f"Feature {idx} has incorrect number of dimensions: {encoder_feature.dim()}")


return encoder_features



def change_input_channels(self, input_channels: int, mode="auto", **kwargs):
if isinstance(self.conv_stem, Conv2dSame):
self.conv_stem = make_n_channel_input_conv2d_same(self.conv_stem, input_channels, mode, **kwargs)
else:
self.conv_stem = make_n_channel_input(self.conv_stem, input_channels, mode, **kwargs)
return self


class DebugUnetDecoder(nn.Module):
def __init__(
self,
encoder_channels,
decoder_channels,
unet_block,
upsample_block=UpsampleBlock,
center=False,
activation='relu'
):
super(DebugUnetDecoder, self).__init__()
self.center = center
self.encoder_channels = encoder_channels
self.decoder_channels = decoder_channels
self.num_stages = len(decoder_channels)
self.up_samples = nn.ModuleList()
self.unet_blocks = nn.ModuleList()
for i in range(self.num_stages):
if i == 0:
in_channels = encoder_channels[-1]
else:
in_channels = decoder_channels[i-1]
out_channels = decoder_channels[i]
self.up_samples.append(
upsample_block(
in_channels=in_channels,
out_channels=out_channels
)
)
if i < len(encoder_channels) - 1:
unet_in_channels = out_channels + encoder_channels[-(i+2)]
else:
unet_in_channels = out_channels
self.unet_blocks.append(
unet_block(
in_channels=unet_in_channels,
out_channels=out_channels,
activation=activation
)
)


def forward(self, encoder_features):
if not isinstance(encoder_features, list):
raise TypeError(f"Decoder expected encoder_features to be a list, but got {type(encoder_features)}")
x = encoder_features[-1]
print(f"Initial x.shape: {x.shape}")
decoded_features = []
for i in range(self.num_stages):
print(f"\nDecoder Stage {i}")
print(f"x.shape before upsampling: {x.shape}")
# Upsample
x = self.up_samples[i](x)
print(f"x.shape after upsampling: {x.shape}")
# Concatenate with encoder feature
if i < len(encoder_features) - 1:
encoder_feat = encoder_features[-(i+2)]
print(f"Encoder Feature {i} shape: {encoder_feat.shape}")
if x.size()[2:] != encoder_feat.size()[2:]:
encoder_feat = F.interpolate(encoder_feat, size=x.size()[2:], mode='nearest')
print(f"Resized encoder_feat to {x.size()[2:]}")
x = torch.cat([x, encoder_feat], dim=1)
print(f"x.shape after concatenation: {x.shape}")
# Pass through UNetBlock
print(f"Passing through UNetBlock with in_channels={x.shape[1]}")
x = self.unet_blocks[i](x)
print(f"x.shape after UNetBlock: {x.shape}")
decoded_features.append(x)
return decoded_features


class DebugSwinTransformer(SwinTransformer):
def __init__(
self,
img_size: Union[int, Tuple[int, int]],
patch_size: int = 4,
in_chans: int = 2,
embed_dim: int = 96,
depths: list = [2, 2, 6, 2],
num_heads: list = [3, 6, 12, 24],
window_size: Tuple[int, int] = (4, 4),
mlp_ratio: float = 4.0,
qkv_bias: bool = True,
qk_scale: Optional[float] = None,
drop_rate: float = 0.0,
attn_drop_rate: float = 0.0,
drop_path_rate: float = 0.1,
norm_layer: Optional[str] = 'torch.nn.LayerNorm', # Fully qualified class path
ape: bool = False,
patch_norm: bool = True,
use_checkpoint: bool = False,
):
"""
DebugSwinTransformer extends the SwinTransformer to include debug statements.


Args:
norm_layer (str): Fully qualified class path for normalization layer (e.g., 'torch.nn.LayerNorm').
All other parameters are the same as SwinTransformer.
"""
# Enhanced assertion for window_size
if isinstance(window_size, int):
assert window_size > 0, f"Invalid window_size: {window_size}"
self.window_size = (window_size, window_size)
elif isinstance(window_size, (tuple, list)):
assert len(window_size) == 2, f"window_size must have two dimensions, got {len(window_size)}"
assert all(isinstance(w, int) and w > 0 for w in window_size), f"Invalid window_size elements: {window_size}"
self.window_size = tuple(window_size)
else:
raise AssertionError(f"Invalid window_size type: {type(window_size)}. Expected int or tuple/list of two ints.")


# Convert norm_layer string to actual class
try:
norm_layer_class = get_class(norm_layer)
except ImportError as e:
logger.error(f"Error loading norm_layer: {e}")
raise


# Initialize the parent SwinTransformer with validated window_size and norm_layer_class
super().__init__(
img_size=img_size,
patch_size=patch_size,
in_chans=in_chans,
embed_dim=embed_dim,
depths=depths,
num_heads=num_heads,
window_size=self.window_size,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop_rate=drop_rate,
attn_drop_rate=attn_drop_rate,
drop_path_rate=drop_path_rate,
norm_layer=norm_layer_class, # Pass the class, not the string
ape=ape,
patch_norm=patch_norm,
use_checkpoint=use_checkpoint,
)


def forward_features(self, x):
logger.debug(f"SwinTransformer.forward_features: window_size={self.window_size}, input_shape={x.shape}")
return super().forward_features(x)


def _attn(self, x):
logger.debug(f"SwinTransformer._attn: window_size={self.window_size}")
# Add assertion to prevent invalid window_size during attention
if isinstance(self.window_size, int):
assert self.window_size > 0, f"Invalid window_size during _attn: {self.window_size}"
else:
assert all(isinstance(w, int) and w > 0 for w in self.window_size), f"Invalid window_size during _attn: {self.window_size}"
return super()._attn(x)



class SwinTransformerEncoder(nn.Module):
def __init__(
self,
img_size: Union[int, Tuple[int, int]],
embed_dim: int = 96,
depths: list = [2, 2, 6, 2],
num_heads: list = [3, 6, 12, 24],
window_size: Tuple[int, int] = (4, 4),
mlp_ratio: float = 4.0,
qkv_bias: bool = True,
qk_scale: Optional[float] = None,
drop_rate: float = 0.0,
attn_drop_rate: float = 0.0,
drop_path_rate: float = 0.1,
norm_layer: Optional[str] = 'torch.nn.LayerNorm', # Fully qualified class path
in_chans: int = 2, # Number of input channels
use_dropout: bool = False,
use_custom_pos_embed: bool = False,
):
"""
Swin Transformer Encoder with enhanced features.


Args:
norm_layer (str): Fully qualified class path for normalization layer (e.g., 'torch.nn.LayerNorm').
All other parameters are the same as before.
"""
super().__init__()
# Validate img_size
assert (
(isinstance(img_size, int) and img_size > 0) or
(isinstance(img_size, (list, tuple)) and
len(img_size) == 2 and
all(isinstance(x, int) and x > 0 for x in img_size))
), f"Invalid img_size={img_size}. Must be a positive int or tuple/list of two positive ints."


self.window_size = window_size


# Convert norm_layer string to actual class
try:
norm_layer_class = get_class(norm_layer)
except ImportError as e:
logger.error(f"Error loading norm_layer: {e}")
raise


self.swin = DebugSwinTransformer(
img_size=img_size, # Now can be int or tuple/list
patch_size=4,
in_chans=in_chans,
embed_dim=embed_dim,
depths=depths,
num_heads=num_heads,
window_size=window_size,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
drop_rate=drop_rate,
attn_drop_rate=attn_drop_rate,
drop_path_rate=drop_path_rate,
norm_layer=norm_layer_class, # Pass the class, not the string
ape=False,
patch_norm=True,
use_checkpoint=False
)


# Initialize advanced weight initialization
self._init_weights()


# Integrate TorchMetrics for monitoring
self.accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=2, average='macro')
self.topk_accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=2, average='macro')


# Optional dropout layer after the transformer
self.dropout = nn.Dropout(drop_rate) if use_dropout else nn.Identity()


self.use_custom_pos_embed = use_custom_pos_embed
if use_custom_pos_embed:
# Adjust positional embedding based on img_size
if isinstance(img_size, int):
pos_embed_size = (img_size // 4, img_size // 4)
else:
pos_embed_size = (img_size[0] // 4, img_size[1] // 4)
self.pos_embed = nn.Parameter(torch.zeros(1, embed_dim, *pos_embed_size))
nn.init.kaiming_normal_(self.pos_embed, mode='fan_out', nonlinearity='relu')


def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=.02)
if m.bias is not None:
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
elif isinstance(m, nn.Conv2d):
if m.weight.dim() > 1:
nn.init.xavier_uniform_(m.weight)
else:
nn.init.normal_(m.weight)
if m.bias is not None:
nn.init.constant_(m.bias, 0)


def forward(self, x):
"""
Forward pass through the Swin Transformer Encoder.


Args:
x (Tensor): Input tensor of shape [B, C, H, W].


Returns:
Tensor: Extracted features from the encoder.
if self.use_custom_pos_embed:
x = x + self.pos_embed
# Use forward_features to retain spatial dimensions
features = self.swin.forward_features(x)
features = self.dropout(features)
return features
"""
def forward(self, x):
features = self.swin.forward_features(x)
features = self.dropout(features)
# Reshape the output to match the expected shape
B, C, H, W = x.shape
features = features.permute(0, 2, 1).reshape(B, -1, H // self.swin.patch_size, W // self.swin.patch_size)
return features


# Enhanced FPNNeck using EnhancedFPNFuse
class FPNNeck(nn.Module):
def __init__(
self,
channels: List[int],
out_channels: int,
use_cbam: bool = False,
use_dropblock: bool = False,
dropblock_prob: float = 0.1,
):
super(FPNNeck, self).__init__()
self.fpn = EnhancedFPNFuse(
input_channels=channels,
output_channels=out_channels,
mode='bilinear',
align_corners=False,
use_cbam=use_cbam,
use_dropblock=use_dropblock,
dropblock_prob=dropblock_prob,
)
self.out_conv = nn.Conv2d(out_channels, out_channels, kernel_size=1)
def forward(self, feature_maps: List[Tensor]) -> Tensor:
x = self.fpn(feature_maps)
x = self.out_conv(x)
return x


# HyperColumnNeck with CBAM and DropBlock
class HyperColumnNeck(nn.Module):
def __init__(self, channels: List[int], dropout_rate=0.0, use_cbam: bool = False, use_dropblock: bool = False, dropblock_prob: float = 0.1):
super(HyperColumnNeck, self).__init__()
self.fuse = EnhancedFPNFuse(
input_channels=channels,
output_channels=sum(channels),
align_corners=True,
use_cbam=use_cbam,
use_dropblock=use_dropblock,
dropblock_prob=dropblock_prob,
)
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))


# Feature Aggregation Network with Additional Enhancements
class FeatureAggregationNetwork(nn.Module):
"""
Advanced Feature Aggregation Network that fuses multiple encoder feature maps using attention mechanisms.
Args:
input_channels (List[int]): List of channel dimensions for each input feature map.
output_channels (int): Number of output channels.
upsample_mode (str): Mode for upsampling ('nearest', 'bilinear', 'bicubic').
use_attention (bool): Whether to use attention mechanisms for feature fusion.
fusion_method (str): Method for feature fusion ('add', 'concat', 'weighted_sum').
"""
def __init__(
self,
input_channels: List[int],
out_channels: int,
num_layers: int,
upsample_mode: str = 'bilinear',
use_attention: bool = True,
fusion_method: str = 'weighted_sum'
):
super(FeatureAggregationNetwork, self).__init__()
self.input_channels = input_channels
self.num_features = len(input_channels)
self.out_channels = out_channels
self.upsample_mode = upsample_mode
self.use_attention = use_attention
self.fusion_method = fusion_method
# Channel alignment layer
#self.channel_align = nn.Conv2d(input_channels[0], out_channels, kernel_size=1)
"""
# Channel alignment layers
self.channel_align_layers = nn.ModuleList([
nn.Conv2d(in_ch, out_channels, kernel_size=1, bias=False) for in_ch in input_channels
])
"""
# Channel alignment layers
self.channel_align_layers = nn.ModuleList([
nn.Conv2d(in_ch, out_channels, kernel_size=1, bias=False) for in_ch in input_channels
])
# Attention modules (if used)
if self.use_attention:
self.attention_layers = nn.ModuleList([
SelfAttentionModule(out_channels)
for _ in range(self.num_features)
])
# Learnable weights for fusion (if using weighted_sum)
if self.fusion_method == 'weighted_sum':
self.fusion_weights = nn.Parameter(torch.ones(self.num_features), requires_grad=True)
# Output convolution
self.output_conv = nn.Sequential(
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
def forward(self, feature_maps: Union[Tensor, List[Tensor]]) -> Tensor:
"""
Forward pass for FeatureAggregationNetwork.
Args:
feature_maps (List[Tensor]): List of feature maps from different encoder stages.
Returns:
Tensor: Aggregated feature map.
"""
if isinstance(feature_maps, Tensor):
if feature_maps.dim() == 3:
# Possibly [channels, height, width], need to add batch dimension
feature_maps = feature_maps.unsqueeze(0)
elif feature_maps.dim() == 4:
# Already has batch dimension
pass
else:
raise ValueError(f"Expected feature_maps with 3 or 4 dimensions, got {feature_maps.dim()} dimensions.")
feature_maps = [feature_maps]
elif isinstance(feature_maps, list):
# Ensure all feature maps have the correct dimensions
for i, fm in enumerate(feature_maps):
if fm.dim() == 3:
feature_maps[i] = fm.unsqueeze(0)
elif fm.dim() != 4:
raise ValueError(f"Feature map at index {i} has {fm.dim()} dimensions, expected 3 or 4.")
else:
raise TypeError(f"Expected feature_maps to be a Tensor or list, but got {type(feature_maps)}")
# Determine the target spatial size (smallest height and width)
target_height = min([fm.size(2) for fm in feature_maps])
target_width = min([fm.size(3) for fm in feature_maps])
target_size = (target_height, target_width)
#aligned_features = []
aligned_features = [self.input_conv(fm) for fm in feature_maps]


for idx, fm in enumerate(feature_maps):
# Spatial alignment
if fm.size(2) != target_height or fm.size(3) != target_width:
fm = F.interpolate(fm, size=target_size, mode=self.upsample_mode, align_corners=False)
logger.debug(f"Resized feature_map[{idx}] to {target_size}")
# Channel alignment
fm = self.channel_align_layers[idx](fm)
# Attention
if self.use_attention:
fm = self.attention_layers[idx](fm)
aligned_features.append(fm)
# Feature fusion
if self.fusion_method == 'add':
fused_feature = sum(aligned_features)
elif self.fusion_method == 'concat':
fused_feature = torch.cat(aligned_features, dim=1)
# Adjust the number of channels to output_channels
fused_feature = nn.Conv2d(
fused_feature.size(1), self.out_channels, kernel_size=1, bias=False
)(fused_feature)
elif self.fusion_method == 'weighted_sum':
# Apply softmax to fusion weights
weights = F.softmax(self.fusion_weights, dim=0)
fused_feature = sum(w * f for w, f in zip(weights, aligned_features))
else:
raise ValueError(f"Unsupported fusion method: {self.fusion_method}")
# Output convolution
output = self.output_conv(fused_feature)
logger.debug(f"FAN output shape: {output.shape}") # Should be 4D
return output

class UnetBlock(nn.Module):
def __init__(self, in_channels, out_channels, activation='relu', use_attention=False, use_dropblock=False, dropblock_prob=0.1):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
self.bn1 = nn.GroupNorm(16, out_channels)
self.activation1 = get_activation_block(activation)
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
self.bn2 = nn.GroupNorm(16, out_channels)
self.activation2 = get_activation_block(activation)
if use_attention:
self.attention = AttentionCBAM(out_channels)
if use_dropblock:
from pytorch_toolbelt.modules import DropBlock2D
self.dropblock = DropBlock2D(block_size=7, drop_prob=dropblock_prob)
# Residual connection
if in_channels != out_channels:
self.residual = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
nn.GroupNorm(16, out_channels)
)
else:
self.residual = nn.Identity()
def forward(self, x):
print(f"[DEBUG] UnetBlock input shape: {x.shape}")
print(f"[DEBUG] Expected input channels: {self.in_channels}, Actual input channels: {x.shape[1]}")
residual = self.residual(x)
print(f"[DEBUG] Residual shape after conv: {residual.shape}")
out = self.conv1(x)
out = self.bn1(out)
out = self.activation1(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.activation2(out)
if hasattr(self, 'attention'):
out = self.attention(out)
if hasattr(self, 'dropblock'):
out = self.dropblock(out)
print(f"[DEBUG] Out shape before adding residual: {out.shape}")
print(f"[DEBUG] Residual shape: {residual.shape}")
# Ensure residual and out have the same shape
if residual.shape != out.shape:
raise ValueError(f"Shape mismatch: residual shape {residual.shape} vs out shape {out.shape}")
out += residual
print(f"[DEBUG] Final output shape: {out.shape}")
return out

config
defaults:
- _self_
- dataset: full_4fold
- augs: flips_light
- sampler: balanced
- normalization: default
- loss: centernet/rfl_bce_mse
- model: centernet/b5_unet_s2
- optimizer: adamw
- scheduler: cos
- train: default
- runner: default
- ema: disabled
- torch: default


slug: b5_unet_s2
activation: silu
num_channels: ${dataset.num_channels}
num_extra_blocks: 0


model:
_target_: xview3.centernet.models.centernet_unet.EnhancedCenterNetUNet
slug: b5_unet_s2
weights: IMAGENET1K_V1
pretrained: True
output_stride: 32
image_size: [1024, 1024]
window_size:
- 4
- 4
num_layers: [2, 5]
num_filters: [64, 128, 256]
kernel_size: [3, 5]
activation: [relu, leaky_relu, silu]
dropout_rate: [0.0, 0.5]
dropout_rate_max: 0.5
optimizer: adamw
learning_rate: 0.001


transformer:
enabled: true
embed_dim: 96
depths: [2, 2, 6, 2]
num_heads: [3, 6, 12, 24]
window_size:
- 4
- 4
mlp_ratio: 4.0
qkv_bias: true
qk_scale: null
drop_rate: 0.0
attn_drop_rate: 0.0
drop_path_rate: 0.1
norm_layer: torch.nn.LayerNorm
use_dropout: false
use_custom_pos_embed: false


decoder:
_target_: xview3.centernet.models.custom_encoder.DebugUnetDecoder
encoder_channels: [32, 64, 128, 256] #[32, 24, 40, 112, 320]
decoder_channels: [512, 384, 256, 128, 64]
unet_block:
_target_: xview3.centernet.models.unet_blocks.UnetBlock
_partial_: true
upsample_block: xview3.centernet.models.custom_encoder.UpsampleBlock
#scale_factor: 2
#mode: bilinear
center: False # Changed from 'bilinear' to 'nearest'
activation: silu


encoder:
_target_: pytorch_toolbelt.modules.encoders.TimmEfficientNetV2
model_name: efficientnetv2_rw_s
pretrained: True #features_only: True
activation: silu
drop_rate: 0.1
drop_path_rate: 0.1
layers:
- 1 # 4
- 2 # 8
- 3 # 16
- 4 # 32


nas:
enabled: true
search_space:
num_layers: [2, 3, 4]
num_filters: [64, 128, 192, 256]
kernel_size: [3, 5]
activation: [relu, leaky_relu, silu]
dropout_rate: [0.0, 0.5]
dropout_rate_max: 0.5
num_trials: 50
metric: "val_acc"
direction: "maximize"


num_channels: ${dataset.num_channels}
num_classes: 2
anchors: [[10, 13], [16, 30], [33, 23]]
input_upsample_factor: 1
neck_out_channels: 256


box_coder:
_target_: xview3.centernet.bboxer.multilabel_circle_coder.MultilabelCircleNetCoder
image_size:
- 1024
- 1024
max_objects: 512
heatmap_encoding: umich
labels_encoding: circle
ignore_value:
_target_: xview3.ignore_value
fixed_radius: 3
labels_radius: 1
ignore_low_confidence_detections: ${dataset.ignore_low_confidence_detections}


head:
decoupled_head:
_target_: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
channels: 256
classifier_dim: 128
objectness_dim: 128
size_dim: 128
offset_dim: 16
dropout_rate: 0.1
activation: silu
num_blocks: 3
yolo_head:
_target_: xview3.centernet.models.heads.yolo_head.YOLOHead
num_classes: 2
anchors: [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119]]
input_image_size: [1024, 1024]
stride: 32
num_anchors_per_scale: 3


fan:
_target_: xview3.centernet.models.necks.FeatureAggregationNetwork
input_channels: [272] #[32, 48, 64, 128] #in_channels: 272 #out_channels: 256 #num_layers: 2
out_channels: 256
num_layers: 2
upsample_mode: bilinear
use_attention: True
fusion_method: weighted_sum


expected_fan_input_channels: 256


# Per-Channel Normalization
normalization:
channels:
vv:
_target_: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
vh:
_target_: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
bathymetry:
_target_: xview3.CubicRootNormalization
wind_direction:
_target_: albumentations.NoOp
wind_speed:
_target_: albumentations.NoOp
mask:
_target_: albumentations.NoOp
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]


# Dataset Configuration
dataset:
normalization:
channels:
vv:
_target_: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
vh:
_target_: xview3.SigmoidNormalization
midpoint: -20.0
temperature: 0.18
bathymetry:
_target_: xview3.CubicRootNormalization
wind_direction:
_target_: albumentations.NoOp
wind_speed:
_target_: albumentations.NoOp
mask:
_target_: albumentations.NoOp
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]
slug: full_4fold
splitter: full
fold: 0
num_folds: 4
# Added the following parameters to address references in box_coder
ignore_low_confidence_detections: true
ignore_low_confidence_labels: false
filter_low_confidence_objects: False
num_channels: 2
input_channels:
- vh
- vv


# DDP Stuff
local_rank: ${oc.env:LOCAL_RANK, 0}
world_size: ${oc.env:WORLD_SIZE, 0}


experiment: null


seed: 42


deterministic: False
detect_anomaly: False
find_unused: False


transfer: null
checkpoint: null


inference:
channels_last: False
use_traced_model: True
tile_size: 2048
tile_step: 1536
accumulate_on_gpu: False
batch_size: 1
fp16: True


# Torch Configuration
torch:
amp:
enabled: false
cudnn:
allow_tf32: True
cuda_allow_tf32: True
detect_anomaly: False
deterministic: False
benchmark: True


# Inference and Training Enhancements
activation_after: "after_model"
tta: "d4"
with_offset: True
yolo_decoders: []


# Validation Configuration
valid:
image_size: [2048, 2048]
batch_size: 16
num_workers: 2


# Training Configuration
train:
slug: default_training
epochs: 4
output_stride: 32
early_stopping:
enabled: true
patience: 2
mode: "max" # Options: 'min' for loss-based, 'max' for accuracy-based
metric: "loss" # The metric to monitor for early stopping
image_size: [512, 512]
batch_size: 16
num_workers: 2
num_samples: 8192
persistent_workers: False



model2:
_target_: xview3.centernet.models.centernet_unet.EnhancedCenterNetUNet
slug: b5_unet_s2
weights: IMAGENET1K_V1
pretrained: True
output_stride: 32
image_size: [1024, 1024]
window_size:
- 4
- 4
num_layers: ${model2.num_layers} # Ensure these are set up for overrides
num_filters: ${model2.num_filters}
kernel_size: ${model2.kernel_size}
activation: ${model2.activation}
dropout_rate: ${model2.dropout_rate}
dropout_rate_max: 0.5
optimizer: ${model2.optimizer}
learning_rate: ${model2.learning_rate}


transformer:
enabled: true
embed_dim: 96
depths: [2, 2, 6, 2]
num_heads: [3, 6, 12, 24]
window_size:
- 4
- 4
mlp_ratio: 4.0
qkv_bias: true
qk_scale: null
drop_rate: 0.0
attn_drop_rate: 0.0
drop_path_rate: 0.1
norm_layer: torch.nn.LayerNorm
use_dropout: false
use_custom_pos_embed: false


decoder:
_target_: xview3.centernet.models.custom_encoder.DebugUNetDecoder
encoder_channels: [32, 24, 40, 112, 320]
channels: [512, 384, 256, 128, 64]
unet_block:
_target_: xview3.centernet.models.unet_blocks.UnetBlock
_partial_: true
upsample_block:
_target_: torch.nn.Upsample
scale_factor: 2
mode: bilinear


encoder:
_target_: pytorch_toolbelt.modules.encoders.TimmEfficientNetV2
model_name: efficientnetv2_rw_s
pretrained: True
activation: silu
drop_rate: 0.1
drop_path_rate: 0.1
layers:
- 1 # 4
- 2 # 8
- 3 # 16
- 4 # 32


nas:
enabled: true
search_space:
num_layers: [2, 5]
num_filters: [64, 128, 256]
kernel_size: [3, 5]
activation: [relu, leaky_relu, silu]
dropout_rate: [0.0, 0.5]
dropout_rate_max: 0.5
num_trials: 50
metric: "val_loss"
direction: "minimize"


num_channels: ${dataset.num_channels}
num_classes: 2
anchors: [[10, 13], [16, 30], [33, 23]]
input_upsample_factor: 1
neck_out_channels: 256


box_coder:
_target_: xview3.centernet.bboxer.multilabel_circle_coder.MultilabelCircleNetCoder
image_size:
- 1024
- 1024
max_objects: 512
heatmap_encoding: umich
labels_encoding: circle
ignore_value:
_target_: xview3.ignore_value
fixed_radius: 3
labels_radius: 1
ignore_low_confidence_detections: ${dataset.ignore_low_confidence_detections}


head:
decoupled_head:
_target_: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
channels: 256
classifier_dim: 128
objectness_dim: 128
size_dim: 128
offset_dim: 16
dropout_rate: 0.1
activation: silu
num_blocks: 3
yolo_head:
_target_: xview3.centernet.models.heads.yolo_head.YOLOHead
num_classes: 2
anchors: [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119]]
input_image_size: [1024, 1024]
stride: 32
num_anchors_per_scale: 3


fan:
_target_: xview3.centernet.models.necks.FeatureAggregationNetwork
input_channels: [272] #[32, 48, 64, 128] #in_channels: 272 #out_channels: 256 #num_layers: 2
out_channels: 256
num_layers: 2
upsample_mode: bilinear
use_attention: True
fusion_method: weighted_sum