migrate to ray_tune (#36)

* update * fix * home * fixes * fixes * fixes * fixes * refactor * Format code with black and isort * fixes * Format code with black and isort * fixes * Format code with black and isort * update * update changelog * fix * Format code with black and isort * update * update * fixes * fixes * fixes Co-authored-by: deepsource-autofix[bot] <62050782+deepsource-autofix[bot]@users.noreply.github.com>
gradsflow · Aug 29, 2021 · 90fb2a3 · 90fb2a3
1 parent 60d245f
commit 90fb2a3
Show file tree

Hide file tree

Showing 21 changed files with 182 additions and 183 deletions.
diff --git a/.github/workflows/latest-changes.yml → .github/workflows/latest-changes.yml.off b/.github/workflows/latest-changes.yml → .github/workflows/latest-changes.yml.off
@@ -24,5 +24,5 @@ jobs:
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
           latest_changes_file: docs/CHANGELOG.md
-          latest_changes_header: '## 0.0.2\n'
+          latest_changes_header: '## 0.0.3\n'
           debug_logs: true
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 ![logo](https://ik.imagekit.io/gradsflow/logo/logo-small_g2MxLWesD.png?updatedAt=1627716948296)
 
-# An AutoML Library made with Optuna and PyTorch Lightning
+# An open-source AutoML Library in PyTorch
 
 [![Maintainability Rating](https://sonarcloud.io/api/project_badges/measure?project=gradsflow_gradsflow&metric=sqale_rating)](https://sonarcloud.io/dashboard?id=gradsflow_gradsflow)
 [![Reliability Rating](https://sonarcloud.io/api/project_badges/measure?project=gradsflow_gradsflow&metric=reliability_rating)](https://sonarcloud.io/dashboard?id=gradsflow_gradsflow)
@@ -27,10 +27,11 @@
 ## Highlights
 
 - 2020-8-25: [Released first version 0.0.1](https://pypi.org/project/gradsflow/) ✨ :tada:
+- 2020-8-29: Migrated to Ray Tune. [Read more...](https://github.com/gradsflow/gradsflow/issues/35)
 
-## What is Gradsflow?
+## What is GradsFlow?
 
-Gradsflow is based on Optuna and PyTorch Lightning ⚡️.
+GradsFlow is based on Ray and PyTorch Lightning ⚡️ (support for other torch frameworks will be added soon).
 It leverages PyTorch Lightning Flash so that you don't have to write any
 PyTorch or Optuna code for model building or hyperparameter tuning 🚀
 
@@ -40,8 +41,8 @@ You can easily integrate any PyTorch/Lightning Flash Model with Gradsflow [AutoM
 - `gradsflow.core`: [Core](https://docs.gradsflow.com/en/latest/gradsflow/core/) defines the building blocks
 of AutoML tasks.
 
-- `gradsflow.taskauto`: [AutoTasks](https://docs.gradsflow.com/en/latest/gradsflow/autotasks/) defines
-different ML/DL tasks which is provided by Gradsflow AutoML API.
+- `gradsflow.autotasks`: [AutoTasks](https://docs.gradsflow.com/en/latest/gradsflow/autotasks/) defines
+different ML/DL tasks which is provided by gradsflow AutoML API.
 
 
 📑 Check out [notebooks examples](https://github.com/gradsflow/gradsflow/tree/main/examples/nbs).
@@ -61,4 +62,4 @@ Read full [**Contributor Covenant Code of Conduct**](https://github.com/gradsflo
 
 ## Acknowledgement
 
-**Gradsflow** is built with help of Optuna and PyTorch Lightning 💜
+**GradsFlow** is built with help of Ray and PyTorch Lightning 💜
diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
@@ -1,6 +1,7 @@
 # Release Notes
 
 ## 0.0.1
+
 * 📝 update example and documentation. Done by [ aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 20 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/20).
 * :tada::sparkles: First Release - v0.0.1 - Refactor API & tested Python 3.7+. Done by [ aniketmaurya](https://github.com/aniketmaurya). Check the [Pull Request 18 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/18).
 * Adding example notebook for AutoSummarization. Done by [the GitHub user gagan3012](https://github.com/gagan3012). Check the [Pull Request 19 with the changes and stuff](https://github.com/gradsflow/gradsflow/pull/19).
@@ -12,8 +13,13 @@
 * Add AutoImageClassifier. PR [#1](https://github.com/gradsflow/gradsflow/pull/1) by [@aniketmaurya](https://github.com/aniketmaurya).
 
 ## 0.0.2
+
 * Fix max steps validation key error. PR [#31](https://github.com/gradsflow/gradsflow/pull/31) by [@aniketmaurya](https://github.com/aniketmaurya).
+
+## 0.0.3
+
 * add optuna visualizations 🎨 . PR [#27](https://github.com/gradsflow/gradsflow/pull/27) by [@aniketmaurya](https://github.com/aniketmaurya).
 * add max_steps for HPO. PR [#25](https://github.com/gradsflow/gradsflow/pull/25) by [@aniketmaurya](https://github.com/aniketmaurya).
 * :memo: update docs & license. PR [#23](https://github.com/gradsflow/gradsflow/pull/23) by [@aniketmaurya](https://github.com/aniketmaurya).
 * fetch best trial model. PR [#21](https://github.com/gradsflow/gradsflow/pull/21) by [@aniketmaurya](https://github.com/aniketmaurya).
+* migrate to ray_tune 🌟. Read more [here](https://github.com/gradsflow/gradsflow/issues/35). PR [#36](https://github.com/gradsflow/gradsflow/pull/36) by [@aniketmaurya](https://github.com/aniketmaurya).
diff --git a/docs/index.md b/docs/index.md
@@ -1,6 +1,6 @@
 ![logo](https://ik.imagekit.io/gradsflow/logo/logo-small_g2MxLWesD.png?updatedAt=1627716948296)
 
-# An AutoML Library made with Optuna and PyTorch Lightning
+# An open-source AutoML Library in PyTorch
 
 [![Maintainability Rating](https://sonarcloud.io/api/project_badges/measure?project=gradsflow_gradsflow&metric=sqale_rating)](https://sonarcloud.io/dashboard?id=gradsflow_gradsflow)
 [![Reliability Rating](https://sonarcloud.io/api/project_badges/measure?project=gradsflow_gradsflow&metric=reliability_rating)](https://sonarcloud.io/dashboard?id=gradsflow_gradsflow)
@@ -27,10 +27,11 @@
 ## Highlights
 
 - 2020-8-25: [Released first version 0.0.1](https://pypi.org/project/gradsflow/) ✨ :tada:
+- 2020-8-29: Migrated to Ray Tune. [Read more...](https://github.com/gradsflow/gradsflow/issues/35)
 
-## What is Gradsflow?
+## What is GradsFlow?
 
-Gradsflow is based on Optuna and PyTorch Lightning ⚡️.
+GradsFlow is based on Ray and PyTorch Lightning ⚡️ (support for other torch frameworks will be added soon).
 It leverages PyTorch Lightning Flash so that you don't have to write any
 PyTorch or Optuna code for model building or hyperparameter tuning 🚀
 
@@ -40,8 +41,8 @@ You can easily integrate any PyTorch/Lightning Flash Model with Gradsflow [AutoM
 - `gradsflow.core`: [Core](https://docs.gradsflow.com/en/latest/gradsflow/core/) defines the building blocks
 of AutoML tasks.
 
-- `gradsflow.taskauto`: [AutoTasks](https://docs.gradsflow.com/en/latest/gradsflow/autotasks/) defines
-different ML/DL tasks which is provided by Gradsflow AutoML API.
+- `gradsflow.autotasks`: [AutoTasks](https://docs.gradsflow.com/en/latest/gradsflow/autotasks/) defines
+different ML/DL tasks which is provided by gradsflow AutoML API.
 
 
 📑 Check out [notebooks examples](https://github.com/gradsflow/gradsflow/tree/main/examples/nbs).
@@ -61,4 +62,4 @@ Read full [**Contributor Covenant Code of Conduct**](https://github.com/gradsflo
 
 ## Acknowledgement
 
-**Gradsflow** is built with help of Optuna and PyTorch Lightning 💜
+**GradsFlow** is built with help of Ray and PyTorch Lightning 💜
diff --git a/docs/overrides/main.html b/docs/overrides/main.html
@@ -9,7 +9,7 @@
   {% endif %}
 
   <meta charset="utf-8">
-  <title>An AutoML Library made with Optuna and PyTorch Lightning</title>
+  <title>An open-source AutoML Library in PyTorch</title>
   <link rel="canonical" href="https://chitra.readthedocs.io/" />
 
   <meta property="og:type" content="website" />

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -7,3 +7,4 @@ mkdocs-autorefs==0.2.1
 mkdocstrings==0.15.2
 tags-macros-plugin @ git+https://github.com/jldiaz/mkdocs-plugin-tags.git@d26e2f124e4f3471639d426459e281080988fe7a
 gradsflow
+ray[tune]
diff --git a/examples/nbs/01-ImageClassification.ipynb b/examples/nbs/01-ImageClassification.ipynb
@@ -9,6 +9,7 @@
    "outputs": [],
    "source": [
     "import os\n",
+    "from pathlib import Path\n",
     "import sys\n",
     "import warnings\n",
     "\n",
@@ -23,18 +24,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import plotly\n",
-    "plotly.offline.init_notebook_mode(connected=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import plotly.io as pio\n",
-    "pio.renderers.default='colab'"
+    "from flash.core.data.utils import download_data\n",
+    "from flash.image import ImageClassificationData\n",
+    "from gradsflow import AutoImageClassifier"
    ]
   },
   {
@@ -43,10 +35,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from flash.core.data.utils import download_data\n",
-    "from flash.image import ImageClassificationData\n",
-    "from gradsflow import AutoImageClassifier\n",
-    "import optuna"
+    "data_dir = \"/Users/aniket/personal/gradsflow/gradsflow/data/\"\n",
+    "\n",
+    "# download_data(\"https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip\", data_dir)"
    ]
   },
   {
@@ -56,8 +47,8 @@
    "outputs": [],
    "source": [
     "datamodule = ImageClassificationData.from_folders(\n",
-    "    train_folder=\"data/hymenoptera_data/train/\",\n",
-    "    val_folder=\"data/hymenoptera_data/val/\",\n",
+    "    train_folder=f\"{data_dir}/hymenoptera_data/train/\",\n",
+    "    val_folder=f\"{data_dir}/hymenoptera_data/val/\",\n",
     ")"
    ]
   },
@@ -75,8 +66,8 @@
     "model = AutoImageClassifier(\n",
     "    datamodule,\n",
     "    suggested_conf=suggested_conf,\n",
-    "    max_epochs=10,\n",
-    "    n_trials=50,\n",
+    "    max_epochs=2,\n",
+    "    n_trials=4,\n",
     "    optimization_metric=\"val_accuracy\",\n",
     "    timeout=50,\n",
     ")\n",
@@ -90,9 +81,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "model.plot_optimization_history()"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
@@ -101,6 +90,22 @@
    "outputs": [],
    "source": []
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.analysis.best_checkpoint"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

diff --git a/examples/src/tasks/image_classifier.py b/examples/src/tasks/image_classifier.py
@@ -12,21 +12,26 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from flash.core.data.utils import download_data
+import os
+
 from flash.image import ImageClassificationData
 
 from gradsflow import AutoImageClassifier
 
 # 1. Create the DataModule
-download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", "./data")
+data_dir = os.getcwd() + "/data"
+# download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", data_dir)
 
 datamodule = ImageClassificationData.from_folders(
-    train_folder="data/hymenoptera_data/train/",
-    val_folder="data/hymenoptera_data/val/",
+    train_folder=f"{data_dir}/hymenoptera_data/train/",
+    val_folder=f"{data_dir}/hymenoptera_data/val/",
 )
 
 model = AutoImageClassifier(
-    datamodule, max_epochs=2, optimization_metric="val_accuracy"
+    datamodule,
+    max_epochs=2,
+    optimization_metric="train_accuracy",
+    max_steps=2,
 )
 print("AutoImageClassifier initialised!")
 

diff --git a/gradsflow/autotasks/autoclassification/image.py b/gradsflow/autotasks/autoclassification/image.py
@@ -63,8 +63,9 @@ class AutoImageClassifier(AutoClassifier):
 
     DEFAULT_BACKBONES = ["ssl_resnet18", "ssl_resnet50"]
 
-    def build_model(self, **kwargs) -> torch.nn.Module:
-        """Build ImageClassifier model from optuna.Trial object or via keyword arguments.
+    def build_model(self, config: dict) -> torch.nn.Module:
+        """Build ImageClassifier model from `ray.tune` hyperparameter configs
+        or via config dictionary arguments.
 
         Arguments:
             backbone [str]: Image classification backbone name - resnet18, resnet50,...
@@ -73,9 +74,9 @@ def build_model(self, **kwargs) -> torch.nn.Module:
             optimizer [str]: PyTorch Optimizers. Check `AutoImageClassification.OPTIMIZER_INDEX`
             learning_rate [float]: Learning rate for the model.
         """
-        backbone = kwargs["backbone"]
-        optimizer = kwargs["optimizer"]
-        learning_rate = kwargs["lr"]
+        backbone = config["backbone"]
+        optimizer = config["optimizer"]
+        learning_rate = config["lr"]
 
         return ImageClassifier(
             self.num_classes,

diff --git a/gradsflow/autotasks/autoclassification/text.py b/gradsflow/autotasks/autoclassification/text.py
@@ -67,8 +67,9 @@ class AutoTextClassifier(AutoClassifier):
         "sgugger/tiny-distilbert-classification",
     ]
 
-    def build_model(self, **kwargs) -> torch.nn.Module:
-        """Build ImageClassifier model from optuna.Trial object or via keyword arguments.
+    def build_model(self, config: dict) -> torch.nn.Module:
+        """Build TextClassifier model from `ray.tune` hyperparameter configs
+        or via config dictionary arguments
 
         Arguments:
             backbone [str]: Image classification backbone name - resnet18, resnet50,...
@@ -77,9 +78,9 @@ def build_model(self, **kwargs) -> torch.nn.Module:
             optimizer [str]: PyTorch Optimizers. Check `AutoImageClassification.OPTIMIZER_INDEX`
             learning_rate [float]: Learning rate for the model.
         """
-        backbone = kwargs["backbone"]
-        optimizer = kwargs["optimizer"]
-        learning_rate = kwargs["lr"]
+        backbone = config["backbone"]
+        optimizer = config["optimizer"]
+        learning_rate = config["lr"]
 
         return TextClassifier(
             self.num_classes,

diff --git a/gradsflow/autotasks/autosummarization.py b/gradsflow/autotasks/autosummarization.py
@@ -69,8 +69,9 @@ class AutoSummarization(AutoClassifier):
         "sshleifer/distilbart-xsum-12-3",
     ]
 
-    def build_model(self, **kwargs) -> torch.nn.Module:
-        """Build ImageClassifier model from optuna.Trial object or via keyword arguments.
+    def build_model(self, config: dict) -> torch.nn.Module:
+        """Build SummarizationModel from `ray.tune` hyperparameter configs
+        or via config dictionary arguments
 
         Arguments:
             backbone [str]: Image classification backbone name -
@@ -80,9 +81,9 @@ def build_model(self, **kwargs) -> torch.nn.Module:
             optimizer [str]: PyTorch Optimizers. Check `AutoImageClassification.OPTIMIZER_INDEX`
             learning_rate [float]: Learning rate for the model.
         """
-        backbone = kwargs["backbone"]
-        optimizer = kwargs["optimizer"]
-        learning_rate = kwargs["lr"]
+        backbone = config["backbone"]
+        optimizer = config["optimizer"]
+        learning_rate = config["lr"]
 
         return SummarizationTask(
             backbone=backbone,

diff --git a/gradsflow/core/autoclassifier.py b/gradsflow/core/autoclassifier.py
@@ -15,9 +15,9 @@
 from abc import abstractmethod
 from typing import Dict, List, Optional, Union
 
-import optuna
 import torch
 from flash.core.data.data_module import DataModule
+from ray import tune
 
 from gradsflow.core.automodel import AutoModel
 
@@ -39,10 +39,9 @@ def __init__(
         optimization_metric: Optional[str] = None,
         suggested_backbones: Union[List, str, None] = None,
         suggested_conf: Optional[dict] = None,
+        tune_confs: Optional[Dict] = None,
         timeout: int = 600,
         prune: bool = True,
-        trainer_confs: Optional[Dict] = None,
-        optuna_confs: Optional[Dict] = None,
     ):
         super().__init__(
             datamodule,
@@ -53,8 +52,7 @@ def __init__(
             suggested_conf=suggested_conf,
             timeout=timeout,
             prune=prune,
-            trainer_confs=trainer_confs,
-            optuna_confs=optuna_confs,
+            tune_confs=tune_confs,
         )
 
         if isinstance(suggested_backbones, (str, list, tuple)):
@@ -73,15 +71,15 @@ def forward(self, x):
         return self.model(x)
 
     # noinspection PyTypeChecker
-    def _get_trial_hparams(self, trial: optuna.Trial) -> Dict[str, str]:
-        """Fetch hyperparameters from current optuna.Trial and returns
-        key-value pair of hparams"""
+    def _create_hparam_config(self) -> Dict[str, str]:
+        """Create hyperparameter config from `ray.tune`
 
-        trial_backbone = trial.suggest_categorical("backbone", self.suggested_backbones)
-        trial_lr = trial.suggest_float("lr", *self.suggested_lr, log=True)
-        trial_optimizer = trial.suggest_categorical(
-            "optimizer", self.suggested_optimizers
-        )
+        Returns:
+             key-value pair of `ray.tune` hparams
+        """
+        trial_backbone = tune.choice(self.suggested_backbones)
+        trial_lr = tune.loguniform(*self.suggested_lr)
+        trial_optimizer = tune.choice(self.suggested_optimizers)
         hparams = {
             "backbone": trial_backbone,
             "lr": trial_lr,
@@ -90,9 +88,9 @@ def _get_trial_hparams(self, trial: optuna.Trial) -> Dict[str, str]:
         return hparams
 
     @abstractmethod
-    def build_model(self, **kwargs) -> torch.nn.Module:
+    def build_model(self, config: dict) -> torch.nn.Module:
         """Every Task implementing AutoClassifier has to implement a
-        build model method that can build `torch.nn.Module` from keyword arguments
+        build model method that can build `torch.nn.Module` from dictionary config
         and return the model.
         """
         raise NotImplementedError