🚀 feature/fix train arguments docs (#44)

* update * add timeout stopper * remove suggested conf * update docs * fix * add test * fix
gradsflow · Aug 29, 2021 · db3e871 · db3e871
1 parent dbb46cb
commit db3e871
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 20 deletions.
diff --git a/gradsflow/autotasks/autoclassification/image.py b/gradsflow/autotasks/autoclassification/image.py
@@ -48,12 +48,7 @@ class AutoImageClassifier(AutoClassifier):
                 val_folder="data/hymenoptera_data/val/",
             )
 
-            suggested_conf = dict(
-                optimizers=["adam", "sgd"],
-                lr=(5e-4, 1e-3),
-            )
             model = AutoImageClassifier(datamodule,
-                                        suggested_conf=suggested_conf,
                                         max_epochs=10,
                                         optimization_metric="val_accuracy",
                                         timeout=300)

diff --git a/gradsflow/autotasks/autoclassification/text.py b/gradsflow/autotasks/autoclassification/text.py
@@ -48,13 +48,8 @@ class AutoTextClassifier(AutoClassifier):
                 val_file="data/imdb/valid.csv",
             )
 
-            suggested_conf = dict(
-                optimizers=["adam", "sgd"],
-                lr=(5e-4, 1e-3),
-            )
             model = AutoTextClassifier(datamodule,
                                        suggested_backbones=['sgugger/tiny-distilbert-classification'],
-                                       suggested_conf=suggested_conf,
                                        max_epochs=10,
                                        optimization_metric="val_accuracy",
                                        timeout=300)

diff --git a/gradsflow/autotasks/autosummarization.py b/gradsflow/autotasks/autosummarization.py
@@ -51,12 +51,7 @@ class AutoSummarization(AutoClassifier):
                 test_file="data/xsum/test.csv",
             )
 
-            suggested_conf = dict(
-                optimizers=["adam", "sgd"],
-                lr=(5e-4, 1e-3),
-            )
             model = AutoSummarization(datamodule,
-                                    suggested_conf=suggested_conf,
                                     max_epochs=10,
                                     optimization_metric="val_accuracy",
                                     timeout=300)

diff --git a/gradsflow/core/automodel.py b/gradsflow/core/automodel.py
@@ -48,16 +48,14 @@ class AutoModel:
     OPTIMIZER_INDEX = module_to_cls_index(torch.optim, True)
     DEFAULT_OPTIMIZERS = ["adam", "sgd"]
     DEFAULT_LR = (1e-5, 1e-2)
-    _BEST_MODEL = "best_model"
-    _CURRENT_MODEL = "current_model"
 
     def __init__(
         self,
         datamodule: DataModule,
         max_epochs: int = 10,
         max_steps: Optional[int] = None,
         optimization_metric: Optional[str] = None,
-        n_trials: int = 100,
+        n_trials: int = 20,
         suggested_conf: Optional[dict] = None,
         timeout: int = 600,
         prune: bool = True,
@@ -136,23 +134,59 @@ def objective(self, config: Dict, trainer_config: Dict):
         return trainer.callback_metrics[self.optimization_metric].item()
 
     def hp_tune(
-        self, ray_config: Optional[dict] = None, trainer_config: Optional[dict] = None
+        self,
+        name: Optional[str] = None,
+        ray_config: Optional[dict] = None,
+        trainer_config: Optional[dict] = None,
+        mode: Optional[str] = None,
+        gpu: Optional[float] = 0,
+        cpu: Optional[float] = None,
+        resume: bool = False,
     ):
         """
         Search Hyperparameter and builds model with the best params
+
+        Args:
+            name Optional[str]: name of the experiment.
+            ray_config dict: configuration passed to `ray.tune.run(...)`
+            trainer_config dict: configuration passed to `pl.trainer.fit(...)`
+            mode Optional[str]: Whether to maximize or mimimize the `optimization_metric`.
+            Values are `max` or `min`
+            gpu Optional[float]: Amount of GPU resource per trial.
+            cpu float: CPU cores per trial
+            resume bool: Whether to resume the training or not.
+
+        !!! note
+            ```python
+                automodel = AutoClassifier(data)  # implements `AutoModel`
+                automodel.hp_tune(name="gflow-example", gpu=1)
+            ```
         """
+
         trainer_config = trainer_config or {}
         ray_config = ray_config or {}
 
         search_space = self._create_search_space()
         trainable = self.objective
 
+        resources_per_trial = {}
+        if gpu:
+            resources_per_trial["gpu"] = gpu
+        if cpu:
+            resources_per_trial["cpu"] = cpu
+
+        mode = mode or "max"
+        timeout_stopper = tune.stopper.TimeoutStopper(self.timeout)
         analysis = tune.run(
             tune.with_parameters(trainable, trainer_config=trainer_config),
+            name=name,
             num_samples=self.n_trials,
             metric=self.optimization_metric,
-            mode="max",
+            mode=mode,
             config=search_space,
+            resources_per_trial=resources_per_trial,
+            resume=resume,
+            stop=timeout_stopper,
             **ray_config,
         )
         self.analysis = analysis

diff --git a/tests/core/test_automodel.py b/tests/core/test_automodel.py
@@ -54,3 +54,12 @@ def test_objective(mock_pl):
     trainer.callback_metrics = {optimization_metric: torch.as_tensor([1])}
 
     model.objective({}, {})
+
+
+@patch("gradsflow.core.automodel.tune.run")
+def test_hp_tune(mock_tune):
+    automodel = AutoModel(datamodule)
+    automodel._create_search_space = MagicMock()
+    automodel.hp_tune(gpu=1, cpu=2)
+
+    mock_tune.assert_called()