HYDRA_FULL_ERROR=1 python train_multilabel2.py +config=configs/multilabel_circle_net.yaml
DEBUG:hydra.core.utils:Setting JobRuntime:name=UNKNOWN_NAME
DEBUG:hydra.core.utils:Setting JobRuntime:name=train_multilabel2
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 919, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 119, in run
    ret = run_job(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 916, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 107, in __init__
    self.scaler = GradScaler()
  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_multilabel2.py", line 62, 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.

Dataset valid_sar_cv
        ignore_low_confidence_detections False
        ignore_low_confidence_labels     False
        filter_low_confidence_objects    False
Train dataset
        Unknown vessels 90
        Unknown fishing 11242
Valid dataset
        Unknown vessels 20
        Unknown fishing 2778
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 919, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 916, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 772, in train
    self.loaders, dataset_callbacks = self.build_loaders(config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 427, in build_loaders
    train_ds, valid_ds, train_sampler, dataset_callbacks = self.build_datasets(config)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel2.py", line 206, in build_datasets
    box_coder=self.box_coder,
AttributeError: 'CenternetPipeline' object has no attribute 'box_coder'
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 


here's train_multiablel2.py

import sys
import traceback
import warnings
from typing import Dict, Any, OrderedDict, Tuple, Optional, List
import catalyst
import os
import numpy as np
import pandas as pd
import torch
import hydra
from hydra.utils import instantiate
from torch.optim import Optimizer
from torch.utils.data import Dataset, Sampler, DataLoader
from torch import nn
from omegaconf import DictConfig
from functools import partial
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, get_optimizable_parameters
from pytorch_toolbelt.utils import fs, transfer_weights, count_parameters
from pytorch_toolbelt.utils.distributed import is_main_process
from xview3.centernet.bboxer import MultilabelCircleNetCoder
from xview3.early_stopping import EarlyStoppingCallback
from xview3.nas import ProgressiveNAS
from xview3.centernet.models.heads.yolo_head import YOLOHead, generate_yolo_targets
from xview3.centernet.models.transformer_classifier import ConvolutionalClassifier
from xview3.centernet.losses.circlenet_loss import CircleNetLoss, YOLOLoss
from xview3.centernet import *
from xview3.pipeline import *
from xview3.factory import *
from xview3.augmentations import *
from xview3.samplers import *
from xview3.dataset import *
from xview3.ema import *
from xview3.evaluation import *
from xview3.inference import *
from xview3.metric import *
from xview3.constants import INPUT_SCENE_ID_KEY, TARGET_CENTERS_KEY, TARGET_FISHING_KEY,TARGET_VESSEL_KEY,INPUT_CONFIDENCE_KEY,TARGET_LENGTHS_KEY
import logging
from torch.amp import autocast, GradScaler
from xview3.utils import TorchDevice
from xview3.evaluation import evaluate_on_scenes
from xview3.inference import predict_multilabel_scenes
from omegaconf import OmegaConf
from datetime import datetime
from combined import CombinedYOLOCenterNetModel 
from xview3.factory import get_detection_model, build_centernet_head

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))

warnings.showwarning = warn_with_traceback

if 'XVIEW3_DIR' not in os.environ:
    os.environ['XVIEW3_DIR'] = '/Volumes/Crucial/xview3_dataset'

def get_optimizer_cls(name: str):
    """
    Maps optimizer name strings to PyTorch optimizer classes.

    Args:
        name (str): Name of the optimizer (e.g., 'adamw', 'sgd', 'adam').

    Returns:
        class: Corresponding PyTorch optimizer class.

    Raises:
        ValueError: If the optimizer name is not supported.
    """
    optimizers = {
        "adamw": torch.optim.AdamW,
        "adam": torch.optim.Adam,
        "sgd": torch.optim.SGD,
        "rmsprop": torch.optim.RMSprop,
        # Add more optimizers as needed
    }
    name = name.lower()
    if name not in optimizers:
        available = ', '.join(optimizers.keys())
        raise ValueError(f"Optimizer '{name}' is not supported. Available optimizers: {available}")
    return optimizers[name]

class CenternetPipeline(Pipeline):
    def __init__(self, cfg: DictConfig):
        super().__init__(cfg)
        # Initialize scaler for mixed precision
        self.device = TorchDevice.choose_torch_device()
        
        # Initialize GradScaler based on device capabilities
        if self.device.type == 'cuda':
            self.scaler = GradScaler()
        elif self.device.type == 'mps':
            # For Apple Silicon/MPS
            self.scaler = GradScaler()
        else:
            # For CPU or other devices, disable GradScaler
            self.scaler = None
        
        self.accumulation_steps = cfg.get('accumulation_steps', 1)

    def build_datasets(self, config) -> Tuple[Dataset, Dataset, Optional[Sampler], List[Callback]]:
        dataset = MultilabelCircleNetDataModule(data_dir=os.path.join(os.environ.get('XVIEW3_DIR', '/Volumes/Crucial/xview3_dataset'), 'full_dataset'))
        
        # Split data into training and validation
        train_df, valid_df, _, shore_root = dataset.train_val_split(
            splitter=config.dataset.splitter,
            fold=config.dataset.fold,
            num_folds=config.dataset.num_folds,
        )

        self.valid_df = valid_df.copy()  # Unmodified ground truth

        # Logging data 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())

        # Data filtering based on configuration
        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")

        if config.dataset.filter_low_confidence_objects:
            train_df = filter_low_confidence_objects(train_df)
            valid_df = filter_low_confidence_objects(valid_df)
            self.master_print("Excluding low-confidence objects from training & validation")
        elif config.dataset.ignore_low_confidence_labels:
            train_df = ignore_low_confidence_objects(train_df)
            valid_df = ignore_low_confidence_objects(valid_df)
            self.master_print("Ignoring low-confidence labels from training & validation")

        normalization = build_normalization(config.normalization)

        # Define augmentations
        augmentations = [
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5),
            A.Resize(height=config.train.train_image_size[0], width=config.train.train_image_size[1], always_apply=True),
        ]

        transform = A.Compose(
            augmentations,
            keypoint_params=A.KeypointParams(
                format='xys',
                label_fields=['confidences', 'is_vessel', 'is_fishing', 'is_near_shore'],
                remove_invisible=False
            )
        )

        individual_augmentations = {
            key: A.Compose(build_augmentations([value]))
            for key, value in config.augs.individual.items()
            if value is not None
        }

        for channel in config.dataset.channels:
            if channel not in normalization:
                normalization[channel] = A.NoOp()

        # Select appropriate dataset sampler
        if config.sampler.sampler_type == "each_object":
            train_ds, train_sampler = dataset.get_random_crop_each_object_dataset(
                train_df,
                train_image_size=tuple(config.train.train_image_size),
                input_channels=config.dataset.channels,
                box_coder=self.box_coder,
                num_samples=config.sampler.num_samples,
                individual_augmentations=individual_augmentations,
                augmentations=augmentations,
                transform=transform,
                normalization=normalization,
                balance_crowd=config.sampler.balance_crowd,
                balance_near_shore=config.sampler.balance_near_shore,
                balance_vessel_type=config.sampler.balance_vessel_type,
                balance_fishing_type=config.sampler.balance_fishing_type,
                balance_per_scene=config.sampler.balance_per_scene,
                balance_location=config.sampler.balance_location,
                channels_last=config.torch.channels_last,
            )
        else:
            train_ds, train_sampler = dataset.get_random_crop_training_dataset(
                train_df,
                train_image_size=tuple(config.train.train_image_size),
                input_channels=config.dataset.channels,
                box_coder=self.box_coder,
                num_samples=config.sampler.num_samples,
                crop_around_ship_p=config.sampler.crop_around_ship_p,
                individual_augmentations=individual_augmentations,
                augmentations=augmentations,
                normalization=normalization,
                balance_near_shore=config.sampler.balance_near_shore,
                balance_crowd=config.sampler.balance_crowd,
                balance_location=config.sampler.balance_location,
                balance_folder=config.sampler.balance_folder,
                channels_last=config.torch.channels_last,
                shore_root=shore_root,
            )

        # Validation dataset
        valid_ds = dataset.get_validation_dataset(
            valid_df,
            valid_crop_size=tuple(config.train.valid_image_size),
            input_channels=config.dataset.channels,
            box_coder=self.box_coder,
            normalization=normalization,
            channels_last=config.torch.channels_last,
        )

        # Logging dataset info
        self.master_print(
            "Train",
            "instances",
            len(train_df),
            "dataset",
            len(train_ds),
            "sampler",
            len(train_sampler) if train_sampler else "N/A",
        )
        self.master_print("Valid", "instances", len(valid_df), "dataset", len(valid_ds))

        return train_ds, valid_ds, train_sampler, []

    def build_loaders(self, config) -> Tuple[Dict[str, DataLoader], List[Callback]]:
        train_ds, valid_ds, train_sampler, dataset_callbacks = self.build_datasets(config)
        
        train_loader = DataLoader(
            train_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            sampler=train_sampler if train_sampler else None,
            pin_memory=True,
            collate_fn=partial(centernet_collate, channels_last=train_ds.channels_last),
        )
        
        valid_loader = DataLoader(
            valid_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            pin_memory=True,
            collate_fn=partial(centernet_collate, channels_last=valid_ds.channels_last),
        )

        return {"train": train_loader, "valid": valid_loader}, dataset_callbacks


    def build_model(self, config: DictConfig, loaders: Dict[str, DataLoader]) -> nn.Module:
        # Instantiate the CenterNetUNetModel using Hydra's instantiate based on _target_
        model = instantiate(config.model)
        
        # Integrate NAS if enabled
        if config.get('nas', {}).get('enabled', False):
            nas = ProgressiveNAS(
                base_model=model,
                search_space=config.nas.search_space,
                evaluation_function=self.evaluate_architecture  # Ensure this method is defined
            )
            model = nas.search(
                num_iterations=config.nas.num_iterations,
                train_loader=loaders['train'],
                val_loader=loaders['valid'],
                device=self.device
            )
            logger.debug(f"Model optimized by NAS: {model}")
        
        logger.debug(f"Combined Model Built: {model}")
        return model

    def build_criterions(self, config: DictConfig) -> Tuple[Dict[str, nn.Module], List[Callback]]:
        criterions = {}
        callbacks = []
        
        # Initialize CenterNet Loss
        criterions['centernet_loss'] = CircleNetLoss().to(self.device)
        
        # Initialize YOLO Loss
        criterions['yolo_loss'] = YOLOLoss(
            num_classes=config.model.num_classes,
            anchors=config.model.anchors
        ).to(self.device)
        
        # Initialize Classification Loss (if applicable)
        criterions['classifier_loss'] = nn.CrossEntropyLoss().to(self.device)

        # Add any necessary callbacks
        return criterions, callbacks

    def build_metrics(self, config: DictConfig, loaders: Dict[str, DataLoader], model: nn.Module) -> List[Callback]:
        callbacks = super().build_metrics(config, loaders, model)
        
        # Initialize custom metrics callbacks
        metrics_callback = CircleNetRocAucMetricsCallback(prefix="train_metrics")  # Ensure this is defined
        callbacks.append(metrics_callback)
        
        return callbacks

    def build_optimizer(self, model: nn.Module, loaders: Dict[str, DataLoader]) -> Tuple[Optimizer, Any, List[Callback]]:
        optimizer_config = self.cfg.optimizer
        optimizer_use_fp16 = optimizer_config.get("fp16", False)

        optimizer_name = optimizer_config["name"].lower()

        wd_on_bias = optimizer_config.get("wd_on_bias", True)
        accumulation_steps = int(optimizer_config.get("accumulation", 1))
        optimizer_params = optimizer_config["params"]

        # Select parameters for optimizer
        if not wd_on_bias:
            optimizer_parameters = [
                {"params": [], "weight_decay": 0.0},
                {"params": [], "weight_decay": optimizer_params.get("weight_decay", 0)}
            ]
            for name, param in model.named_parameters():
                if param.requires_grad:
                    if name.endswith(".bias"):
                        optimizer_parameters[0]["params"].append(param)
                    else:
                        optimizer_parameters[1]["params"].append(param)
        else:
            optimizer_parameters = get_optimizable_parameters(model)

        # Get optimizer class
        optimizer_cls = get_optimizer_cls(optimizer_name)
        optimizer = optimizer_cls(
            optimizer_parameters,
            **optimizer_params,
        )

        # Define optimizer callbacks
        if optimizer_use_fp16:
            opt_callback = AMPOptimizerCallback(accumulation_steps=accumulation_steps)
        else:
            opt_callback = OptimizerCallback(accumulation_steps=accumulation_steps, decouple_weight_decay=False)

        # Define scheduler
        scheduler_params = self.cfg.scheduler
        num_epochs = self.cfg.train.epochs

        scheduler_name = scheduler_params["scheduler_name"].lower()
        scheduler = get_scheduler(
            optimizer,
            learning_rate=optimizer_params["lr"],
            num_epochs=num_epochs,
            batches_in_epoch=len(loaders["train"]),
            **scheduler_params,
        )

        # Define scheduler callback based on scheduler type
        if isinstance(scheduler, (torch.optim.lr_scheduler.CyclicLR, OneCycleLRWithWarmup)):
            scheduler_callback = SchedulerCallback(mode="batch")
        else:
            scheduler_callback = SchedulerCallback(mode="epoch")

        callbacks = [opt_callback, scheduler_callback]

        # Logging optimizer and scheduler info
        self.master_print("Optimizer        :", optimizer_name)
        self.master_print("Learning rate    :", optimizer_params["lr"])
        self.master_print("Weight decay     :", optimizer_params.get("weight_decay", 0))
        self.master_print("WD on bias       :", wd_on_bias)
        self.master_print("Accumulation     :", accumulation_steps)
        self.master_print("Scheduler        :", scheduler_name)

        return optimizer, scheduler, callbacks


    def build_distributed_params(self, config: DictConfig) -> Dict:
        fp16 = config.optimizer.get("fp16", False)

        if self.distributed:
            local_rank = torch.distributed.get_rank()
            distributed_params = {"rank": local_rank, "syncbn": True}
            if fp16:
                distributed_params["amp"] = True
            if self.cfg.find_unused:
                distributed_params["find_unused_parameters"] = True
        else:
            distributed_params = {"amp": fp16} if fp16 else {}

        return distributed_params

    def build_quircks(self, config: DictConfig) -> List[Callback]:
        extra_callbacks = []

        if config.ema.callback:
            cb = hydra.utils.instantiate(config.ema.callback)
            extra_callbacks.append(cb)
            self.master_print("Using EMA Callback", cb)

        return extra_callbacks

    def get_experiment_dir(self, config: DictConfig) -> str:
        if self.cfg.experiment:
            return self.cfg.experiment

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

        fold = config.dataset.get("fold")
        if fold is not None:
            experiment_slug += f"_fold{fold}"

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

    def build_loaders(self, config: DictConfig) -> Tuple[Dict[str, DataLoader], List[Callback]]:
        train_ds, valid_ds, train_sampler, dataset_callbacks = self.build_datasets(config)
        
        train_loader = DataLoader(
            train_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            sampler=train_sampler if train_sampler else None,
            pin_memory=True,
            collate_fn=partial(centernet_collate, channels_last=train_ds.channels_last),
        )
        
        valid_loader = DataLoader(
            valid_ds,
            batch_size=config.get('batch_size', 8),
            num_workers=config.get('num_workers', 4),
            pin_memory=True,
            collate_fn=partial(centernet_collate, channels_last=valid_ds.channels_last),
        )

        return {"train": train_loader, "valid": valid_loader}, dataset_callbacks

    def build_model(self, config: DictConfig, loaders: Dict[str, DataLoader]) -> nn.Module:
        # Instantiate the CenterNetUNetModel using Hydra's instantiate based on _target_
        model = instantiate(config.model)
        
        # Integrate NAS if enabled
        if config.get('nas', {}).get('enabled', False):
            nas = ProgressiveNAS(
                base_model=model,
                search_space=config.nas.search_space,
                evaluation_function=self.evaluate_architecture  # Ensure this method is defined
            )
            model = nas.search(
                num_iterations=config.nas.num_iterations,
                train_loader=loaders['train'],
                val_loader=loaders['valid'],
                device=self.device
            )
            logger.debug(f"Model optimized by NAS: {model}")
        
        logger.debug(f"Combined Model Built: {model}")
        return model

    def build_criterions(self, config: DictConfig) -> Tuple[Dict[str, nn.Module], List[Callback]]:
        criterions = {}
        callbacks = []
        
        # Initialize CenterNet Loss
        criterions['centernet_loss'] = CircleNetLoss().to(self.device)
        
        # Initialize YOLO Loss
        criterions['yolo_loss'] = YOLOLoss(
            num_classes=config.model.num_classes,
            anchors=config.model.anchors
        ).to(self.device)
        
        # Initialize Classification Loss
        criterions['classifier_loss'] = nn.CrossEntropyLoss().to(self.device)

        # Add any necessary callbacks (e.g., Mixup)
        return criterions, callbacks

    def build_metrics(self, config: DictConfig, loaders: Dict[str, DataLoader], model: nn.Module) -> List[Callback]:
        callbacks = super().build_metrics(config, loaders, model)
        
        # Initialize custom metrics callbacks
        metrics_callback = CircleNetRocAucMetricsCallback(prefix="train_metrics")  # Ensure this is defined
        callbacks.append(metrics_callback)
        
        return callbacks

    def build_optimizer(self, model: nn.Module, loaders: Dict[str, DataLoader]) -> Tuple[Optimizer, Any, List[Callback]]:
        optimizer_config = self.cfg.optimizer
        optimizer_use_fp16 = optimizer_config.get("fp16", False)

        optimizer_name = optimizer_config["name"].lower()

        wd_on_bias = optimizer_config.get("wd_on_bias", True)
        accumulation_steps = int(optimizer_config.get("accumulation", 1))
        optimizer_params = optimizer_config["params"]

        # Select parameters for optimizer
        if not wd_on_bias:
            optimizer_parameters = [
                {"params": [], "weight_decay": 0.0},
                {"params": [], "weight_decay": optimizer_params.get("weight_decay", 0)}
            ]
            for name, param in model.named_parameters():
                if param.requires_grad:
                    if name.endswith(".bias"):
                        optimizer_parameters[0]["params"].append(param)
                    else:
                        optimizer_parameters[1]["params"].append(param)
        else:
            optimizer_parameters = get_optimizable_parameters(model)

        # Get optimizer class
        optimizer_cls = get_optimizer_cls(optimizer_name)  # Implement this function
        optimizer = optimizer_cls(
            optimizer_parameters,
            **optimizer_params,
        )

        # Define optimizer callbacks
        if optimizer_use_fp16:
            opt_callback = AMPOptimizerCallback(accumulation_steps=accumulation_steps)
        else:
            opt_callback = OptimizerCallback(accumulation_steps=accumulation_steps, decouple_weight_decay=False)

        # Define scheduler
        scheduler_params = self.cfg.scheduler
        num_epochs = self.cfg.train.epochs

        scheduler_name = scheduler_params["scheduler_name"].lower()
        scheduler = get_scheduler(
            optimizer,
            learning_rate=optimizer_params["lr"],
            num_epochs=num_epochs,
            batches_in_epoch=len(loaders["train"]),
            **scheduler_params,
        )

        # Define scheduler callback based on scheduler type
        if isinstance(scheduler, (torch.optim.lr_scheduler.CyclicLR, OneCycleLRWithWarmup)):
            scheduler_callback = SchedulerCallback(mode="batch")
        else:
            scheduler_callback = SchedulerCallback(mode="epoch")

        callbacks = [opt_callback, scheduler_callback]

        # Logging optimizer and scheduler info
        self.master_print("Optimizer        :", optimizer_name)
        self.master_print("Learning rate    :", optimizer_params["lr"])
        self.master_print("Weight decay     :", optimizer_params.get("weight_decay", 0))
        self.master_print("WD on bias       :", wd_on_bias)
        self.master_print("Accumulation     :", accumulation_steps)
        self.master_print("Scheduler        :", scheduler_name)

        return optimizer, scheduler, callbacks

    def train_step(self, config ,batch: Dict[str, Any], batch_idx: int) -> Dict[str, Any]:
        self.model.train()
        if self.scaler:
            with autocast(device_type=self.device.type):
                images = batch['INPUT_IMAGE_KEY'].to(self.device)
                encoded_targets = self.encode_batch_targets(batch, self.device)
                
                # Forward pass
                outputs = self.model(images)
                
                # Compute CenterNet Loss
                centernet_loss = self.criterion['centernet_loss'](
                    outputs['centernet']['objectness_map'],
                    encoded_targets['CENTERNET']['objectness_map'],
                    outputs['centernet']['vessel_map'],
                    encoded_targets['CENTERNET']['vessel_map'],
                    outputs['centernet']['fishing_map'],
                    encoded_targets['CENTERNET']['fishing_map'],
                    outputs['centernet']['size_map'],
                    encoded_targets['CENTERNET']['size_map'],
                    outputs['centernet']['offset_map'],
                    encoded_targets['CENTERNET']['offset_map'],
                )
                
                # Compute YOLO Loss
                yolo_loss = self.criterion['yolo_loss'](
                    outputs['yolo'],
                    encoded_targets['YOLO']
                )
                
                # Compute Classification Loss
                classifier_loss = self.criterion['classifier_loss'](
                    outputs['classifier'],
                    batch['classification_labels'].to(self.device)  # Ensure this key exists
                )
                
                # Combine Losses with weights (configurable)
                total_loss = (
                    config.loss.centernet_weight * centernet_loss +
                    config.loss.yolo_weight * yolo_loss +
                    config.loss.classifier_weight * classifier_loss
                )
            
            # Scale the loss
            scaled_loss = total_loss / self.accumulation_steps
            self.scaler.scale(scaled_loss).backward()
        
            if (batch_idx + 1) % self.accumulation_steps == 0:
                self.scaler.step(self.optimizer)
                self.scaler.update()
                self.optimizer.zero_grad()

        else:
            # Without mixed precision
            images = batch['INPUT_IMAGE_KEY'].to(self.device)
            encoded_targets = self.encode_batch_targets(batch, self.device)
            outputs = self.model(images)
            # Compute CenterNet Loss
            centernet_loss = self.criterion['centernet_loss'](
                outputs['centernet']['objectness_map'],
                encoded_targets['CENTERNET']['objectness_map'],
                outputs['centernet']['vessel_map'],
                encoded_targets['CENTERNET']['vessel_map'],
                outputs['centernet']['fishing_map'],
                encoded_targets['CENTERNET']['fishing_map'],
                outputs['centernet']['size_map'],
                encoded_targets['CENTERNET']['size_map'],
                outputs['centernet']['offset_map'],
                encoded_targets['CENTERNET']['offset_map'],
            )
            
            # Compute YOLO Loss
            yolo_loss = self.criterion['yolo_loss'](
                outputs['yolo'],
                encoded_targets['YOLO']
            )
            
            # Compute Classification Loss
            classifier_loss = self.criterion['classifier_loss'](
                outputs['classifier'],
                batch['classification_labels'].to(self.device)  # Ensure this key exists
            )
            
            # Combine Losses with weights (configurable)
            total_loss = (
                config.loss.centernet_weight * centernet_loss +
                config.loss.yolo_weight * yolo_loss +
                config.loss.classifier_weight * classifier_loss
            )
            total_loss = centernet_loss + yolo_loss

            # Backpropagate normally
            total_loss.backward()
        
            if (batch_idx + 1) % self.accumulation_steps == 0:
                self.optimizer.step()
                self.optimizer.zero_grad()
        
        return {'loss': total_loss}
    

    def evaluate_architecture(self, model: nn.Module, train_loader: DataLoader, val_loader: DataLoader, device: torch.device) -> float:
        if model is None or not hasattr(self, 'criterion'):
            raise RuntimeError("Model or criterion not initialized.")
        
        model.to(device)
        model.eval()
        total_centernet_loss = 0.0
        total_yolo_loss = 0.0
        total_classifier_loss = 0.0
        total_samples = 0

        with torch.no_grad():
            for batch in val_loader:
                images = batch['INPUT_IMAGE_KEY'].to(device)
                encoded_targets = self.encode_batch_targets(batch, device)

                # Forward pass
                outputs = model(images)

                # Compute CenterNet Loss
                centernet_loss = self.criterion['centernet_loss'](
                    outputs['centernet']['objectness_map'],
                    encoded_targets['CENTERNET']['objectness_map'],
                    outputs['centernet']['vessel_map'],
                    encoded_targets['CENTERNET']['vessel_map'],
                    outputs['centernet']['fishing_map'],
                    encoded_targets['CENTERNET']['fishing_map'],
                    outputs['centernet']['size_map'],
                    encoded_targets['CENTERNET']['size_map'],
                    outputs['centernet']['offset_map'],
                    encoded_targets['CENTERNET']['offset_map'],
                )

                # Compute YOLO Loss
                yolo_loss = self.criterion['yolo_loss'](
                    outputs['yolo'],
                    encoded_targets['YOLO']
                )

                # Compute Classification Loss
                classifier_loss = self.criterion['classifier_loss'](
                    outputs['classifier'],
                    batch['classification_labels'].to(device)  # Ensure this key exists
                )

                # Aggregate Losses
                total_centernet_loss += centernet_loss.item() * images.size(0)
                total_yolo_loss += yolo_loss.item() * images.size(0)
                total_classifier_loss += classifier_loss.item() * images.size(0)
                total_samples += images.size(0)

        average_centernet_loss = total_centernet_loss / total_samples
        average_yolo_loss = total_yolo_loss / total_samples
        average_classifier_loss = total_classifier_loss / total_samples

        logger.debug(f"[NAS Evaluation] Average CenterNet Loss: {average_centernet_loss}")
        logger.debug(f"[NAS Evaluation] Average YOLO Loss: {average_yolo_loss}")
        logger.debug(f"[NAS Evaluation] Average Classifier Loss: {average_classifier_loss}")

        # Define a metric to minimize; lower loss is better
        return average_centernet_loss + average_yolo_loss + average_classifier_loss

    def encode_batch_targets(self, batch: Dict[str, Any], device: torch.device) -> Dict[str, Any]:
        """
        Encodes both CenterNet and YOLO targets from the batch.
        """
        # Extract relevant batch components
        centers = batch[TARGET_CENTERS_KEY].cpu().numpy()
        confidences = batch[INPUT_CONFIDENCE_KEY]
        is_vessel = batch[TARGET_VESSEL_KEY].cpu().numpy()
        is_fishing = batch[TARGET_FISHING_KEY].cpu().numpy()
        lengths = batch[TARGET_LENGTHS_KEY].cpu().numpy()

        # Heatmap and other CenterNet targets are already in batch
        heatmap = batch['centernet_objectness_map'].to(device)
        vessel_map = batch['centernet_vessel_map'].to(device)
        fishing_map = batch['centernet_fishing_map'].to(device)
        size_map = batch['centernet_size_map'].to(device)
        offset_map = batch['centernet_offset_map'].to(device)

        # YOLO targets
        yolo_targets = batch['yolo_targets'].to(device)

        # Classification labels
        classification_labels = batch['classification_labels'].to(device)

        encoded_targets = {
            'CENTERNET': {
                'objectness_map': heatmap,
                'vessel_map': vessel_map,
                'fishing_map': fishing_map,
                'size_map': size_map,
                'offset_map': offset_map,
            },
            'YOLO': yolo_targets,
            'classifier_labels': classification_labels
        }

        return encoded_targets

    def train(self):
        config = self.cfg

        # Build data loaders
        self.loaders, dataset_callbacks = self.build_loaders(config)

        # Build model
        model = self.build_model(config, self.loaders)
        self.model = model  # Assign to self.model for access in other methods

        # Build criterions
        criterions, criterions_callbacks = self.build_criterions(config)

        # Build metrics
        metrics_callbacks = self.build_metrics(config, self.loaders, model)

        # Build optimizer
        optimizer, scheduler, optimizer_callbacks = self.build_optimizer(model, self.loaders)

        # Combine all callbacks
        all_callbacks = dataset_callbacks + criterions_callbacks + metrics_callbacks + optimizer_callbacks + [TimerCallback()]

        # Get experiment directory
        experiment_dir = self.get_experiment_dir(config)
        self.experiment_dir = experiment_dir

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

        # Setup runner
        runner_config = config.get("runner", {})
        input_key = runner_config.get("input_key", INPUT_SCENE_ID_KEY)
        output_key = runner_config.get("output_key", None)

        if input_key and not isinstance(input_key, str):
            input_key = OmegaConf.to_container(input_key, resolve=True)
        if output_key and not isinstance(output_key, str):
            output_key = OmegaConf.to_container(output_key, resolve=True)

        runner = catalyst.runners.SupervisedRunner(
            input_key=input_key,
            output_key=output_key,
            device=self.device
        )

        # Initialize GradScaler for mixed precision training
        self.scaler = GradScaler()

        # Set up gradient accumulation
        self.accumulation_steps = config.get('accumulation_steps', 1)

        # Prepare all callbacks
        all_callbacks += self.build_quircks(config)

        # Train the model
        runner.train(
            model=model,
            criterion=criterions,  # Use the initialized criterions
            optimizer=optimizer,
            scheduler=scheduler,
            loaders=self.loaders,
            callbacks=all_callbacks,
            logdir=experiment_dir,
            num_epochs=config.train.epochs,
            verbose=True,
            main_metric=config.train.main_metric,
            minimize_metric=True,  # Assuming lower loss is better
            checkpoint_data=self.get_checkpoint_data(),
            # Use the train_step defined for custom training behavior
            # Note: Catalyst's Runner handles the training loop; ensure compatibility
        )

        # Finalize experiment
        self.on_experiment_finished()

    def get_checkpoint_data(self):
        return {
            "config": OmegaConf.to_container(self.cfg, resolve=True),
            "version": get_version(),
        }

    def on_experiment_finished(self):
        pass

    def build_quircks(self, config: DictConfig) -> List[Callback]:
        extra_callbacks = []

        if config.ema.callback:
            cb = hydra.utils.instantiate(config.ema.callback)
            extra_callbacks.append(cb)
            self.master_print("Using EMA Callback", cb)

        return extra_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="mps")

class CombinedYOLOCenterNetModel(nn.Module):
    def __init__(self, backbone, centernet_head, yolo_head):
        super().__init__()
        self.backbone = backbone
        self.centernet_head = centernet_head
        self.yolo_head = yolo_head
    
    def forward(self, x):
        features = self.backbone(x)
        
        # CenterNet predictions
        centernet_outputs = self.centernet_head(features)
        
        # YOLO predictions
        yolo_outputs = self.yolo_head(features)
        
        # Combine outputs into a dictionary
        return {
            'CENTERNET_OUTPUT_OBJECTNESS_MAP': centernet_outputs['objectness_map'],
            'CENTERNET_OUTPUT_VESSEL_MAP': centernet_outputs['vessel_map'],
            'CENTERNET_OUTPUT_FISHING_MAP': centernet_outputs['fishing_map'],
            'CENTERNET_OUTPUT_SIZE': centernet_outputs['size_map'],
            'CENTERNET_OUTPUT_OFFSET': centernet_outputs['offset_map'],
            'CENTERNET_OUTPUT_YOLO': yolo_outputs,
        }
    
@hydra.main(config_path="configs", config_name="multilabel_circle_net", version_base=None)
def main(cfg: DictConfig) -> None:
    torch.mps.empty_cache()
    torch.manual_seed(int(cfg.seed))

    torch.set_anomaly_enabled(cfg.torch.detect_anomaly)

    torch.backends.cudnn.deterministic = cfg.torch.deterministic
    torch.backends.cudnn.benchmark = cfg.torch.benchmark

    # The flag below controls whether to allow TF32 on matmul. This flag defaults to True.
    torch.backends.cuda.matmul.allow_tf32 = cfg.torch.cuda_allow_tf32
    # The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
    torch.backends.cudnn.allow_tf32 = cfg.torch.cudnn.allow_tf32

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

    CenternetPipeline(cfg).train()

if __name__ == "__main__":
    main()

#test