@@ -24,6 +24,7 @@ from super_gradients.training.datasets.segmentation_datasets import CityscapesDa
 from super_gradients.common.factories.samplers_factory import SamplersFactory
 from super_gradients.training.utils.distributed_training_utils import wait_for_the_master, get_local_rank
 from super_gradients.common.abstractions.abstract_logger import get_logger
+from super_gradients.training.utils.utils import override_default_params_without_nones
 
 logger = get_logger(__name__)
 
@@ -83,7 +84,7 @@ def _process_dataloader_params(cfg, dataloader_params, dataset, train):
         default_dataloader_params["sampler"] = {"DistributedSampler": {}}
         default_dataloader_params = _instantiate_sampler(dataset, default_dataloader_params)
 
-    dataloader_params = _override_default_params_without_nones(dataloader_params, default_dataloader_params)
+    dataloader_params = override_default_params_without_nones(dataloader_params, default_dataloader_params)
     if get_param(dataloader_params, "batch_sampler"):
         sampler = dataloader_params.pop("sampler")
         batch_size = dataloader_params.pop("batch_size")
@@ -92,13 +93,6 @@ def _process_dataloader_params(cfg, dataloader_params, dataset, train):
     return dataloader_params
 
 
-def _override_default_params_without_nones(params, default_params):
-    for key, val in default_params.items():
-        if key not in params.keys() or params[key] is None:
-            params[key] = val
-    return params
-
-
 def _instantiate_sampler(dataset, dataloader_params):
     sampler_name = list(dataloader_params["sampler"].keys())[0]
     dataloader_params["sampler"][sampler_name]["dataset"] = dataset

from .training_hyperparams import cifar10_resnet_train_params, cityscapes_ddrnet_train_params, \
    cityscapes_regseg48_train_params, \
    cityscapes_stdc_base_train_params, \
    cityscapes_stdc_seg50_train_params, \
    cityscapes_stdc_seg75_train_params, \
    coco2017_ssd_lite_mobilenet_v2_train_params, \
    coco2017_yolox_train_params, \
    coco_segmentation_shelfnet_lw_train_params, \
    imagenet_efficientnet_train_params, \
    imagenet_mobilenetv2_train_params, \
    imagenet_mobilenetv3_base_train_params, \
    imagenet_mobilenetv3_large_train_params, \
    imagenet_mobilenetv3_small_train_params, \
    imagenet_regnetY_train_params, \
    imagenet_repvgg_train_params, \
    imagenet_resnet50_train_params, \
    imagenet_resnet50_kd_train_params, \
    imagenet_vit_base_train_params, \
    imagenet_vit_large_train_params, \
    get

__all__ = ["cifar10_resnet_train_params", "cityscapes_ddrnet_train_params",
           "cityscapes_regseg48_train_params",
           "cityscapes_stdc_base_train_params",
           "cityscapes_stdc_seg50_train_params",
           "cityscapes_stdc_seg75_train_params",
           "coco2017_ssd_lite_mobilenet_v2_train_params",
           "coco2017_yolox_train_params",
           "coco_segmentation_shelfnet_lw_train_params",
           "imagenet_efficientnet_train_params",
           "imagenet_mobilenetv2_train_params",
           "imagenet_mobilenetv3_base_train_params",
           "imagenet_mobilenetv3_large_train_params",
           "imagenet_mobilenetv3_small_train_params",
           "imagenet_regnetY_train_params",
           "imagenet_repvgg_train_params",
           "imagenet_resnet50_train_params",
           "imagenet_resnet50_kd_train_params",
           "imagenet_vit_base_train_params",
           "imagenet_vit_large_train_params",
           "get"]

import hydra
import pkg_resources
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from super_gradients.training.utils.utils import override_default_params_without_nones
from super_gradients.common.abstractions.abstract_logger import get_logger
from typing import Dict

logger = get_logger(__name__)


def get(config_name, overriding_params: Dict = None) -> Dict:
    """
    Class for creating training hyper parameters dictionary, taking defaults from yaml
     files in src/super_gradients/recipes.

    :param overriding_params: Dict, dictionary like object containing entries to override in the recipe's training
     hyper parameters dictionary.
    :param config_name: yaml config filename in recipes (for example coco2017_yolox).
    """
    if overriding_params is None:
        overriding_params = dict()
    GlobalHydra.instance().clear()
    with initialize_config_dir(config_dir=pkg_resources.resource_filename("super_gradients.recipes", "")):
        cfg = compose(config_name=config_name)
        cfg = hydra.utils.instantiate(cfg)
        training_params = cfg.training_hyperparams
        training_params = override_default_params_without_nones(overriding_params, training_params)
        return training_params


def cifar10_resnet_train_params(overriding_params: Dict = None):
    return get("cifar10_resnet", overriding_params)


def cityscapes_ddrnet_train_params(overriding_params: Dict = None):
    return get("cityscapes_ddrnet", overriding_params)


def cityscapes_regseg48_train_params(overriding_params: Dict = None):
    return get("cityscapes_regseg48", overriding_params)


def cityscapes_stdc_base_train_params(overriding_params: Dict = None):
    return get("cityscapes_stdc_base", overriding_params)


def cityscapes_stdc_seg50_train_params(overriding_params: Dict = None):
    return get("cityscapes_stdc_seg50", overriding_params)


def cityscapes_stdc_seg75_train_params(overriding_params: Dict = None):
    return get("cityscapes_stdc_seg75", overriding_params)


def coco2017_ssd_lite_mobilenet_v2_train_params(overriding_params: Dict = None):
    return get("coco2017_ssd_lite_mobilenet_v2", overriding_params)


def coco2017_yolox_train_params(overriding_params: Dict = None):
    return get("coco2017_yolox", overriding_params)


def coco_segmentation_shelfnet_lw_train_params(overriding_params: Dict = None):
    return get("coco_segmentation_shelfnet_lw", overriding_params)


def imagenet_efficientnet_train_params(overriding_params: Dict = None):
    return get("imagenet_efficientnet", overriding_params)


def imagenet_mobilenetv2_train_params(overriding_params: Dict = None):
    return get("imagenet_mobilenetv2", overriding_params)


def imagenet_mobilenetv3_base_train_params(overriding_params: Dict = None):
    return get("imagenet_mobilenetv3_base", overriding_params)


def imagenet_mobilenetv3_large_train_params(overriding_params: Dict = None):
    return get("imagenet_mobilenetv3_large", overriding_params)


def imagenet_mobilenetv3_small_train_params(overriding_params: Dict = None):
    return get("imagenet_mobilenetv3_small", overriding_params)


def imagenet_regnetY_train_params(overriding_params: Dict = None):
    return get("imagenet_regnetY", overriding_params)


def imagenet_repvgg_train_params(overriding_params: Dict = None):
    return get("imagenet_repvgg", overriding_params)


def imagenet_resnet50_train_params(overriding_params: Dict = None):
    return get("imagenet_resnet50", overriding_params)


def imagenet_resnet50_kd_train_params(overriding_params: Dict = None):
    return get("imagenet_resnet50_kd", overriding_params)


def imagenet_vit_base_train_params(overriding_params: Dict = None):
    return get("imagenet_vit_base", overriding_params)


def imagenet_vit_large_train_params(overriding_params: Dict = None):
    return get("imagenet_vit_large", overriding_params)

@@ -2,7 +2,7 @@ import math
 import time
 from functools import lru_cache
 from pathlib import Path
-from typing import Mapping, Optional, Tuple, Union, List
+from typing import Mapping, Optional, Tuple, Union, List, Dict
 from zipfile import ZipFile
 import os
 from jsonschema import validate
@@ -444,3 +444,16 @@ def get_image_size_from_path(img_path: str) -> Tuple[int, int]:
     """Get the image size of an image at a specific path"""
     with open(img_path, 'rb') as f:
         return exif_size(Image.open(f))
+
+
+def override_default_params_without_nones(params: Dict, default_params: Dict) -> Dict:
+    """
+    Helper method for overriding default dictionary's entries excluding entries with None values.
+    :param params: dict, output dictionary which will take the defaults.
+    :param default_params: dict, dictionary for the defaults.
+    :return: dict, params after manipulation,
+    """
+    for key, val in default_params.items():
+        if key not in params.keys() or params[key] is None:
+            params[key] = val
+    return params

@@ -4,7 +4,7 @@ import unittest
 from tests.integration_tests.ema_train_integration_test import EMAIntegrationTest
 from tests.unit_tests import ZeroWdForBnBiasTest, SaveCkptListUnitTest, TestAverageMeter, \
     TestRepVgg, TestWithoutTrainTest, OhemLossTest, EarlyStopTest, SegmentationTransformsTest, \
-    TestConvBnRelu, FactoriesTest, InitializeWithDataloadersTest
+    TestConvBnRelu, FactoriesTest, InitializeWithDataloadersTest, TrainingParamsTest
 from tests.end_to_end_tests import TestTrainer
 from tests.unit_tests.phase_delegates_test import ContextMethodsTest
 from tests.unit_tests.random_erase_test import RandomEraseTest
@@ -76,6 +76,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetSubclassing))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TestDetectionDatasetCaching))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(MultiScaleTest))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(TrainingParamsTest))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

@@ -16,10 +16,10 @@ from tests.unit_tests.segmentation_transforms_test import SegmentationTransforms
 from tests.unit_tests.pretrained_models_unit_test import PretrainedModelsUnitTest
 from tests.unit_tests.conv_bn_relu_test import TestConvBnRelu
 from tests.unit_tests.initialize_with_dataloaders_test import InitializeWithDataloadersTest
-
+from tests.unit_tests.training_params_factory_test import TrainingParamsTest
 
 __all__ = ['TestDatasetInterface', 'ZeroWdForBnBiasTest', 'SaveCkptListUnitTest',
            'AllArchitecturesTest', 'TestAverageMeter', 'TestRepVgg', 'TestWithoutTrainTest',
            'StrictLoadEnumTest', 'TrainWithInitializedObjectsTest', 'TestAutoAugment',
            'OhemLossTest', 'EarlyStopTest', 'SegmentationTransformsTest', 'PretrainedModelsUnitTest', 'TestConvBnRelu',
-           'FactoriesTest', 'InitializeWithDataloadersTest']
+           'FactoriesTest', 'InitializeWithDataloadersTest', 'TrainingParamsTest']

import unittest
from super_gradients.training import training_hyperparams


class TrainingParamsTest(unittest.TestCase):
    def test_get_train_params(self):
        train_params = training_hyperparams.coco2017_yolox_train_params()
        self.assertTrue(train_params["loss"] == "yolox_loss")
        self.assertTrue(train_params["max_epochs"] == 300)

    def test_get_train_params_with_overrides(self):
        train_params = training_hyperparams.coco2017_yolox_train_params(overriding_params={"max_epochs": 5})
        self.assertTrue(train_params["loss"] == "yolox_loss")
        self.assertTrue(train_params["max_epochs"] == 5)


if __name__ == '__main__':
    unittest.main()