autogluon.multimodal 1.2.1b20250303__py3-none-any.whl → 1.2.1b20250305__py3-none-any.whl

autogluon/multimodal/learners/few_shot_svm.py

@@ -18,19 +18,10 @@ from autogluon.core.metrics import Scorer
 from autogluon.core.utils.loaders import load_pd
 
 from ..constants import CLIP, COLUMN_FEATURES, HF_TEXT, TIMM_IMAGE, Y_PRED, Y_TRUE
-from ..data import BaseDataModule, MultiModalFeaturePreprocessor
-from ..utils import (
-    CustomUnpickler,
-    LogFilter,
-    apply_log_filter,
-    compute_score,
-    data_to_df,
-    extract_from_output,
-    get_available_devices,
-    logits_to_prob,
-    select_model,
-    turn_on_off_feature_column_info,
-)
+from ..data import BaseDataModule, MultiModalFeaturePreprocessor, data_to_df, turn_on_off_feature_column_info
+from ..models import select_model
+from ..optim import compute_score
+from ..utils import LogFilter, apply_log_filter, extract_from_output, get_available_devices, logits_to_prob
 from .base import BaseLearner
 
 logger = logging.getLogger(__name__)
@@ -62,7 +53,7 @@ class FewShotSVMLearner(BaseLearner):
                     "model.hf_text.checkpoint_name": "sentence-transformers/all-mpnet-base-v2",
                     "model.hf_text.pooling_mode": "mean",
                     "env.per_gpu_batch_size": 32,
-                    "env.eval_batch_size_ratio": 4,
+                    "env.inference_batch_size_ratio": 4,
                 }
         presets
             Presets regarding model quality, e.g., best_quality, high_quality, and medium_quality.
@@ -589,7 +580,7 @@ class FewShotSVMLearner(BaseLearner):
     ):
         predictor = super().load(path=path, resume=resume, verbosity=verbosity)
         with open(os.path.join(path, "svm.pkl"), "rb") as fp:
-            params = CustomUnpickler(fp).load()
+            params = pickle.load(fp)  # nosec B301
         svm = make_pipeline(StandardScaler(), SVC(gamma="auto"))
         svm.set_params(**params)
         predictor._svm = svm

autogluon/multimodal/learners/matching.py

@@ -46,41 +46,43 @@ from ..constants import (
     Y_PRED_PROB,
     Y_TRUE,
 )
-from ..data import BaseDataModule, MultiModalFeaturePreprocessor, infer_column_types
-from ..optimization import MatcherLitModule, get_matcher_loss_func, get_matcher_miner_func, get_metric
-from ..presets import matcher_presets
-from ..problem_types import PROBLEM_TYPES_REG
-from ..utils import (
-    CustomUnpickler,
-    assign_feature_column_names,
-    average_checkpoints,
+from ..data import (
+    BaseDataModule,
+    MultiModalFeaturePreprocessor,
+    create_fusion_data_processors,
+    data_to_df,
+    infer_column_types,
+    infer_dtypes_by_model_names,
+    init_df_preprocessor,
+)
+from ..models import is_lazy_weight_tensor, select_model
+from ..optim import (
+    MatcherLitModule,
     compute_ranking_score,
     compute_score,
+    get_matcher_loss_func,
+    get_matcher_miner_func,
+    get_torchmetric,
+)
+from ..utils import (
+    average_checkpoints,
     compute_semantic_similarity,
     convert_data_for_ranking,
-    create_fusion_data_processors,
     create_siamese_model,
     customize_model_names,
-    data_to_df,
     extract_from_output,
     get_config,
     get_dir_ckpt_paths,
     get_load_ckpt_paths,
     get_local_pretrained_config_paths,
-    get_minmax_mode,
     hyperparameter_tune,
-    infer_dtypes_by_model_names,
-    init_df_preprocessor,
-    is_lazy_weight_tensor,
-    load_text_tokenizers,
+    matcher_presets,
     on_fit_end_message,
     save_pretrained_model_configs,
-    save_text_tokenizers,
-    select_model,
     split_hyperparameters,
     update_config_by_rules,
-    upgrade_config,
 )
+from ..utils.problem_types import PROBLEM_TYPES_REG
 from .base import BaseLearner
 
 pl_logger = logging.getLogger("lightning")
@@ -88,9 +90,9 @@ pl_logger.propagate = False  # https://github.com/Lightning-AI/lightning/issues/
 logger = logging.getLogger(__name__)
 
 
-class MultiModalMatcher(BaseLearner):
+class MatchingLearner(BaseLearner):
     """
-    MultiModalMatcher is a framework to learn/extract embeddings for multimodal data including image, text, and tabular.
+    MatchingLearner is a framework to learn/extract embeddings for multimodal data including image, text, and tabular.
     These embeddings can be used e.g. with cosine-similarity to find items with similar semantic meanings.
     This can be useful for computing the semantic similarity of two items, semantic search, paraphrase mining, etc.
     """
@@ -448,7 +450,7 @@ class MultiModalMatcher(BaseLearner):
             # top_k_average is called inside hyperparameter_tune() when building the final predictor.
             self.top_k_average(
                 save_path=self._save_path,
-                top_k_average_method=self._config.optimization.top_k_average_method,
+                top_k_average_method=self._config.optim.top_k_average_method,
                 standalone=standalone,
                 clean_ckpts=clean_ckpts,
             )
@@ -476,7 +478,7 @@ class MultiModalMatcher(BaseLearner):
         **kwargs,
     ):
         """
-        Fit MultiModalMatcher. Train the model to learn embeddings to simultaneously maximize and minimize
+        Fit MatchingLearner. Train the model to learn embeddings to simultaneously maximize and minimize
         the semantic similarities of positive and negative pairs.
         The data may contain image, text, numeric, or categorical features.
 
@@ -538,7 +540,7 @@ class MultiModalMatcher(BaseLearner):
 
         Returns
         -------
-        An "MultiModalMatcher" object (itself).
+        An "MatchingLearner" object (itself).
         """
         self.setup_save_path(save_path=save_path)
         training_start = self.on_fit_start(presets=presets)
@@ -805,7 +807,7 @@ class MultiModalMatcher(BaseLearner):
             label_processors_count = {k: len(v) for k, v in label_processors.items()}
             logger.debug(f"label_processors_count: {label_processors_count}")
 
-        validation_metric, custom_metric_func = get_metric(
+        validation_metric, custom_metric_func = get_torchmetric(
             metric_name=self._validation_metric_name,
             num_classes=self._output_shape,
             is_matching=self._pipeline in matcher_presets.list_keys(),
@@ -863,17 +865,17 @@ class MultiModalMatcher(BaseLearner):
             validate_data=self._tuning_data,
             id_mappings=id_mappings,
         )
-        optimization_kwargs = dict(
-            optim_type=config.optimization.optim_type,
-            lr_choice=config.optimization.lr_choice,
-            lr_schedule=config.optimization.lr_schedule,
-            lr=config.optimization.learning_rate,
-            lr_decay=config.optimization.lr_decay,
-            end_lr=config.optimization.end_lr,
-            lr_mult=config.optimization.lr_mult,
-            weight_decay=config.optimization.weight_decay,
-            warmup_steps=config.optimization.warmup_steps,
-            track_grad_norm=OmegaConf.select(config, "optimization.track_grad_norm", default=-1),
+        optim_kwargs = dict(
+            optim_type=config.optim.optim_type,
+            lr_choice=config.optim.lr_choice,
+            lr_schedule=config.optim.lr_schedule,
+            lr=config.optim.lr,
+            lr_decay=config.optim.lr_decay,
+            end_lr=config.optim.end_lr,
+            lr_mult=config.optim.lr_mult,
+            weight_decay=config.optim.weight_decay,
+            warmup_steps=config.optim.warmup_steps,
+            track_grad_norm=config.optim.track_grad_norm,
         )
         metrics_kwargs = dict(
             validation_metric=validation_metric,
@@ -893,7 +895,7 @@ class MultiModalMatcher(BaseLearner):
             loss_func=loss_func,
             miner_func=miner_func,
             **metrics_kwargs,
-            **optimization_kwargs,
+            **optim_kwargs,
         )
         callbacks = self.get_callbacks_per_run(save_path=save_path, config=config, litmodule=litmodule)
         tb_logger = self.get_tb_logger(save_path=save_path)
@@ -1025,7 +1027,7 @@ class MultiModalMatcher(BaseLearner):
                     ingredients = [top_k_model_paths[0]]
                 else:
                     raise ValueError(
-                        f"The key for 'optimization.top_k_average_method' is not supported. "
+                        f"The key for 'optim.top_k_average_method' is not supported. "
                         f"We only support '{GREEDY_SOUP}', '{UNIFORM_SOUP}' and '{BEST}'. "
                         f"The provided value is '{top_k_average_method}'."
                     )
@@ -1202,7 +1204,7 @@ class MultiModalMatcher(BaseLearner):
             df_preprocessor=df_preprocessor,
             data_processors=data_processors,
             per_gpu_batch_size=batch_size,
-            num_workers=self._config.env.num_workers_evaluation,
+            num_workers=self._config.env.num_workers_inference,
             predict_data=data,
             id_mappings=id_mappings,
         )
@@ -1931,18 +1933,14 @@ class MultiModalMatcher(BaseLearner):
         # Save text tokenizers before saving data processors
         query_processors = copy.deepcopy(query_processors)
         if TEXT in query_processors:
-            query_processors[TEXT] = save_text_tokenizers(
-                text_processors=query_processors[TEXT],
-                path=path,
-            )
+            for per_text_processor in query_processors[TEXT]:
+                per_text_processor.save_tokenizer(path)
 
         # Save text tokenizers before saving data processors
         response_processors = copy.deepcopy(response_processors)
         if TEXT in response_processors:
-            response_processors[TEXT] = save_text_tokenizers(
-                text_processors=response_processors[TEXT],
-                path=path,
-            )
+            for per_text_processor in response_processors[TEXT]:
+                per_text_processor.save_tokenizer(path)
 
         data_processors = {
             QUERY: query_processors,
@@ -1955,7 +1953,7 @@ class MultiModalMatcher(BaseLearner):
         with open(os.path.join(path, f"assets.json"), "w") as fp:
             json.dump(
                 {
-                    "class_name": self.__class__.__name__,
+                    "learner_class": self.__class__.__name__,
                     "query": self._query,
                     "response": self._response,
                     "match_label": self._match_label,
@@ -1990,7 +1988,7 @@ class MultiModalMatcher(BaseLearner):
 
     @staticmethod
     def _load_metadata(
-        matcher: MultiModalMatcher,
+        matcher: MatchingLearner,
         path: str,
         resume: Optional[bool] = False,
         verbosity: Optional[int] = 3,
@@ -2013,11 +2011,8 @@ class MultiModalMatcher(BaseLearner):
         with open(os.path.join(path, "assets.json"), "r") as fp:
             assets = json.load(fp)
 
-        query_config = upgrade_config(query_config, assets["version"])
-        response_config = upgrade_config(response_config, assets["version"])
-
         with open(os.path.join(path, "df_preprocessor.pkl"), "rb") as fp:
-            df_preprocessor = CustomUnpickler(fp).load()
+            df_preprocessor = pickle.load(fp)  # nosec B301
 
         query_df_preprocessor = df_preprocessor[QUERY]
         response_df_preprocessor = df_preprocessor[RESPONSE]
@@ -2025,7 +2020,7 @@ class MultiModalMatcher(BaseLearner):
 
         try:
             with open(os.path.join(path, "data_processors.pkl"), "rb") as fp:
-                data_processors = CustomUnpickler(fp).load()
+                data_processors = pickle.load(fp)  # nosec B301
 
             query_processors = data_processors[QUERY]
             response_processors = data_processors[RESPONSE]
@@ -2033,32 +2028,20 @@ class MultiModalMatcher(BaseLearner):
 
             # Load text tokenizers after loading data processors.
             if TEXT in query_processors:
-                query_processors[TEXT] = load_text_tokenizers(
-                    text_processors=query_processors[TEXT],
-                    path=path,
-                )
-            # backward compatibility. Add feature column names in each data processor.
-            query_processors = assign_feature_column_names(
-                data_processors=query_processors,
-                df_preprocessor=query_df_preprocessor,
-            )
+                for per_text_processor in query_processors[TEXT]:
+                    per_text_processor.load_tokenizer(path)
+
             # Only keep the modalities with non-empty processors.
             query_processors = {k: v for k, v in query_processors.items() if len(v) > 0}
 
             # Load text tokenizers after loading data processors.
             if TEXT in response_processors:
-                response_processors[TEXT] = load_text_tokenizers(
-                    text_processors=response_processors[TEXT],
-                    path=path,
-                )
-            # backward compatibility. Add feature column names in each data processor.
-            response_processors = assign_feature_column_names(
-                data_processors=response_processors,
-                df_preprocessor=response_df_preprocessor,
-            )
+                for per_text_processor in response_processors[TEXT]:
+                    per_text_processor.load_tokenizer(path)
+
             # Only keep the modalities with non-empty processors.
             response_processors = {k: v for k, v in response_processors.items() if len(v) > 0}
-        except:  # backward compatibility. reconstruct the data processor in case something went wrong.
+        except:  # reconstruct the data processor in case something went wrong.
             query_processors = None
             response_processors = None
             label_processors = None
@@ -2069,19 +2052,14 @@ class MultiModalMatcher(BaseLearner):
         matcher._label_column = assets["label_column"]
         matcher._problem_type = assets["problem_type"]
         matcher._pipeline = assets["pipeline"]
-        if "presets" in assets:
-            matcher._presets = assets["presets"]
+        matcher._presets = assets["presets"]
         matcher._eval_metric_name = assets["eval_metric_name"]
         matcher._verbosity = verbosity
         matcher._resume = resume
         matcher._save_path = path  # in case the original exp dir is copied to somewhere else
         matcher._pretrained_path = path
-        if "pretrained" in assets:
-            matcher._pretrained = assets["pretrained"]
-        if "fit_called" in assets:
-            matcher._fit_called = assets["fit_called"]
-        else:
-            matcher._fit_called = True  # backward compatible
+        matcher._pretrained = assets["pretrained"]
+        matcher._fit_called = assets["fit_called"]
         matcher._config = config
         matcher._query_config = query_config
         matcher._response_config = response_config
@@ -2094,10 +2072,7 @@ class MultiModalMatcher(BaseLearner):
         matcher._query_processors = query_processors
         matcher._response_processors = response_processors
         matcher._label_processors = label_processors
-        if "minmax_mode" in assets:
-            matcher._minmax_mode = assets["minmax_mode"]
-        else:
-            matcher._minmax_mode = get_minmax_mode(matcher._validation_metric_name)
+        matcher._minmax_mode = assets["minmax_mode"]
 
         return matcher
 

autogluon/multimodal/learners/ner.py

@@ -7,15 +7,16 @@ from typing import Callable, Dict, List, Optional, Union
 
 import lightning.pytorch as pl
 import pandas as pd
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import DictConfig
 from torch import nn
 
 from autogluon.core.metrics import Scorer
 
 from ..constants import NER, NER_RET, Y_PRED, Y_TRUE
 from ..data import MultiModalFeaturePreprocessor
-from ..optimization import NerLitModule, get_metric
-from ..utils import compute_score, create_fusion_model, extract_from_output, merge_bio_format
+from ..models import create_fusion_model
+from ..optim import NerLitModule, compute_score, get_minmax_mode, get_torchmetric, infer_metrics
+from ..utils import extract_from_output, merge_bio_format
 from .base import BaseLearner
 
 logger = logging.getLogger(__name__)
@@ -86,7 +87,7 @@ class NERLearner(BaseLearner):
             self._output_shape = output_shape  # since ner infers output_shape in fit_per_run(), the learners needs to update the attribute afterwards.
 
     def get_validation_metric_per_run(self, output_shape: int):
-        validation_metric, custom_metric_func = get_metric(
+        validation_metric, custom_metric_func = get_torchmetric(
             metric_name=self._validation_metric_name,
             num_classes=output_shape,
             problem_type=self._problem_type,
@@ -109,38 +110,38 @@ class NERLearner(BaseLearner):
             )
         return model
 
-    def get_optimization_kwargs_per_run(self, config, validation_metric, custom_metric_func, loss_func):
+    def get_optim_kwargs_per_run(self, config, validation_metric, custom_metric_func, loss_func):
         return dict(
-            optim_type=config.optimization.optim_type,
-            lr_choice=config.optimization.lr_choice,
-            lr_schedule=config.optimization.lr_schedule,
-            lr=config.optimization.learning_rate,
-            lr_decay=config.optimization.lr_decay,
-            end_lr=config.optimization.end_lr,
-            lr_mult=config.optimization.lr_mult,
-            weight_decay=config.optimization.weight_decay,
-            warmup_steps=config.optimization.warmup_steps,
-            track_grad_norm=OmegaConf.select(config, "optimization.track_grad_norm", default=-1),
+            optim_type=config.optim.optim_type,
+            lr_choice=config.optim.lr_choice,
+            lr_schedule=config.optim.lr_schedule,
+            lr=config.optim.lr,
+            lr_decay=config.optim.lr_decay,
+            end_lr=config.optim.end_lr,
+            lr_mult=config.optim.lr_mult,
+            weight_decay=config.optim.weight_decay,
+            warmup_steps=config.optim.warmup_steps,
+            track_grad_norm=config.optim.track_grad_norm,
             validation_metric=validation_metric,
             validation_metric_name=self._validation_metric_name,
             custom_metric_func=custom_metric_func,
             loss_func=loss_func,
-            efficient_finetune=OmegaConf.select(config, "optimization.efficient_finetune"),
-            skip_final_val=OmegaConf.select(config, "optimization.skip_final_val", default=False),
+            peft=config.optim.peft,
+            skip_final_val=config.optim.skip_final_val,
         )
 
     def get_litmodule_per_run(
         self,
         model: Optional[nn.Module] = None,
         peft_param_names: Optional[List[str]] = None,
-        optimization_kwargs: Optional[dict] = None,
+        optim_kwargs: Optional[dict] = None,
         is_train=True,
     ):
         if is_train:
             return NerLitModule(
                 model=model,
                 trainable_param_names=peft_param_names,
-                **optimization_kwargs,
+                **optim_kwargs,
             )
         else:
             return NerLitModule(model=self._model)
@@ -214,7 +215,7 @@ class NERLearner(BaseLearner):
             advanced_hyperparameters=advanced_hyperparameters,
         )
         validation_metric, custom_metric_func = self.get_validation_metric_per_run(output_shape=output_shape)
-        loss_func = self.get_loss_func_per_run(config=config)
+        loss_func, _ = self.get_loss_func_per_run(config=config)
         if max_time == timedelta(seconds=0):
             return dict(
                 config=config,
@@ -230,7 +231,7 @@ class NERLearner(BaseLearner):
             per_gpu_batch_size=config.env.per_gpu_batch_size,
             num_workers=config.env.num_workers,
         )
-        optimization_kwargs = self.get_optimization_kwargs_per_run(
+        optim_kwargs = self.get_optim_kwargs_per_run(
             config=config,
             validation_metric=validation_metric,
             custom_metric_func=custom_metric_func,
@@ -239,7 +240,7 @@ class NERLearner(BaseLearner):
         litmodule = self.get_litmodule_per_run(
             model=model,
             peft_param_names=peft_param_names,
-            optimization_kwargs=optimization_kwargs,
+            optim_kwargs=optim_kwargs,
         )
         callbacks = self.get_callbacks_per_run(save_path=save_path, config=config, litmodule=litmodule)
         plugins = self.get_plugins_per_run(model=model, peft_param_names=peft_param_names)

autogluon/multimodal/learners/object_detection.py

@@ -5,25 +5,30 @@ from datetime import timedelta
 from typing import Dict, List, Optional, Union
 
 import pandas as pd
-from omegaconf import DictConfig, OmegaConf
+from omegaconf import DictConfig
 from torch import nn
 
 from ..constants import BBOX, DDP, MAP, MULTI_IMAGE_MIX_DATASET, OBJECT_DETECTION, XYWH
-from ..data import BaseDataModule, MultiImageMixDataset, MultiModalFeaturePreprocessor, infer_rois_column_type
-from ..optimization import LitModule, MMDetLitModule
+from ..data import (
+    BaseDataModule,
+    MultiImageMixDataset,
+    MultiModalFeaturePreprocessor,
+    infer_rois_column_type,
+    split_train_tuning_data,
+)
+from ..models import create_fusion_model
+from ..optim import MMDetLitModule
 from ..utils import (
     check_if_packages_installed,
     cocoeval,
     convert_pred_to_xywh,
     convert_result_df,
-    create_fusion_model,
     extract_from_output,
     from_coco_or_voc,
     get_detection_classes,
     object_detection_data_to_df,
     save_result_coco_format,
     setup_save_path,
-    split_train_tuning_data,
 )
 from .base import BaseLearner
 
@@ -303,18 +308,18 @@ class ObjectDetectionLearner(BaseLearner):
 
         return num_gpus
 
-    def get_optimization_kwargs_per_run(self, config, validation_metric, custom_metric_func):
+    def get_optim_kwargs_per_run(self, config, validation_metric, custom_metric_func):
         return dict(
-            optim_type=config.optimization.optim_type,
-            lr_choice=config.optimization.lr_choice,
-            lr_schedule=config.optimization.lr_schedule,
-            lr=config.optimization.learning_rate,
-            lr_decay=config.optimization.lr_decay,
-            end_lr=config.optimization.end_lr,
-            lr_mult=config.optimization.lr_mult,
-            weight_decay=config.optimization.weight_decay,
-            warmup_steps=config.optimization.warmup_steps,
-            track_grad_norm=OmegaConf.select(config, "optimization.track_grad_norm", default=-1),
+            optim_type=config.optim.optim_type,
+            lr_choice=config.optim.lr_choice,
+            lr_schedule=config.optim.lr_schedule,
+            lr=config.optim.lr,
+            lr_decay=config.optim.lr_decay,
+            end_lr=config.optim.end_lr,
+            lr_mult=config.optim.lr_mult,
+            weight_decay=config.optim.weight_decay,
+            warmup_steps=config.optim.warmup_steps,
+            track_grad_norm=config.optim.track_grad_norm,
             validation_metric=validation_metric,
             validation_metric_name=self._validation_metric_name,
             custom_metric_func=custom_metric_func,
@@ -323,7 +328,7 @@ class ObjectDetectionLearner(BaseLearner):
     def get_litmodule_per_run(
         self,
         model: Optional[nn.Module] = None,
-        optimization_kwargs: Optional[dict] = None,
+        optim_kwargs: Optional[dict] = None,
         is_train=True,
     ):
         if self._problem_type == OBJECT_DETECTION:
@@ -334,7 +339,7 @@ class ObjectDetectionLearner(BaseLearner):
         if is_train:
             return LightningModule(
                 model=model,
-                **optimization_kwargs,
+                **optim_kwargs,
             )
         else:
             return LightningModule(model=self._model)
@@ -395,14 +400,14 @@ class ObjectDetectionLearner(BaseLearner):
             num_workers=config.env.num_workers,
             model_config=model.config,
         )
-        optimization_kwargs = self.get_optimization_kwargs_per_run(
+        optim_kwargs = self.get_optim_kwargs_per_run(
             config=config,
             validation_metric=validation_metric,
             custom_metric_func=custom_metric_func,
         )
         litmodule = self.get_litmodule_per_run(
             model=model,
-            optimization_kwargs=optimization_kwargs,
+            optim_kwargs=optim_kwargs,
         )
         callbacks = self.get_callbacks_per_run(save_path=save_path, config=config, litmodule=litmodule)
         plugins = self.get_plugins_per_run(model=model)
@@ -524,7 +529,7 @@ class ObjectDetectionLearner(BaseLearner):
             df_preprocessor=df_preprocessor,
             data_processors=data_processors,
             per_gpu_batch_size=batch_size,
-            num_workers=self._config.env.num_workers_evaluation,
+            num_workers=self._config.env.num_workers_inference,
             predict_data=data,
             is_train=False,
         )

autogluon/multimodal/learners/semantic_segmentation.py

@@ -7,19 +7,17 @@ import numpy as np
 import pandas as pd
 import torch
 import torch.nn.functional as F
-from omegaconf import OmegaConf
 from PIL import Image
 from scipy.special import softmax
 
 from autogluon.core.metrics import Scorer
 
-from ..constants import LABEL, LOGITS, SEMANTIC_MASK, SEMANTIC_SEGMENTATION, SEMANTIC_SEGMENTATION_IMG
-from ..optimization.lit_semantic_seg import SemanticSegmentationLitModule
-from ..optimization.semantic_seg_metrics import Balanced_Error_Rate_Pred as Balanced_Error_Rate
-from ..optimization.semantic_seg_metrics import Binary_IoU_Pred as Binary_IoU
-from ..optimization.semantic_seg_metrics import COD_METRICS_NAMES_Pred as COD_METRICS_NAMES
-from ..optimization.semantic_seg_metrics import Multiclass_IoU_Pred as Multiclass_IoU
-from ..optimization.utils import get_loss_func, get_norm_layer_param_names, get_trainable_params_efficient_finetune
+from ..constants import LABEL, LOGITS, SEMANTIC_MASK, SEMANTIC_SEGMENTATION_IMG
+from ..optim import SemanticSegmentationLitModule, get_loss_func, get_norm_layer_param_names, get_peft_param_names
+from ..optim.metrics.semantic_seg_metrics import Balanced_Error_Rate_Pred as Balanced_Error_Rate
+from ..optim.metrics.semantic_seg_metrics import Binary_IoU_Pred as Binary_IoU
+from ..optim.metrics.semantic_seg_metrics import COD_METRICS_NAMES_Pred as COD_METRICS_NAMES
+from ..optim.metrics.semantic_seg_metrics import Multiclass_IoU_Pred as Multiclass_IoU
 from ..utils import extract_from_output, setup_save_path
 from .base import BaseLearner
 
@@ -122,24 +120,24 @@ class SemanticSegmentationLearner(BaseLearner):
     @staticmethod
     def get_peft_param_names_per_run(model, config):
         peft_param_names = None
-        peft = OmegaConf.select(config, "optimization.efficient_finetune")
+        peft = config.optim.peft
         if peft:
             norm_param_names = get_norm_layer_param_names(model)
-            peft_param_names = get_trainable_params_efficient_finetune(
+            peft_param_names = get_peft_param_names(
                 norm_param_names,
-                efficient_finetune=peft,
-                extra_params=OmegaConf.select(config, "optimization.extra_trainable_params"),
+                peft=peft,
+                extra_params=config.optim.extra_trainable_params,
             )
         return peft_param_names
 
     def get_loss_func_per_run(self, config, mixup_active=None):
         loss_func = get_loss_func(
             problem_type=self._problem_type,
-            loss_func_name=OmegaConf.select(config, "optimization.loss_function"),
-            config=config.optimization,
+            loss_func_name=config.optim.loss_func,
+            config=config.optim,
             num_classes=self._output_shape,
         )
-        return loss_func
+        return loss_func, None
 
     def evaluate_semantic_segmentation(
         self,
@@ -240,7 +238,7 @@ class SemanticSegmentationLearner(BaseLearner):
         model=None,
         model_postprocess_fn=None,
         peft_param_names=None,
-        optimization_kwargs=None,
+        optim_kwargs=None,
         distillation_kwargs=None,
         is_train=True,
     ):
@@ -249,13 +247,13 @@ class SemanticSegmentationLearner(BaseLearner):
                 model=model,
                 model_postprocess_fn=model_postprocess_fn,
                 trainable_param_names=peft_param_names,
-                **optimization_kwargs,
+                **optim_kwargs,
             )
         else:
             return SemanticSegmentationLitModule(
                 model=self._model,
                 model_postprocess_fn=self._model_postprocess_fn,
-                **optimization_kwargs,
+                **optim_kwargs,
             )
 
     def on_predict_start(self, data: pd.DataFrame):