Refactor profile (#251)

abhash-er · web-flow · commit 1bc7aa966ecc · 2025-07-02T14:28:31.000+02:00
* refactor(Profiles): change searchspace_type to trainer_preset

* fix(scripts): fix the profiles of all scripts following refactor

refactor changes searchspace_type to trainer_preset
diff --git a/README.md b/README.md
@@ -49,10 +49,10 @@ We define modular, differentiable NAS components within our library. Below is a
 ```python 
 from confopt.profile import DARTSProfile
 from confopt.train import Experiment
-from confopt.enums import SearchSpaceType, DatasetType
+from confopt.enums import TrainerPresetType, SearchSpaceType, DatasetType
 
 profile = DARTSProfile(
-        searchspace_type=SearchSpaceType.DARTS,
+        trainer_preset=TrainerPresetType.DARTS,
         epochs=3
 )
 
diff --git a/docs/source/README.md b/docs/source/README.md
@@ -53,10 +53,10 @@ Below is a snippet that demonstrates how we run a vanilla-DARTS experiment.
 ```python 
 from confopt.profile import DARTSProfile
 from confopt.train import Experiment
-from confopt.enums import SearchSpaceType, DatasetType
+from confopt.enums import TrainerPresetType, SearchSpaceType, DatasetType
 
 profile = DARTSProfile(
-        searchspace_type=SearchSpaceType.DARTS,
+        trainer_preset=TrainerPresetType.DARTS,
         epochs=3
 )
 
diff --git a/examples/demo.py b/examples/demo.py
@@ -2,7 +2,7 @@
 
 from confopt.profile import DARTSProfile, DRNASProfile, GDASProfile, ReinMaxProfile
 from confopt.train import Experiment
-from confopt.enums import DatasetType, SearchSpaceType
+from confopt.enums import DatasetType, TrainerPresetType, SearchSpaceType
 
 if __name__ == "__main__":
     searchspace = SearchSpaceType.DARTS
@@ -19,7 +19,7 @@
     }
 
     profile = DRNASProfile(
-        searchspace_type=searchspace,
+        trainer_preset=TrainerPresetType.DARTS,
         epochs=10,
         oles=True,
         calc_gm_score=True,
diff --git a/examples/demo_advanced.py b/examples/demo_advanced.py
@@ -2,13 +2,13 @@
 
 from confopt.profile import GDASProfile
 from confopt.train import Experiment
-from confopt.enums import DatasetType, SearchSpaceType
+from confopt.enums import DatasetType, TrainerPresetType, SearchSpaceType
 
 if __name__ == "__main__":
     search_space = SearchSpaceType.DARTS
 
     profile = GDASProfile(
-        searchspace_type=search_space,
+        trainer_preset=TrainerPresetType.DARTS,
         epochs=10,
         perturbation="random",
         entangle_op_weights=True,
diff --git a/examples/demo_light.py b/examples/demo_light.py
@@ -2,11 +2,11 @@
 
 from confopt.profile import GDASProfile
 from confopt.train import Experiment
-from confopt.enums import SearchSpaceType, DatasetType
+from confopt.enums import SearchSpaceType, TrainerPresetType, DatasetType
 
 if __name__ == "__main__":
     profile = GDASProfile(
-        searchspace_type=SearchSpaceType.DARTS,
+        trainer_preset=TrainerPresetType.DARTS,
         epochs=3,
     )
 
diff --git a/examples/demo_taskonomy.py b/examples/demo_taskonomy.py
@@ -8,7 +8,7 @@
 if __name__ == "__main__":
     domain = "class_object"
     profile = GDASProfile(
-        searchspace_type="tnb101", epochs=3, searchspace_domain=domain
+        trainer_preset="tnb101", epochs=3, searchspace_domain=domain
     )
     profile.configure_searchspace(num_classes=get_num_classes("taskonomy", domain))
     experiment = Experiment(
diff --git a/examples/example_synthetic.py b/examples/example_synthetic.py
@@ -78,7 +78,7 @@ def get_profile(args: argparse.Namespace) -> Callable:  # type: ignore
     searchspace = SearchSpaceType.BABYDARTS
     dataset = DatasetType.SYNTHETIC
 
-    profile = get_profile(args)(searchspace_type=searchspace, epochs=args.search_epochs)
+    profile = get_profile(args)(trainer_preset=searchspace.value, epochs=args.search_epochs)
 
     profile.configure_synthetic_dataset(
         signal_width=args.signal_width,
diff --git a/examples/examples_lora.py b/examples/examples_lora.py
@@ -12,7 +12,7 @@
     seed = 100
 
     profile = DARTSProfile(
-        searchspace_type=searchspace,
+        trainer_preset="darts",
         is_partial_connection=True,
         perturbation="random",
         sampler_sample_frequency="step",
diff --git a/examples/experiment_drnas.py b/examples/experiment_drnas.py
@@ -7,7 +7,7 @@
 
 from confopt.profile import DiscreteProfile, DRNASProfile
 from confopt.train import Experiment
-from confopt.enums import DatasetType, SearchSpaceType
+from confopt.enums import DatasetType, SearchSpaceType, TrainerPresetType
 
 dataset_size = {
     "cifar10": 10,
@@ -65,7 +65,7 @@ def read_args() -> argparse.Namespace:
 
     # Sampler and Perturbator have different sample_frequency
     profile = DRNASProfile(
-        searchspace_type=searchspace,
+        trainer_preset=TrainerPresetType(args.searchspace),
         is_partial_connection=args.searchspace == "darts",
         epochs=args.search_epochs,
         sampler_sample_frequency="step",
@@ -95,7 +95,7 @@ def read_args() -> argparse.Namespace:
     }
     profile.configure_trainer(**train_config)
     discrete_profile = DiscreteProfile(
-        searchspace_type=searchspace, epochs=args.eval_epochs, train_portion=0.9
+        trainer_preset=args.searchspace, epochs=args.eval_epochs, train_portion=0.9
     )
     discrete_profile.configure_trainer(batch_size=64)
 
diff --git a/examples/experiment_example.py b/examples/experiment_example.py
@@ -2,7 +2,7 @@
 
 from confopt.profile import GDASProfile
 from confopt.train import Experiment
-from confopt.enums import DatasetType, SearchSpaceType
+from confopt.enums import DatasetType, SearchSpaceType, TrainerPresetType
 
 if __name__ == "__main__":
     searchspace = SearchSpaceType("nb201")
@@ -11,7 +11,7 @@
 
     # Sampler and Perturbator have different sample_frequency
     profile = GDASProfile(
-        searchspace_type=searchspace,
+        trainer_preset=TrainerPresetType("nb201"),
         is_partial_connection=True,
         perturbation="random",
         sampler_sample_frequency="step",
diff --git a/examples/notebooks/demo_notebook.ipynb b/examples/notebooks/demo_notebook.ipynb
@@ -25,7 +25,7 @@
    "source": [
     "from confopt.profile import DARTSProfile\n",
     "from confopt.train import Experiment\n",
-    "from confopt.enums import SearchSpaceType, DatasetType"
+    "from confopt.enums import SearchSpaceType, DatasetType, TrainerPresetType"
    ]
   },
   {
@@ -51,7 +51,7 @@
    "outputs": [],
    "source": [
     "profile = DARTSProfile(\n",
-    "    searchspace_type=SearchSpaceType.DARTS,\n",
+    "    trainer_preset=TrainerPresetType.DARTS,\n",
     "    epochs=3,\n",
     "    seed=100,\n",
     "    is_partial_connection=True,\n",
@@ -115,7 +115,7 @@
     "from confopt.profile import GDASProfile\n",
     "\n",
     "profile = GDASProfile(\n",
-    "    searchspace_type=SearchSpaceType.NB201,\n",
+    "    trainer_preset=TrainerPresetType.NB201,\n",
     "    epochs=10,\n",
     "    perturbation=\"random\",\n",
     "    lora_rank=1,\n",
@@ -323,7 +323,7 @@
     "from confopt.profile import DiscreteProfile\n",
     "\n",
     "discrete_profile = DiscreteProfile(\n",
-    "    searchspace_type=SearchSpaceType.DARTS,\n",
+    "    trainer_preset=TrainerPresetType.DARTS,\n",
     "    epochs=10,\n",
     "    seed=100,\n",
     "    batch_size=96,\n",
diff --git a/scripts/benchsuite_experiments/benchsuite.py b/scripts/benchsuite_experiments/benchsuite.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from confopt.enums import DatasetType, SearchSpaceType
+from confopt.enums import DatasetType, SearchSpaceType, TrainerPresetType
 from confopt.profile import BaseProfile, DARTSProfile, DiscreteProfile
 from confopt.searchspace.darts.core.genotypes import PRIMITIVES
 from confopt.train import Experiment
@@ -180,7 +180,7 @@ def configure_discrete_profile_with_hp_set(
 
 if __name__ == "__main__":
     profile = DARTSProfile(
-        searchspace_type=SearchSpaceType.DARTS,
+        trainer_preset=TrainerPresetType.DARTS,
         epochs=10,
     )
 
diff --git a/scripts/benchsuite_experiments/supernet_search.py b/scripts/benchsuite_experiments/supernet_search.py
@@ -207,7 +207,7 @@ def read_config(file_path: str) -> dict[str, str]:
 
     epochs = 3 if args.dryrun is True else epochs_profiles[args.optimizer]
     profile = profile_classes[args.optimizer](
-        searchspace_type=SEARCHSPACE,
+        trainer_preset=SEARCHSPACE.value,
         epochs=epochs,
         is_partial_connection=args.pcdarts,
         perturbation=args.sdarts,
diff --git a/scripts/benchsuite_experiments/train_genotype.py b/scripts/benchsuite_experiments/train_genotype.py
@@ -108,7 +108,7 @@ def set_profile_genotype(discrete_profile: DiscreteProfile, path: str) -> None:
 def main(args: argparse.Namespace, hpset: int) -> None:
 
     discrete_profile = DiscreteProfile(
-        searchspace_type="darts",
+        trainer_preset="darts",
         epochs=args.epochs,
         seed=args.seed,
     )
diff --git a/scripts/train_genotype_ddp.py b/scripts/train_genotype_ddp.py
@@ -1,9 +1,10 @@
+from confopt.enums import TrainerPresetType
 from confopt.profile.profiles import DiscreteProfile
 from confopt.train import Experiment
 from confopt.train.experiment import DatasetType, SearchSpaceType
 
 if __name__ == "__main__":
-    profile = DiscreteProfile(searchspace_type=SearchSpaceType.DARTS)
+    profile = DiscreteProfile(trainer_preset=TrainerPresetType.DARTS)
     config = profile.get_trainer_config()
     profile.configure_trainer(use_ddp=True)
     config.update({"genotype": profile.get_genotype()})
diff --git a/src/confopt/enums.py b/src/confopt/enums.py
@@ -15,6 +15,18 @@ def __str__(self) -> str:
         return self.value
 
 
+class TrainerPresetType(Enum):
+    DARTS = "darts"
+    NB201 = "nb201"
+    NB1SHOT1 = "nb1shot1"
+    TNB101 = "tnb101"
+    BABYDARTS = "baby_darts"
+    RobustDARTS = "robust_darts"
+
+    def __str__(self) -> str:
+        return self.value
+
+
 class SamplerType(Enum):
     DARTS = "darts"
     DRNAS = "drnas"
diff --git a/src/confopt/profile/base.py b/src/confopt/profile/base.py
@@ -5,7 +5,7 @@
 
 import torch
 
-from confopt.enums import SamplerType, SearchSpaceType
+from confopt.enums import SamplerType, TrainerPresetType
 
 DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 # TODO Change this to real data
@@ -113,7 +113,7 @@ class BaseProfile:
     def __init__(  # noqa: C901, PLR0912, PLR0915
         self,
         sampler_type: str | SamplerType,
-        searchspace_type: str | SearchSpaceType,
+        trainer_preset: str | TrainerPresetType,
         epochs: int = 50,
         *,
         sampler_sample_frequency: str = "step",
@@ -154,23 +154,23 @@ def __init__(  # noqa: C901, PLR0912, PLR0915
             AssertionError: If any of the provided configurations are invalid or \
                 inconsistent.
         """
-        self.searchspace_type = (
-            SearchSpaceType(searchspace_type)
-            if isinstance(searchspace_type, str)
-            else searchspace_type
+        self.trainer_preset = (
+            TrainerPresetType(trainer_preset)
+            if isinstance(trainer_preset, str)
+            else trainer_preset
         )
         self.sampler_type = (
             SamplerType(sampler_type) if isinstance(sampler_type, str) else sampler_type
         )
 
         assert isinstance(
-            self.searchspace_type, SearchSpaceType
-        ), f"Illegal value {self.searchspace_type} for searchspace_type"
+            self.trainer_preset, TrainerPresetType
+        ), f"Illegal value {self.trainer_preset} for trainer_preset"
         assert isinstance(
             self.sampler_type, SamplerType
         ), f"Illegal value {self.sampler_type} for sampler_type"
 
-        if self.searchspace_type == SearchSpaceType.TNB101:
+        if self.trainer_preset == TrainerPresetType.TNB101:
             assert searchspace_domain in [
                 "class_object",
                 "class_scene",
@@ -179,7 +179,7 @@ def __init__(  # noqa: C901, PLR0912, PLR0915
             assert (
                 searchspace_domain is None
             ), "searchspace_domain is not required for this searchspace"
-        if searchspace_type in ("nb1shot1", SearchSpaceType.NB1SHOT1):
+        if trainer_preset in ("nb1shot1", TrainerPresetType.NB1SHOT1):
             assert searchspace_subspace in [
                 "S1",
                 "S2",
@@ -486,7 +486,7 @@ def get_config(self) -> dict:
                 "toggle_probability": self.lora_toggle_probability,
             },
             "sampler_type": self.sampler_type,
-            "searchspace": self.searchspace_type.value,
+            "searchspace": self.trainer_preset.value,
             "searchspace_domain": self.searchspace_domain,
             "weight_type": weight_type,
             "oles": self.oles_config,
@@ -578,24 +578,25 @@ def _initialize_partial_connector_config(self) -> None:
 
     def _initialize_trainer_config(self) -> None:
         """Initialize the configuration for the trainer based on the
-        searchspace_type.
+        trainer_preset.
 
         Args:
             None
 
         Returns:
             None
         """
-        if self.searchspace_type == SearchSpaceType.NB201:
+        if self.trainer_preset == TrainerPresetType.NB201:
             self._initialize_trainer_config_nb201()
-        elif self.searchspace_type in (
-            SearchSpaceType.BABYDARTS,
-            SearchSpaceType.DARTS,
+        elif self.trainer_preset in (
+            TrainerPresetType.BABYDARTS,
+            TrainerPresetType.DARTS,
+            TrainerPresetType.RobustDARTS,
         ):
             self._initialize_trainer_config_darts()
-        elif self.searchspace_type == SearchSpaceType.NB1SHOT1:
+        elif self.trainer_preset == TrainerPresetType.NB1SHOT1:
             self._initialize_trainer_config_1shot1()
-        elif self.searchspace_type == SearchSpaceType.TNB101:
+        elif self.trainer_preset == TrainerPresetType.TNB101:
             self._initialize_trainer_config_tnb101()
 
     def _initialize_dropout_config(self) -> None:
@@ -946,7 +947,7 @@ def configure_searchspace(self, **config: Any) -> None:
 
         Args:
             **config: Arbitrary keyword arguments. Possible depend on the \
-                the search space type. For more information please check the \
+                the trainer preset type. For more information please check the \
                 Parameters of the supernet of each search space.
 
 
@@ -1058,7 +1059,7 @@ def get_run_description(self) -> str:
             str: A string describing the run configuration.
         """
         run_configs = []
-        run_configs.append(f"ss_{self.searchspace_type}")
+        run_configs.append(f"ss_{self.trainer_preset}")
         if self.entangle_op_weights:
             run_configs.append("type_we")
         else:
diff --git a/src/confopt/profile/profiles.py b/src/confopt/profile/profiles.py
diff --git a/src/confopt/train/experiment.py b/src/confopt/train/experiment.py
diff --git a/tests/test_profiles.py b/tests/test_profiles.py

Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@`
`8`	`8`	`if __name__ == "__main__":`
`9`	`9`	`domain = "class_object"`
`10`	`10`	`profile = GDASProfile(`
`11`		`- searchspace_type="tnb101", epochs=3, searchspace_domain=domain`
	`11`	`+ trainer_preset="tnb101", epochs=3, searchspace_domain=domain`
`12`	`12`	`)`
`13`	`13`	`profile.configure_searchspace(num_classes=get_num_classes("taskonomy", domain))`
`14`	`14`	`experiment = Experiment(`