you are a super advanced machine learning software engineer that able to deeply understand and debug complex projects and help me enhance creatively and fix issues using most advaced methods, libraries and techniques to have a great performance and accuracy in my project which is using TIFF SAR images to detect vessels and if fishing or not, you do this smartly without removing, ruining or affecting any function , method or file structure!
in this project it used CenterNetUNet model and we upgraded it to integrate NAS, YOLO, FBN, FAN, Optuna and SwinTransformer

I think that @nas.py @train_multilabel3.py i think loading config.model and updating the config by nas is confused 

##the error


DEBUG:rasterio._env:GDAL_DATA found in environment.
DEBUG:rasterio._env:PROJ_DATA found in environment.
DEBUG:rasterio._env:Started GDALEnv: self=<rasterio._env.GDALEnv object at 0x360e00b20>.
DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x360e00790>
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=16802, row_off=15456, width=512, height=512)
DEBUG:rasterio._io:IO window xoff=16802.0 yoff=15456.0 width=512.0 height=512.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x360e00790>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x360e00b20> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x360e00b20>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x360e00790>
[DEBUG] Channel 0 shape: (512, 512)
[DEBUG] Channel 1 shape: (512, 512)
(84.0, 176.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
(276.0, 268.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
[2024-10-05 17:11:32,296][__main__][INFO] - Building model with config:
[2024-10-05 17:11:32,302][__main__][INFO] - _target_: xview3.centernet.models.CenterNetUNetModel.from_config
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: null
  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
  - 2
  - 3
  - 4
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_acc
  direction: maximize
num_channels: 2
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:
  enabled: true
  in_channels:
  - 32
  - 64
  - 128
  - 256
  out_channels: 256
  num_layers: 2
config:
  slug: b5_unet_s2
  activation: silu
  num_channels: ${dataset.num_channels}
  num_extra_blocks: 0
  encoder:
    _target_: pytorch_toolbelt.modules.encoders.TimmB5Encoder
    pretrained: true
    activation: silu
    layers:
    - 1
    - 2
    - 3
    - 4
  decoder:
    upsample_type: shuffle
    block_type: IRBlock
    channels:
    - 256
    - 384
    - 512
  head:
    _target_: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
    classifier_dim: 128
    objectness_dim: 128
    size_dim: 64
    offset_dim: 16
    dropout_rate: 0.1
    activation: silu
    num_groups: 4
    num_blocks: 3
  box_coder:
    _target_: xview3.centernet.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}

[I 2024-10-05 17:11:32,303] A new study created in memory with name: no-name-fdf90701-8eea-4018-98ba-cf19605b1172
Available keys in model config: dict_keys(['_target_', 'slug', 'weights', 'pretrained', 'output_stride', 'image_size', 'window_size', 'num_layers', 'num_filters', 'kernel_size', 'activation', 'dropout_rate', 'dropout_rate_max', 'optimizer', 'learning_rate', 'transformer', 'decoder', 'encoder', 'nas', 'num_channels', 'num_classes', 'anchors', 'input_upsample_factor', 'neck_out_channels', 'box_coder', 'head', 'fan', 'config'])
Complete model configuration:
 _target_: xview3.centernet.models.CenterNetUNetModel.from_config
slug: b5_unet_s2
weights: IMAGENET1K_V1
pretrained: true
output_stride: 32
image_size:
- 1024
- 1024
window_size:
- 4
- 4
num_layers: 4
num_filters: 128
kernel_size: 5
activation: leaky_relu
dropout_rate: 0.4996974369604577
dropout_rate_max: 0.5
optimizer: SGD
learning_rate: 5.430983971647737e-05
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: null
  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
  - 2
  - 3
  - 4
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_acc
  direction: maximize
num_channels: 2
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:
  enabled: true
  in_channels:
  - 32
  - 64
  - 128
  - 256
  out_channels: 256
  num_layers: 2
config:
  slug: b5_unet_s2
  activation: silu
  num_channels: ${dataset.num_channels}
  num_extra_blocks: 0
  encoder:
    _target_: pytorch_toolbelt.modules.encoders.TimmB5Encoder
    pretrained: true
    activation: silu
    layers:
    - 1
    - 2
    - 3
    - 4
  decoder:
    upsample_type: shuffle
    block_type: IRBlock
    channels:
    - 256
    - 384
    - 512
  head:
    _target_: xview3.centernet.models.heads.DecoupledHeadGroupNormLateShuffle
    classifier_dim: 128
    objectness_dim: 128
    size_dim: 64
    offset_dim: 16
    dropout_rate: 0.1
    activation: silu
    num_groups: 4
    num_blocks: 3
  box_coder:
    _target_: xview3.centernet.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}

[DEBUG] Number of heads instantiated: 2
[DEBUG] Head 0 type: <class 'xview3.centernet.models.heads.decoupled_head.DecoupledHeadGroupNormLateShuffle'>
[DEBUG] Head 1 type: <class 'xview3.centernet.models.heads.yolo_head.YOLOHead'>
[2024-10-05 17:11:32,592][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub (timm/efficientnetv2_rw_s.ra2_in1k)
[2024-10-05 17:11:32,820][timm.models._hub][INFO] - [timm/efficientnetv2_rw_s.ra2_in1k] Safe alternative available for 'pytorch_model.bin' (as 'model.safetensors'). Loading weights using safetensors.
[2024-10-05 17:11:32,840][timm.models._builder][WARNING] - Unexpected keys (bn2.bias, bn2.num_batches_tracked, bn2.running_mean, bn2.running_var, bn2.weight, classifier.bias, classifier.weight, conv_head.weight) found while loading pretrained weights. This may be expected if model is being adapted.
[2024-10-05 17:11:32,842][xview3.centernet.models.centernet_unet][INFO] - Adjusted encoder's first conv layer to have 2 input channels.
[2024-10-05 17:11:33,088][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub (timm/efficientnetv2_rw_s.ra2_in1k)
[2024-10-05 17:11:33,218][timm.models._hub][INFO] - [timm/efficientnetv2_rw_s.ra2_in1k] Safe alternative available for 'pytorch_model.bin' (as 'model.safetensors'). Loading weights using safetensors.
[2024-10-05 17:11:33,228][timm.models._builder][INFO] - Converted input conv conv_stem pretrained weights from 3 to 2 channel(s)
[2024-10-05 17:11:33,240][timm.models._builder][WARNING] - Unexpected keys (bn2.bias, bn2.num_batches_tracked, bn2.running_mean, bn2.running_var, bn2.weight, classifier.bias, classifier.weight, conv_head.weight) found while loading pretrained weights. This may be expected if model is being adapted.
[2024-10-05 17:11:33,241][xview3.centernet.models.centernet_unet][INFO] - Adjusted encoder's first conv layer to have 2 input channels.
[W 2024-10-05 17:11:33,293] Trial 0 failed with parameters: {'num_layers': 4, 'num_filters': 128, 'kernel_size': 5, 'activation': 'leaky_relu', 'dropout_rate': 0.4996974369604577, 'learning_rate': 5.430983971647737e-05, 'optimizer': 'SGD'} because of the following error: ConfigAttributeError("Key 'model' is not in struct\n    full_key: model.model\n    object_type=dict").
Traceback (most recent call last):
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/optuna/study/_optimize.py", line 197, in _run_trial
    value_or_values = func(trial)
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 240, in <lambda>
    lambda trial: self.evaluate_architecture(trial, config, train_loader, val_loader),
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 388, in evaluate_architecture
    model = self.nas.create_model(config, trial)  # Ensure create_model is accessible
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 110, in create_model
    model = CenterNetUNetModel.from_config(updated_config.model)
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 625, in from_config
    output_stride=config.model.get('output_stride', 32),
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 359, in __getattr__
    self._format_and_raise(key=key, value=None, cause=e)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/base.py", line 231, in _format_and_raise
    format_and_raise(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 819, in format_and_raise
    _raise(ex, cause)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 797, in _raise
    raise ex.with_traceback(sys.exc_info()[2])  # set env var OC_CAUSE=1 for full trace
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 351, in __getattr__
    return self._get_impl(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 442, in _get_impl
    node = self._get_child(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/basecontainer.py", line 73, in _get_child
    child = self._get_node(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 475, in _get_node
    self._validate_get(key)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 164, in _validate_get
    self._format_and_raise(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/base.py", line 231, in _format_and_raise
    format_and_raise(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 899, in format_and_raise
    _raise(ex, cause)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 797, in _raise
    raise ex.with_traceback(sys.exc_info()[2])  # set env var OC_CAUSE=1 for full trace
omegaconf.errors.ConfigAttributeError: Key 'model' is not in struct
    full_key: model.model
    object_type=dict
[W 2024-10-05 17:11:33,295] Trial 0 failed with value None.
Error executing job with overrides: []
Traceback (most recent call last):
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 821, 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 "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 817, in main
    pipeline.train()
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 535, in train
    self.model = self.build_model(config, train_loader, valid_loader)
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 340, in build_model
    best_trial = self.nas.optimize(
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 239, in optimize
    study.optimize(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/optuna/study/study.py", line 475, in optimize
    _optimize(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/optuna/study/_optimize.py", line 63, in _optimize
    _optimize_sequential(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/optuna/study/_optimize.py", line 160, in _optimize_sequential
    frozen_trial = _run_trial(study, func, catch)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/optuna/study/_optimize.py", line 248, in _run_trial
    raise func_err
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/optuna/study/_optimize.py", line 197, in _run_trial
    value_or_values = func(trial)
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 240, in <lambda>
    lambda trial: self.evaluate_architecture(trial, config, train_loader, val_loader),
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel3.py", line 388, in evaluate_architecture
    model = self.nas.create_model(config, trial)  # Ensure create_model is accessible
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 110, in create_model
    model = CenterNetUNetModel.from_config(updated_config.model)
  File "/Volumes/Crucial/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/centernet/models/centernet_unet.py", line 625, in from_config
    output_stride=config.model.get('output_stride', 32),
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 359, in __getattr__
    self._format_and_raise(key=key, value=None, cause=e)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/base.py", line 231, in _format_and_raise
    format_and_raise(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 819, in format_and_raise
    _raise(ex, cause)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 797, in _raise
    raise ex.with_traceback(sys.exc_info()[2])  # set env var OC_CAUSE=1 for full trace
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 351, in __getattr__
    return self._get_impl(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 442, in _get_impl
    node = self._get_child(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/basecontainer.py", line 73, in _get_child
    child = self._get_node(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 475, in _get_node
    self._validate_get(key)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/dictconfig.py", line 164, in _validate_get
    self._format_and_raise(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/base.py", line 231, in _format_and_raise
    format_and_raise(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 899, in format_and_raise
    _raise(ex, cause)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/omegaconf/_utils.py", line 797, in _raise
    raise ex.with_traceback(sys.exc_info()[2])  # set env var OC_CAUSE=1 for full trace
omegaconf.errors.ConfigAttributeError: Key 'model' is not in struct
    full_key: model.model
    object_type=dict
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 


##nas.py
# nas.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import optuna
from torch.optim import Adam, SGD, AdamW
from torch.utils.data import DataLoader
from typing import Any, Callable, Dict, Tuple, Optional
import random
import logging
from xview3.utils import TorchDevice
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder
import numpy as np
import pandas as pd
import torchmetrics
from xview3.centernet.models.dropath import DropPath
from xview3.centernet.models.custom_encoder import CustomTimmEncoder
from xview3.centernet.models.necks import FeatureAggregationNetwork
from xview3.centernet.models.heads import DecoupledHeadGroupNormLateShuffle, YOLOHead
from xview3.common import FeaturePyramidNetwork
from xview3.centernet.models.unet_blocks import SqueezeExcitation
from xview3.centernet.constants import *
from omegaconf import OmegaConf
device = TorchDevice.choose_torch_device()
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)

class FocalLoss(nn.Module):
    """
    Focal Loss for addressing class imbalance by down-weighting easy examples.
    """
    def __init__(self, alpha=2.0, beta=4.0, eps=1e-6):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.beta = beta
        self.eps = eps

    def forward(self, pred, gt):
        pos_inds = gt.eq(1).float()
        neg_inds = gt.lt(1).float()
        neg_weights = torch.pow(1 - gt, self.beta)

        # Ensure pred and gt have the same spatial dimensions
        if pred.size() != gt.size():
            raise ValueError(f"pred and gt must have the same size, but got pred: {pred.size()} and gt: {gt.size()}")

        pred = torch.clamp(pred, self.eps, 1 - self.eps)
        pos_loss = -torch.log(pred) * torch.pow(1 - pred, self.alpha) * pos_inds
        neg_loss = -torch.log(1 - pred) * torch.pow(pred, self.alpha) * neg_weights * neg_inds

        num_pos = pos_inds.float().sum()
        loss = (pos_loss + neg_loss).sum() / num_pos.clamp(min=1)
        return loss


class OptunaNAS:
    """
    Optuna-based Neural Architecture Search (NAS) for optimizing hyperparameters of the CenterNet-UNet model.
    """
    def __init__(self, search_space: Dict[str, Any], evaluate_architecture: Callable, max_epochs: int = 10, use_metrics: bool = True, direction='maximize'):
        """
        Initializes the OptunaNAS.

        Args:
            search_space (Dict[str, Any]): Dictionary defining the search space.
            evaluation_function (Callable): Function to evaluate the model, should return a float metric.
            max_epochs (int): Maximum number of epochs for training each trial.
            use_metrics (bool): Whether to compute additional metrics.
        """
        self.search_space = search_space
        self.evaluate_architecture = evaluate_architecture
        self.direction = direction
        self.max_epochs = max_epochs
        self.use_metrics = use_metrics

        # Initialize box coder
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(1024, 1024),
            output_stride=2,
            max_objects=512,
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )

    def create_model(self, config: Dict[str, Any], trial: optuna.Trial) -> nn.Module:
        """
        Creates and returns a CenterNet-UNet model based on the given hyperparameters.

        Args:
            config (Dict[str, Any]): Configuration dictionary.
            trial (optuna.Trial): Optuna trial object.

        Returns:
            nn.Module: The instantiated CenterNet-UNet model.
        """
        from xview3.centernet.models.centernet_unet import CenterNetUNetModel 
        # Suggest hyperparameters from the search space
        hyperparams = self.suggest_hyperparameters(trial)

        # Update the configuration with suggested hyperparameters
        updated_config = copy.deepcopy(config)
        for key, value in hyperparams.items():
            OmegaConf.update(updated_config, f"model.{key}", value)
        

        # Instantiate the model from the updated configuration
        model = CenterNetUNetModel.from_config(updated_config.model)
        model.to(device)

        return model

    
    def suggest_hyperparameters(self, trial: optuna.Trial) -> Dict[str, Any]:
        """
        Suggests hyperparameters using Optuna trial.

        Args:
            trial (optuna.Trial): Optuna trial object.

        Returns:
            Dict[str, Any]: Suggested hyperparameters.
        """

        hyperparams = {}
        # Suggest hyperparameters based on the search space
        hyperparams['num_layers'] = trial.suggest_int('num_layers', self.search_space['num_layers'][0], self.search_space['num_layers'][1])
        hyperparams['num_filters'] = trial.suggest_categorical('num_filters', self.search_space['num_filters'])
        hyperparams['kernel_size'] = trial.suggest_categorical('kernel_size', self.search_space['kernel_size'])
        hyperparams['activation'] = trial.suggest_categorical('activation', self.search_space['activation'])
        hyperparams['dropout_rate'] = trial.suggest_float('dropout_rate', 0.0, self.search_space['dropout_rate_max'])
        hyperparams['learning_rate'] = trial.suggest_float('learning_rate', 1e-5, 1e-2, log=True)
        hyperparams['optimizer'] = trial.suggest_categorical('optimizer', ['Adam', 'SGD', 'AdamW'])
        
        return hyperparams
    
    
    def evaluate_architecture(self, trial: optuna.Trial, config: Dict[str, Any], train_loader: DataLoader, val_loader: DataLoader) -> float:
        """
        Objective function for each Optuna trial.

        Args:
            trial (optuna.Trial): Optuna trial object.
            config (Dict[str, Any]): Configuration dictionary.
            train_loader (DataLoader): Training data loader.
            val_loader (DataLoader): Validation data loader.

        Returns:
            float: Validation metric to be optimized.
        """
        model = self.create_model(config, trial)
        model.to(device)

        # Define optimizer
        learning_rate = config.model.learning_rate  # Updated key
        optimizer_choice = config.model.optimizer   # Updated key
        
        if optimizer_choice.lower() == 'adamw':    # Ensuring case-insensitivity
            optimizer = AdamW(model.parameters(), lr=learning_rate)   # Corrected optimizer import and case
        elif optimizer_choice.lower() == 'adam':
            optimizer = Adam(model.parameters(), lr=learning_rate)
        elif optimizer_choice.lower() == 'sgd':
            optimizer = SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=1e-4)
        else:
            raise ValueError(f"Unsupported optimizer type: {optimizer_choice}")
        # Define loss function
        criterion = FocalLoss(alpha=2.0, beta=4.0, eps=1e-6)

        # Initialize mixed precision scaler if available
        scaler = torch.cuda.amp.GradScaler() if torch.cuda.is_available() else None

        # Define scheduler
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.max_epochs)

        # Initialize metrics
        accuracy_metric = torchmetrics.Accuracy().to(device)
        f1_score = torchmetrics.F1(num_classes=2, average='macro').to(device)

        # Training loop
        for epoch in range(self.max_epochs):
            model.train()
            epoch_loss = 0.0
            for batch in train_loader:
                inputs, targets = self.get_batch(batch)
                inputs = inputs.to(device)
                targets = self.prepare_targets(targets).to(device)

                optimizer.zero_grad()

                with torch.cuda.amp.autocast() if scaler else torch.enable_grad():
                    outputs = model(inputs)
                    loss = self.calculate_loss(outputs, targets)

                if scaler:
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    loss.backward()
                    optimizer.step()

                epoch_loss += loss.item()

            scheduler.step()

            # Validation
            val_acc, val_f1 = self.evaluate(model, val_loader, accuracy_metric, f1_score)
            logger.info(f"Trial {trial.number} Epoch {epoch+1}/{self.max_epochs} - Loss: {epoch_loss / len(train_loader):.4f} Val Acc: {val_acc:.4f} Val F1: {val_f1:.4f}")

            # Report intermediate objective value
            trial.report(val_acc, epoch)

            # Handle pruning based on the intermediate value
            if trial.should_prune():
                raise optuna.exceptions.TrialPruned()

        # Final validation after training
        final_val_acc, final_val_f1 = self.evaluate(model, val_loader, accuracy_metric, f1_score)
        logger.info(f"Trial {trial.number} - Final Validation Accuracy: {final_val_acc:.4f}, F1 Score: {final_val_f1:.4f}")

        return final_val_acc  # Adjust based on your optimization goal
        
    def optimize(self, n_trials: int, config: Dict[str, Any], train_loader: DataLoader, val_loader: DataLoader) -> optuna.Study:
        """
        Runs the Optuna optimization process.

        Args:
            n_trials (int): Number of Optuna trials.
            config (Dict[str, Any]): Configuration dictionary.
            train_loader (DataLoader): Training data loader.
            val_loader (DataLoader): Validation data loader.

        Returns:
            optuna.Trial: The best trial object.
        """
        study = optuna.create_study(direction=self.direction)
        study.optimize(
            lambda trial: self.evaluate_architecture(trial, config, train_loader, val_loader),
            n_trials=n_trials
        )
        return study.best_trial

    def calculate_loss(self, outputs: Dict[str, torch.Tensor], targets: Dict[str, torch.Tensor]) -> torch.Tensor:
        """
        Calculates the composite loss for multi-task learning.

        Args:
            outputs (Dict[str, torch.Tensor]): Model outputs.
            targets (Dict[str, torch.Tensor]): Ground truth targets.

        Returns:
            torch.Tensor: Total loss.
        """
        total_loss = 0.0

        # Example loss components, adjust based on your outputs and targets
        # Heatmap loss (Focal Loss)
        if CENTERNET_OUTPUT_OBJECTNESS_MAP in outputs and 'CENTERNET_TARGET_OBJECTNESS_MAP' in targets:
            heatmap_loss = FocalLoss(alpha=2.0, beta=4.0, eps=1e-6)(outputs[CENTERNET_OUTPUT_OBJECTNESS_MAP], targets['CENTERNET_TARGET_OBJECTNESS_MAP'])
            total_loss += heatmap_loss

        # Vessel loss (Binary Cross-Entropy)
        if CENTERNET_OUTPUT_VESSEL_MAP in outputs and 'CENTERNET_TARGET_VESSEL_MAP' in targets:
            vessel_loss = F.binary_cross_entropy_with_logits(outputs[CENTERNET_OUTPUT_VESSEL_MAP], targets['CENTERNET_TARGET_VESSEL_MAP'])
            total_loss += vessel_loss

        # Fishing loss (Binary Cross-Entropy)
        if CENTERNET_OUTPUT_FISHING_MAP in outputs and 'CENTERNET_TARGET_FISHING_MAP' in targets:
            fishing_loss = F.binary_cross_entropy_with_logits(outputs[CENTERNET_OUTPUT_FISHING_MAP], targets['CENTERNET_TARGET_FISHING_MAP'])
            total_loss += fishing_loss

        # Size loss (L1 Loss)
        if CENTERNET_OUTPUT_SIZE in outputs and 'CENTERNET_TARGET_SIZE' in targets:
            size_loss = F.l1_loss(outputs[CENTERNET_OUTPUT_SIZE], targets['CENTERNET_TARGET_SIZE'])
            total_loss += size_loss

        # Offset loss (L1 Loss)
        if CENTERNET_OUTPUT_OFFSET in outputs and 'CENTERNET_TARGET_OFFSET' in targets:
            offset_loss = F.l1_loss(outputs[CENTERNET_OUTPUT_OFFSET], targets['CENTERNET_TARGET_OFFSET'])
            total_loss += offset_loss

        # YOLO loss (Assuming you have YOLO outputs and targets)
        if CENTERNET_OUTPUT_YOLO in outputs and CENTERNET_TARGET_YOLO in targets:
            yolo_loss = self.box_coder.yolo_loss(outputs[CENTERNET_OUTPUT_YOLO], targets[CENTERNET_TARGET_YOLO])
            total_loss += yolo_loss

        return total_loss

    def evaluate(self, model: nn.Module, val_loader: DataLoader, accuracy_metric: torchmetrics.Metric, f1_metric: torchmetrics.Metric) -> Tuple[float, float]:
        """
        Evaluates the model on the validation set and computes the desired metrics.

        Args:
            model (nn.Module): The model to evaluate.
            val_loader (DataLoader): Validation data loader.
            accuracy_metric (torchmetrics.Metric): Accuracy metric instance.
            f1_metric (torchmetrics.Metric): F1 score metric instance.

        Returns:
            Tuple[float, float]: Validation accuracy and F1 score.
        """
        model.eval()
        accuracy_metric.reset()
        f1_metric.reset()
        with torch.no_grad():
            for batch in val_loader:
                inputs, targets = self.get_batch(batch)
                inputs = inputs.to(device)
                targets = self.prepare_targets(targets).to(device)

                outputs = model(inputs)

                # Compute metrics for vessel classification
                if CENTERNET_OUTPUT_VESSEL_MAP in outputs and 'CENTERNET_TARGET_VESSEL_MAP' in targets:
                    preds = torch.sigmoid(outputs[CENTERNET_OUTPUT_VESSEL_MAP]) > 0.5
                    targets_binary = targets['CENTERNET_TARGET_VESSEL_MAP'] > 0.5
                    accuracy_metric.update(preds, targets_binary)
                    f1_metric.update(preds, targets_binary)

        val_acc = accuracy_metric.compute().item()
        val_f1 = f1_metric.compute().item()
        return val_acc, val_f1

    def get_batch(self, batch: Tuple[Any, Any]) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
        """
        Processes a batch from the DataLoader.

        Args:
            batch (Tuple[Any, Any]): A batch from the DataLoader.

        Returns:
            Tuple[torch.Tensor, Dict[str, torch.Tensor]]: Processed inputs and prepared targets.
        """
        inputs, targets = batch
        return inputs, targets
    

    def prepare_targets(self, targets: Dict[str, Any]) -> Dict[str, torch.Tensor]:
        """
        Prepares targets for the loss function.

        Args:
            targets (Dict[str, Any]): Raw targets from the DataLoader.

        Returns:
            Dict[str, torch.Tensor]: Prepared targets as tensors.
        """
        prepared = {}
        for key, value in targets.items():
            prepared[key] = value.to(dtype=torch.float32)
        return prepared

    def run(self, n_trials: int, config: Dict[str, Any], train_loader: DataLoader, val_loader: DataLoader) -> optuna.Study:
        """
        Runs the Optuna optimization process.

        Args:
            n_trials (int): Number of Optuna trials.
            config (Dict[str, Any]): Configuration dictionary.
            train_loader (DataLoader): Training data loader.
            val_loader (DataLoader): Validation data loader.

        Returns:
            optuna.Study: The Optuna study object containing optimization results.
        """
        # Define the Optuna study with advanced sampler and pruner
        sampler = optuna.samplers.TPESampler(n_startup_trials=10, n_ei_candidates=24)
        pruner = optuna.pruners.MedianPruner(n_startup_trials=5, n_warmup_steps=5)
        study = optuna.create_study(direction='maximize', sampler=sampler, pruner=pruner)
        study.optimize(lambda trial: self.evaluate_architecture(trial, config, train_loader, val_loader), n_trials=n_trials)
        return study