(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % HYDRA_FULL_ERROR=1 python train_multilabel.py
DEBUG:hydra.core.utils:Setting JobRuntime:name=UNKNOWN_NAME
DEBUG:hydra.core.utils:Setting JobRuntime:name=train_multilabel
Dataset full_4fold
        ignore_low_confidence_detections True
        ignore_low_confidence_labels     True
        filter_low_confidence_objects    False
Train dataset
        Unknown vessels 7788
        Unknown fishing 31133
Valid dataset
        Unknown vessels 36
        Unknown fishing 3934
Preparing training dataset: 100%|█████████████████████████████████████| 425/425 [00:00<00:00, 491.31it/s]
Preparing validation dataset: 100%|█████████████████████████████████████| 12/12 [00:00<00:00, 567.80it/s]
Available keys in model config: dict_keys(['_target_', 'slug', 'weights', 'pretrained', 'output_stride', 
Available keys in model config: dict_keys(['_target_', 'slug', 'weights', 'pretrained', 'output_stride', 'image_size', 'decoder', 'encoder', 'nas', 'num_channels', 'num_classes', 'anchors', 'input_upsample_fact
Available keys in model config: dict_keys(['_target_', 'slug', 'weights', 'pretrained', 'output_stride', 'image_size', 'decoder', 'encoder', 'nas', 'num_channels', 'num_classes', 'anchors', 'input_upsample_factor', 'neck_out_channels', 'box_coder', 'head', '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
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
    dropout_rate_max: 0.5
  num_trials: 50
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.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
    input_image_size:
    - 640
    - 640
    stride: 32
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.DecoupledHeadGroupNormLateShuff
[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-04 17:17:27,071][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub 
[2024-10-04 17:17:27,071][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub (timm/efficientnetv2_rw_s.ra2_in1k)
[2024-10-04 17:17:27,272][timm.models._hub][INFO] - [timm/efficientnetv2_rw_s.ra2_in1k] Safe alternative 
[2024-10-04 17:17:27,272][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-04 17:17:27,398][timm.models._builder][WARNING] - Unexpected keys (bn2.bias, bn2.num_batches_tra
[2024-10-04 17:17:27,398][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
[2024-10-04 17:17:27,398][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-04 17:17:27,648][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub 
[2024-10-04 17:17:27,648][timm.models._builder][INFO] - Loading pretrained weights from Hugging Face hub (timm/efficientnetv2_rw_s.ra2_in1k)
[2024-10-04 17:17:27,779][timm.models._hub][INFO] - [timm/efficientnetv2_rw_s.ra2_in1k] Safe alternative 
[2024-10-04 17:17:27,779][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-04 17:17:27,791][timm.models._builder][INFO] - Converted input conv conv_stem pretrained weights
[2024-10-04 17:17:27,791][timm.models._builder][INFO] - Converted input conv conv_stem pretrained weights from 3 to 2 channel(s)
[2024-10-04 17:17:27,805][timm.models._builder][WARNING] - Unexpected keys (bn2.bias, bn2.num_batches_tra
[2024-10-04 17:17:27,805][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
[2024-10-04 17:17:27,805][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-04 17:17:27,807][xview3.centernet.models.centernet_unet][INFO] - Adjusted encoder's first conv l
[2024-10-04 17:17:27,807][xview3.centernet.models.centernet_unet][INFO] - Adjusted encoder's first conv layer 'initial_conv' to have 2 input channels.
[2024-10-04 17:17:32,705][__main__][INFO] - Total parameters in the model: {'total': 263815591, 'trainabl
[2024-10-04 17:17:32,705][__main__][INFO] - Total parameters in the model: {'total': 263815591, 'trainable': 263815591, 'encoder': 21657072, 'decoder': 14193536}
Losses
   losses/ MultilabelCenterTrimmedLoss weight 1.0
     target ['CENTERNET_TARGET_OBJECTNESS_MAP', 'CENTERNET_TARGET_VESSEL_MAP', 'CENTERNET_TARGET_FISHING_
     target ['CENTERNET_TARGET_OBJECTNESS_MAP', 'CENTERNET_TARGET_VESSEL_MAP', 'CENTERNET_TARGET_FISHING_MAP', 'CENTERNET_TARGET_SIZE', 'CENTERNET_TARGET_OFFSET']
     output ['CENTERNET_OUTPUT_OBJECTNESS_MAP', 'CENTERNET_OUTPUT_VESSEL_MAP', 'CENTERNET_OUTPUT_FISHING_
     output ['CENTERNET_OUTPUT_OBJECTNESS_MAP', 'CENTERNET_OUTPUT_VESSEL_MAP', 'CENTERNET_OUTPUT_FISHING_MAP', 'CENTERNET_OUTPUT_SIZE', 'CENTERNET_OUTPUT_OFFSET']
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 7
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 723, 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
  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
  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_
  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 <lam
  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 7
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 719, in main
    pipeline.train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 4
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 458, in train
    self.optimizer, self.scheduler, optimizer_callbacks, self.scaler = self.build_optimizer(self.model, t
    self.optimizer, self.scheduler, optimizer_callbacks, self.scaler = self.build_optimizer(self.model, train_loader)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 4
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 403, in build_optimizer
    scaler = torch.amp.GradScaler(enabled=optimizer_config.get("fp16", False)) if device != TorchDevice.C
    scaler = torch.amp.GradScaler(enabled=optimizer_config.get("fp16", False)) if device != TorchDevice.CPU_DEVICE else None
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/amp/grad_scaler.py", line 132, in __in
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/amp/grad_scaler.py", line 132, in __init__
    warnings.warn(
  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 7
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 71, in warn_with_traceback
    traceback.print_stack(file=log)
/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/torch/amp/grad_scaler.py:132: UserWarning:

torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.

Error executing job with overrides: []
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 7
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 723, 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
  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
  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_
  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_
  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 <lam
  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_va
  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 7
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 719, in main
    pipeline.train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 5
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 563, in train
    runner.train(
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 9
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 93, in train
    super().train(*args, **kwargs)
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/runner.py", line 144, in tr
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/catalyst/runners/runner.py", line 144, in train
    raise NotImplementedError("CheckpointCallback already exist")
NotImplementedError: CheckpointCallback already exist


just ignore the cuda issue and focus on checkpoint ones
###train_multilabel.py

from torch.utils.tensorboard import SummaryWriter
import sys
import traceback
import warnings
from typing import Dict, Any, Tuple, Optional, List
import os
import numpy as np
import pandas as pd
import torch
import hydra
from catalyst import dl
from catalyst.callbacks import MetricCallback, CheckpointCallback
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from catalyst.contrib.nn import Lookahead
from catalyst.dl import OptimizerCallback, SchedulerCallback
from torch.optim import Optimizer
from torch.utils.data import DataLoader
from torch import nn
from omegaconf import DictConfig, OmegaConf
from xview3.nas import OptunaNAS
import optuna
import albumentations as A
from catalyst.callbacks import ControlFlowCallback, CriterionCallback, SchedulerCallback, TimerCallback
from catalyst.core import Callback
from pytorch_toolbelt.utils.catalyst import (
    ShowPolarBatchesCallback,
    BestMetricCheckpointCallback,
    HyperParametersCallback,
    MixupCriterionCallback,
    MixupInputCallback,
)
from pytorch_toolbelt.optimization.functional import (
    freeze_model,
    get_lr_decay_parameters,
)
from pytorch_toolbelt.utils import fs, transfer_weights, count_parameters, to_numpy
from pytorch_toolbelt.utils.distributed import is_main_process
from xview3.centernet.bboxer import MultilabelCircleNetCoder
from xview3.early_stopping import EarlyStoppingCallback
from xview3.centernet.multilabel_dataset import MultilabelCircleNetDataModule
from xview3.centernet.models.heads.yolo_head import YOLOHead, generate_yolo_targets
from xview3.centernet.models.heads.yolo_decoder import YOLODecoder
from xview3.centernet.models.transformer_classifier import SwinTransformer, SwinTransformerEncoder
from xview3.centernet.losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet import *
from xview3.pipeline import Pipeline
from xview3.factory import *
from xview3.augmentations import *
from xview3.samplers import *
from xview3.dataset import read_multichannel_image, stack_multichannel_image
from xview3.ema import *
from xview3.evaluation import evaluate_on_scenes
from xview3.inference import predict_multilabel_scenes
from xview3.metric import *
from xview3.constants import INPUT_SCENE_ID_KEY
from xview3.centernet.constants import CENTERNET_OUTPUT_YOLO
import logging
from torch.cuda.amp import autocast, GradScaler
from xview3.utils import TorchDevice
from xview3.evaluation import evaluate_on_scenes
from datetime import datetime
import torch.nn.functional as F

logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)


def warn_with_traceback(message, category, filename, lineno, file=None, line=None):
    log = file if hasattr(file, "write") else sys.stderr
    traceback.print_stack(file=log)
    log.write(warnings.formatwarning(message, category, filename, lineno))


def safe_divide(numerator, denominator, default=0):
    return numerator / denominator if denominator != 0 else default


warnings.showwarning = warn_with_traceback

def get_optimizable_parameters(model):
    """
    Returns all trainable parameters of the model.
    """
    return [p for p in model.parameters() if p.requires_grad]

def callback_exists(callbacks, callback_type):
            return any(isinstance(cb, callback_type) or (hasattr(cb, 'callback') and isinstance(cb.callback, callback_type)) for cb in callbacks)


class CallbackOrder:
    OPTIMIZER = 550
    EXTERNAL = 1000
    METRICS = 1050
    
class MixedPrecisionCallback(dl.Callback):
    def __init__(self, scaler, device):
        super().__init__(order=99) 
        self.scaler = scaler
        self.device = device

    def on_batch_start(self, runner):
        runner.model.train()
        runner.optimizer.zero_grad()

    def on_batch_end(self, runner):
        if self.scaler and self.device != TorchDevice.CPU_DEVICE:
            self.scaler.step(runner.optimizer)
            self.scaler.update()
        else:
            runner.optimizer.step()

class CenternetPipeline(Pipeline):
    def __init__(self, cfg: DictConfig):
        super().__init__(cfg)
        self.box_coder: Optional[MultilabelCircleNetCoder] = None
        self.model: Optional[nn.Module] = None
        self.criterion: Optional[nn.Module] = None
        self.scaler: Optional[GradScaler] = GradScaler() if torch.cuda.is_available() else None
        self.device = TorchDevice.choose_torch_device()
        self.experiment_dir = ""
        self.yolo_decoder: Optional[YOLODecoder] = None
        self.nas: Optional[OptunaNAS] = None

    def build_datasets(self, config: DictConfig) -> Tuple[DataLoader, DataLoader, Optional[DataLoader], List[Callback]]:
        """
        Builds and returns DataLoaders for training and validation datasets.
        """
        dataset_module = MultilabelCircleNetDataModule(config=config)

        train_image_size = config.train.image_size if 'image_size' in config.train else [512, 512]
        valid_image_size = config.valid.image_size if 'image_size' in config.valid else [2048, 2048]

        train_df, valid_df, _, shore_root = dataset_module.train_val_split(
            splitter=config.dataset.splitter,
            fold=config.dataset.fold,
            num_folds=config.dataset.num_folds,
        )

        # Validate output_stride
        output_stride = config.model.get('output_stride', 32)
        if output_stride <= 0:
            raise ValueError(f"Output stride must be a positive integer, got {output_stride}")

        self.box_coder = MultilabelCircleNetCoder(
            image_size=(valid_image_size[0], valid_image_size[1]),
            output_stride=output_stride,
            max_objects=config.model.get('max_objects', 512),
            heatmap_encoding=config.model.get('heatmap_encoding', "umich"),
            labels_encoding=config.model.get('labels_encoding', "circle"),
            ignore_value=config.model.get('ignore_value', -1)
        )

        normalization = getattr(config.dataset, 'normalization', None)
        if normalization is None:
            # Define default normalization
            normalization = {
                'mean': [0.485, 0.456, 0.406],
                'std': [0.229, 0.224, 0.225]
            }
            logger.warning("Normalization not defined in config. Using default normalization.")

        # Check and validate heatmap_size
        heatmap_size = valid_image_size[0] // output_stride
        if heatmap_size <= 0:
            raise ValueError(f"Heatmap size must be a positive integer, got {heatmap_size}")

        self.valid_df = valid_df.copy()  # Keep a copy of unmodified groundtruth

        # Logging dataset info
        self.master_print("Dataset", config.dataset.slug)
        self.master_print("\tignore_low_confidence_detections", config.dataset.ignore_low_confidence_detections)
        self.master_print("\tignore_low_confidence_labels    ", config.dataset.ignore_low_confidence_labels)
        self.master_print("\tfilter_low_confidence_objects   ", config.dataset.filter_low_confidence_objects)

        self.master_print("Train dataset")
        self.master_print("\tUnknown vessels", np.isnan(train_df.is_vessel.values.astype(float)).sum())
        self.master_print("\tUnknown fishing", np.isnan(train_df.is_fishing.values.astype(float)).sum())
        self.master_print("Valid dataset")
        self.master_print("\tUnknown vessels", np.isnan(valid_df.is_vessel.values.astype(float)).sum())
        self.master_print("\tUnknown fishing", np.isnan(valid_df.is_fishing.values.astype(float)).sum())

        # Check mutually exclusive configurations
        if config.dataset.filter_low_confidence_objects and config.dataset.ignore_low_confidence_labels:
            raise ValueError("filter_low_confidence_objects and ignore_low_confidence_labels are mutually exclusive")

        # Data augmentation and normalization
        transforms = A.Compose([
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.Rotate(limit=15, p=0.5),
            A.RandomBrightnessContrast(p=0.5),
            A.Normalize(mean=normalization['mean'], std=normalization['std']),
        ], bbox_params=A.BboxParams(format='coco', label_fields=['labels']))

        # Initialize datasets
        train_dataset, train_sampler = dataset_module.get_random_crop_training_dataset(
            train_df=train_df,
            train_image_size=train_image_size,
            input_channels=config.dataset.channels,
            box_coder=self.box_coder,
            normalization=transforms,
            num_samples=config.train.get('num_samples', None),
            individual_augmentations=config.train.get('individual_augmentations', {}),
            augmentations=config.train.get('augs', []),
            balance_crowd=config.train.get('balance_crowd', False),
            balance_location=config.train.get('balance_location', False),
            channels_last=config.dataset.get('channels_last', False),
            read_image_fn=read_multichannel_image,
        )

        valid_dataset = dataset_module.get_validation_dataset(
            valid_df=valid_df,
            valid_crop_size=valid_image_size,
            input_channels=config.dataset.channels,
            box_coder=self.box_coder,
            normalization=transforms,
            channels_last=config.dataset.get('channels_last', False),
            read_image_fn=read_multichannel_image,
        )

        # Setting sampler to None for distributed training or default sampling
        train_sampler = None

        # Prepare callbacks (e.g., for logging, early stopping)
        dataset_callbacks = []

        # Initialize DataLoaders
        train_loader = DataLoader(
            train_dataset,
            sampler=train_sampler,
            batch_size=config.train.batch_size,
            num_workers=config.train.num_workers,
            pin_memory=True,
            shuffle=(train_sampler is None)
        )
        valid_loader = DataLoader(
            valid_dataset,
            batch_size=config.train.get('batch_size', config.train.batch_size),
            num_workers=config.train.get('num_workers', config.train.num_workers),
            pin_memory=True,
            shuffle=False
        )

        return train_loader, valid_loader, None, dataset_callbacks

    def build_model(self, config: DictConfig, loaders: Dict[str, DataLoader]) -> nn.Module:
        """
        Builds the model using CenterNetUNet with integrated heads and NAS if configured.
        """
        # Instantiate CenterNetUNetModel from config
        self.model = CenterNetUNetModel.from_config(config.model)
        self.model = self.model.to(self.device)

        # Initialize YOLODecoder
        self.yolo_decoder = YOLODecoder(
            anchors=config.model.anchors,
            num_classes=config.model.num_classes,
            input_image_size=tuple(config.model.get('image_size', (1024, 1024))),
            stride=config.model.output_stride,
            conf_threshold=config.inference.get('thresholds', {}).get('objectness', 0.5),
            iou_threshold=config.inference.get('thresholds', {}).get('iou', 0.4)
        ).to(self.device)

        # Make sure all model parameters are correctly loaded
        self.model = self.model.to(self.device)
        self.yolo_decoder = self.yolo_decoder.to(self.device)

        logger.info(f"Total parameters in the model: {count_parameters(self.model)}")

        return self.model

    def build_criterions(self, config: DictConfig) -> Tuple[Dict[str, nn.Module], List[Callback]]:
        """
        Builds and returns loss functions and their corresponding callbacks.
        """
        criterions, criterions_callbacks = super().build_criterions(config)

        # Initialize combined loss (assuming model has a calculate_loss method)
        self.criterion = CircleNetLoss().to(self.device)
        self.yolo_loss = YOLOLoss(
            num_classes=config.model.num_classes,
            anchors=config.model.anchors
        ).to(self.device)

        criterions['combined_loss'] = self.criterion
        # Optionally, add YOLO loss separately if needed
        # criterions['yolo_loss'] = self.yolo_loss

        return criterions, criterions_callbacks

    def evaluate_architecture(self, model: nn.Module, train_loader: DataLoader, val_loader: DataLoader) -> float:
        """
        Evaluates a given model architecture on the validation set.
        This function is used by OptunaNAS to determine the performance of different architectures.
        """
        model.to(self.device)
        model.eval()
        total_loss = 0.0
        total_samples = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(val_loader):
                try:
                    inputs = batch['images'].to(self.device)

                    # Ensure correct input shape
                    if inputs.dim() == 3:
                        inputs = inputs.unsqueeze(0)

                    # Create 3-channel input for models expecting 3 channels
                    if inputs.shape[1] == 2:
                        third_channel = (inputs[:, 0, :, :] + inputs[:, 1, :, :]) / 2
                        inputs = torch.cat([inputs, third_channel.unsqueeze(1)], dim=1)

                    logger.debug(f"[DEBUG] Input shape: {inputs.shape}")

                    encoded_targets = self.encode_batch_targets(batch)

                    # Forward pass
                    outputs = model(inputs)

                    # Decode YOLO outputs
                    yolo_out = outputs.get(CENTERNET_OUTPUT_YOLO, None)
                    if yolo_out is not None:
                        decoded_yolo = self.yolo_decoder.decode(yolo_out)
                        # Depending on your evaluation metric, you might want to process decoded_yolo

                    # Calculate loss
                    loss = self.criterion(outputs, encoded_targets)
                    total_loss += loss.item() * inputs.size(0)
                    total_samples += inputs.size(0)

                except Exception as e:
                    logger.error(f"Error in batch {batch_idx}: {str(e)}")
                    logger.debug(f"Input shape: {inputs.shape}")
                    raise  # Re-raise the exception for full traceback

        # Compute average loss
        average_loss = safe_divide(total_loss, total_samples, default=float('inf'))
        logger.debug(f"[DEBUG] Architecture Evaluation - Average Validation Loss: {average_loss}")

        return -average_loss  # Optuna minimizes the objective, so return negative loss

    def encode_batch_targets(self, batch: Dict[str, Any]) -> Dict[str, torch.Tensor]:
        """
        Encodes batch targets for training.
        """
        inputs, targets = batch['images'], batch['targets']
        encoded_targets = self.box_coder.encode_batch(targets)
        return {k: v.to(self.device) for k, v in encoded_targets.items()}
    """
    def build_optimizer(self, model: nn.Module, loaders: Dict[str, DataLoader]) -> Tuple[Optimizer, Any, List[Callback]]:
        optimizer = torch.optim.AdamW(get_optimizable_parameters(model), lr=self.cfg.train.learning_rate, weight_decay=self.cfg.train.weight_decay)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.cfg.train.epochs)
        optimizer_callbacks = [SchedulerCallback()]
        return optimizer, scheduler, optimizer_callbacks
    """
    
    
    def build_optimizer(self, model: torch.nn.Module, train_loader: torch.utils.data.DataLoader) -> Tuple[torch.optim.Optimizer, Any, List[dl.Callback], Optional[GradScaler]]:
        optimizer_config = self.cfg.optimizer
        lr = optimizer_config.params.lr
        weight_decay = optimizer_config.params.weight_decay
        eps = optimizer_config.params.eps

        parameters = get_optimizable_parameters(model)

        optimizer = torch.optim.AdamW(
            parameters,
            lr=lr,
            weight_decay=weight_decay,
            eps=eps
        )

        if optimizer_config.get("lookahead", False):
            optimizer = Lookahead(optimizer)

        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader) * self.cfg.train.epochs)

        optimizer_callbacks = [
            dl.OptimizerCallback(metric_key="loss"),
            dl.SchedulerCallback(mode="batch")
        ]

        device = TorchDevice.choose_torch_device()
        scaler = torch.amp.GradScaler(enabled=optimizer_config.get("fp16", False)) if device != TorchDevice.CPU_DEVICE else None

        return optimizer, scheduler, optimizer_callbacks, scaler


    def get_experiment_dir(self, config: DictConfig) -> str:
        """
        Determines and returns the directory path for logging and saving model checkpoints.
        """
        if config.get("experiment", None):
            return config.experiment

        current_time = datetime.now().strftime("%y%m%d_%H_%M")
        experiment_slug = f"{current_time}_{config.model.slug}_{config.dataset.slug}_{config.train.slug}_"

        # Append unique identifiers based on folds or other config parameters
        fold = config.dataset.get("fold", None)
        if fold is not None:
            experiment_slug += f"fold{fold}_"

        # Include augmentation slug
        experiment_slug += f"{config.augs.slug}"

        log_dir = os.path.join("runs", experiment_slug)
        return log_dir

    def get_checkpoint_data(self) -> Dict[str, Any]:
        """
        Retrieves data to be saved in checkpoints.
        """
        checkpoint_data = {
            "model_state_dict": self.model.state_dict(),
            "optimizer_state_dict": self.optimizer.state_dict(),
            "scheduler_state_dict": self.scheduler.state_dict(),
            "scaler_state_dict": self.scaler.state_dict() if self.scaler else None,
            # Add any other necessary state info here
        }
        return checkpoint_data

    def on_experiment_finished(self):
        """
        Handles post-training steps such as saving the final model.
        """
        final_model_path = os.path.join(self.experiment_dir, "final_model.pth")
        torch.save(self.model.state_dict(), final_model_path)
        self.master_print(f"Final model saved to {final_model_path}")


    def train(self):
        config = self.cfg
        train_loader, valid_loader, _, dataset_callbacks = self.build_datasets(config)
        self.model = self.build_model(config, train_loader)
        criterions, criterions_callbacks = self.build_criterions(config)
        metric_callbacks = self.build_metrics(config, train_loader, self.model)

        self.optimizer, self.scheduler, optimizer_callbacks, self.scaler = self.build_optimizer(self.model, train_loader)

        self.experiment_dir = self.get_experiment_dir(config)

        if self.is_master:
            os.makedirs(self.experiment_dir, exist_ok=True)
            dst_config_fname = os.path.join(self.experiment_dir, "config.yaml")
            with open(dst_config_fname, "w") as f:
                OmegaConf.save(self.cfg, f)

        device = TorchDevice.choose_torch_device()
        mixed_precision_callback = MixedPrecisionCallback(self.scaler, device)

        runner = dl.SupervisedRunner(
            input_key=config.runner.input_key,
            output_key=config.runner.output_key,
            device=device
        )

        callbacks = [mixed_precision_callback] + dataset_callbacks + metric_callbacks + optimizer_callbacks + criterions_callbacks
        callbacks = [cb for cb in callbacks if not isinstance(cb, CheckpointCallback) and not (hasattr(cb, 'callback') and isinstance(cb.callback, CheckpointCallback))]
        

        # Add a callback for monitoring the main metric
        main_metric_callback = MetricCallback(
            prefix="main_metric",
            metric_fn=lambda outputs, targets: outputs[config.runner.main_metric],
            input_key="outputs",
            target_key="targets"
        )
        callbacks.append(main_metric_callback)

        # Add checkpoint callback
        if not callback_exists(callbacks, CheckpointCallback):
            checkpoint_callback = CheckpointCallback(
                save_n_best=3,
                resume=True,
                resume_dir=self.experiment_dir
            )
            callbacks.append(checkpoint_callback)


        # Add early stopping callback
        early_stopping = EarlyStoppingCallback(
            patience=config.train.get("early_stopping_patience", 10),
            min_delta=config.train.get("early_stopping_min_delta", 1e-5),
            #metric_key=f"main_metric/{config.runner.main_metric}",
            minimize=config.runner.main_metric_minimize
        )
        callbacks.append(early_stopping)

        # Add gradient clipping callback
        class GradientClippingCallback(dl.Callback):
            def __init__(self, max_norm=config.train.get("max_grad_norm", 10.0)):
                super().__init__(order=CallbackOrder.OPTIMIZER)
                self.max_norm = max_norm

            def on_batch_end(self, runner: "IRunner"):
                if (runner.batch_step + 1) % config.train.get("accumulation_steps", 1) == 0:
                    torch.nn.utils.clip_grad_norm_(runner.model.parameters(), self.max_norm)

        callbacks.append(GradientClippingCallback())

        # Add model architecture logging callback
        class ModelArchCallback(dl.Callback):
            def on_epoch_start(self, runner: "IRunner"):
                if runner.epoch == 0:
                    runner.logger.info(f"Model Architecture:\n{runner.model}")

        callbacks.append(ModelArchCallback(order=CallbackOrder.EXTERNAL))

        # Add custom metric logging callback
        class CustomMetricCallback(dl.Callback):
            def on_epoch_end(self, runner: "IRunner"):
                epoch_loss = np.mean(runner.batch_metrics["loss"])
                runner.loader_metrics["custom_metric"] = epoch_loss * runner.global_epoch

        callbacks.append(CustomMetricCallback(order=CallbackOrder.METRICS))

        # Add Tensorboard logging
        log_dir = os.path.join(os.getcwd(), "tensorboard_logs")
        os.makedirs(log_dir, exist_ok=True)
        tb_writer = SummaryWriter(log_dir=log_dir)

        # Create a custom TensorBoard callback
        class TensorboardLoggerCallback(dl.Callback):
            def __init__(self, writer):
                super().__init__(order=CallbackOrder.EXTERNAL)
                self.writer = writer

            def on_epoch_end(self, runner):
                for k, v in runner.loader_metrics.items():
                    self.writer.add_scalar(f'metrics/{k}', v, runner.epoch)

        callbacks.append(TensorboardLoggerCallback(tb_writer))

        runner.train(
            model=self.model,
            criterion=self.criterion,
            optimizer=self.optimizer,
            scheduler=self.scheduler,
            loaders={
                "train": train_loader,
                "valid": valid_loader
            },
            num_epochs=config.train.epochs,
            callbacks=callbacks,
            logdir=self.experiment_dir,
            verbose=True,
            load_best_on_end=True,
            checkpoint_data=self.get_checkpoint_data(),
            fp16=config.train.get("fp16", False),
            #amp=config.train.get("amp", False)
        )

        self.on_experiment_finished()

    def train_step(self, runner: 'catalyst.core.runner.Runner') -> None:
        inputs, targets = self.encode_batch_targets(runner.input)
        with autocast(enabled=self.cfg.train.amp.enabled and runner.device != TorchDevice.CPU_DEVICE):
            outputs = self.model(inputs)
            loss = self.model.calculate_loss(outputs, targets)

        runner.batch_metrics["loss"] = loss.item()

        if self.scaler and runner.device != TorchDevice.CPU_DEVICE:
            self.scaler.scale(loss).backward()
            if (runner.step + 1) % self.cfg.train.accumulation_steps == 0:
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.train.max_grad_norm)
        else:
            loss.backward()
            if (runner.step + 1) % self.cfg.train.accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.train.max_grad_norm)


    def build_metrics(self, config: DictConfig, loaders: Dict, model: nn.Module) -> List[Callback]:
        callbacks = [
            CircleNetRocAucMetricsCallback(prefix="extra_metrics"),
            MultilabelCircleNetDecodeCallback(box_coder=self.box_coder),
        ]

        show_batches = config.train.get("show", False)

        callbacks += [
            ControlFlowCallback(
                MultilabelCircleNetDetectionMetrics(
                    groundtruth_df=self.valid_df,
                    #shore_root=self.shore_root,
                    prefix="",
                    objectness_thresholds=[0.25, 0.275, 0.3, 0.325, 0.350, 0.375, 0.4, 0.425, 0.45, 0.475, 0.5, 0.525, 0.55],
                ),
                loaders="valid",
            ),
        ]
        
        if self.is_master:
            callbacks.append(
                BestMetricCheckpointCallback(
                    target_metric="metrics/aggregate",
                    target_metric_minimize=False,
                    save_n_best=3,
                    checkpoints_dir=self.experiment_dir
                )
            )
        
        if config.train.early_stopping > 0:
            callbacks.append(
                EarlyStoppingCallback(
                    metrics=["metrics/aggregate"],
                    minimize=False,
                    min_delta=1e-6,
                    patience=config.train.early_stopping,
                )
            )

        if self.is_master and show_batches:
            # Add ShowPolarBatchesCallback if needed
            # You'll need to implement or import the visualize_predictions_multilabel_centernet function
            pass

        if self.is_master:
            # Get the output_stride from the box_coder
            output_stride = self.box_coder.output_stride if hasattr(self.box_coder, 'output_stride') else None

            callbacks.append(
                HyperParametersCallback(
                    hparam_dict=dict(
                        model=str(config.model.config.slug),
                        model_output_stride=output_stride,
                        box_coder_heatmap_encoding=str(config.model.config.box_coder.heatmap_encoding),
                        box_coder_labels_encoding=str(config.model.config.box_coder.labels_encoding),
                        box_coder_fixed_radius=str(config.model.config.box_coder.fixed_radius),
                        box_coder_labels_radius=str(config.model.config.box_coder.labels_radius),
                        optimizer=str(config.optimizer.name),
                        optimizer_lr=float(config.optimizer.params.lr),
                        optimizer_eps=float(config.optimizer.params.get('eps', 'None')),
                        optimizer_wd=float(config.optimizer.params.weight_decay),
                        optimizer_scheduler=str(config.scheduler.scheduler_name),
                        dataset=str(config.dataset.slug),
                        dataset_fold=int(config.dataset.fold) if config.dataset.fold is not None else "None",
                        dataset_filter_low_confidence_objects=bool(config.dataset.filter_low_confidence_objects),
                        dataset_ignore_low_confidence_labels=bool(config.dataset.ignore_low_confidence_labels),
                        dataset_ignore_low_confidence_detections=bool(config.dataset.ignore_low_confidence_detections),
                        dataset_augmentations=str(config.augs.slug),
                        dataset_train_image_size=f"{config.train.train_image_size}",
                        sampler=str(config.sampler.slug),
                        sampler_num_samples=int(config.sampler.num_samples),
                        sampler_type=str(config.sampler.sampler_type),
                        sampler_balance_crowd=bool(config.sampler.balance_crowd),
                        balance_near_shore=bool(config.sampler.balance_near_shore),
                        balance_vessel_type=bool(config.sampler.balance_vessel_type),
                        balance_fishing_type=bool(config.sampler.balance_fishing_type),
                        balance_location=bool(config.sampler.balance_location),
                        loss=str(config.loss.slug),
                    )
                )
            )

        return callbacks


class TrainOnlyCriterionCallback(CriterionCallback):
    def on_batch_end(self, runner) -> None:
        if runner.is_train_loader:
            return super(TrainOnlyCriterionCallback, self).on_batch_end(runner)
        else:
            runner.batch_metrics[self.prefix] = torch.tensor(0, device="cpu")


@hydra.main(config_path="configs", config_name="multilabel_circle_net", version_base=None)
def main(cfg: DictConfig) -> None:
    """
    Entry point for training the model.
    """
    # Seed setup
    torch.manual_seed(int(cfg.seed))
    torch.set_anomaly_enabled(cfg.torch.detect_anomaly)

    # CuDNN configurations
    torch.backends.cudnn.deterministic = cfg.torch.deterministic
    torch.backends.cudnn.benchmark = cfg.torch.benchmark

    # Allow TF32 on matmul and cuDNN
    torch.backends.cuda.matmul.allow_tf32 = cfg.torch.cuda_allow_tf32
    torch.backends.cudnn.allow_tf32 = cfg.torch.cudnn.allow_tf32

    if cfg.dataset.data_dir is None:
        raise ValueError("--data-dir must be set")

    # Initialize Pipeline
    pipeline = CenternetPipeline(cfg)
    pipeline.train()


if __name__ == "__main__":
    main()