diff --git a/examples/nbs/01-ImageClassification.ipynb b/examples/nbs/01-ImageClassification.ipynb index 44824eb8..b90e121e 100644 --- a/examples/nbs/01-ImageClassification.ipynb +++ b/examples/nbs/01-ImageClassification.ipynb @@ -2,7 +2,11 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, "source": [ "\"Open\n", "\n", @@ -15,6 +19,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "pycharm": { + "name": "#%%\n" + }, "tags": [] }, "outputs": [], @@ -30,7 +37,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [], "source": [ "import torchvision\n", @@ -43,7 +54,11 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, "source": [ "Let's use `CalTech101` dataset provided by `torchvision`" ] @@ -51,7 +66,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [ { "name": "stdout", @@ -75,7 +94,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [], "source": [ "train_data, val_data = random_split_dataset(data, 0.01)\n", @@ -86,7 +109,11 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, "source": [ "If you want to run Gradsflow on a remote server then first setup [ray cluster](https://docs.ray.io/en/master/cluster/index.html) and initialize ray with the remote address." ] @@ -94,7 +121,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [], "source": [ "# ray.init(address=\"REMOTE_IP_ADDR\")\n", @@ -103,7 +134,11 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, "source": [ "To train an image classifier create an object of `AutoImageClassifier` and provide number of trials and timeout." ] @@ -111,7 +146,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [ { "name": "stdout", @@ -126,10 +165,9 @@ " train_dataloader=train_dl,\n", " val_dataloader=val_dl,\n", " num_classes=num_classes,\n", - " max_epochs=5,\n", + " max_epochs=2,\n", " optimization_metric=\"train_loss\",\n", - " max_steps=1,\n", - " n_trials=1,\n", + " n_trials=4,\n", ")\n", "print(\"AutoImageClassifier initialised!\")" ] @@ -137,7 +175,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [ { "name": "stderr", @@ -758,7 +800,11 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [], "source": [ "# ray.shutdown()" diff --git a/examples/nbs/02-TextClassification.ipynb b/examples/nbs/02-TextClassification.ipynb index 3f518c9a..144d2110 100644 --- a/examples/nbs/02-TextClassification.ipynb +++ b/examples/nbs/02-TextClassification.ipynb @@ -71,7 +71,11 @@ "cell_type": "code", "execution_count": null, "id": "9f5d0474", - "metadata": {}, + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, "outputs": [], "source": [ "suggested_conf = dict(\n", @@ -83,9 +87,8 @@ " datamodule,\n", " suggested_backbones=[\"sgugger/tiny-distilbert-classification\"],\n", " suggested_conf=suggested_conf,\n", - " max_epochs=1,\n", + " max_epochs=2,\n", " optimization_metric=\"val_accuracy\",\n", - " timeout=5,\n", ")\n", "\n", "print(\"AutoTextClassifier initialised!\")\n", diff --git a/examples/nbs/03-TextSummarization.ipynb b/examples/nbs/03-TextSummarization.ipynb index ad4c7904..af626692 100644 --- a/examples/nbs/03-TextSummarization.ipynb +++ b/examples/nbs/03-TextSummarization.ipynb @@ -68,7 +68,7 @@ " suggested_conf=suggested_conf,\n", " max_epochs=1,\n", " optimization_metric=\"train_loss\",\n", - " timeout=5,\n", + " timeout=600,\n", ")\n", "\n", "print(\"AutoSummarization initialised!\")\n", diff --git a/examples/nbs/04-RayDataset.ipynb b/examples/nbs/04-RayDataset.ipynb index 1d36e6dd..4e9714b4 100644 --- a/examples/nbs/04-RayDataset.ipynb +++ b/examples/nbs/04-RayDataset.ipynb @@ -59,7 +59,7 @@ { "cell_type": "code", "execution_count": null, - "id": "1c550a6f", + "id": "b94e30dd", "metadata": { "pycharm": { "name": "#%%\n" @@ -77,21 +77,13 @@ { "cell_type": "code", "execution_count": null, - "id": "ef8c7353", + "id": "3822e6bb", "metadata": { "pycharm": { "name": "#%%\n" } }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2021-09-26 01:55:36,868\tINFO services.py:1263 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8266\u001b[39m\u001b[22m\n" - ] - } - ], + "outputs": [], "source": [ "transforms = get_augmentations()\n", "\n", diff --git a/examples/src/tasks/text_classification.py b/examples/src/tasks/text_classification.py index c974572c..e0029b97 100644 --- a/examples/src/tasks/text_classification.py +++ b/examples/src/tasks/text_classification.py @@ -1,8 +1,11 @@ +import ray from flash.core.data.utils import download_data from flash.text import TextClassificationData from gradsflow import AutoTextClassifier +ray.init(address="auto") + download_data("https://pl-flash-data.s3.amazonaws.com/imdb.zip", "./data/") print("Creating datamodule...") @@ -14,14 +17,15 @@ optimizers=["adam"], lr=(5e-4, 1e-3), ) + model = AutoTextClassifier( datamodule, suggested_backbones=["sgugger/tiny-distilbert-classification"], suggested_conf=suggested_conf, - max_epochs=1, + max_epochs=2, optimization_metric="val_accuracy", - timeout=5, ) print("AutoTextClassifier initialised!") model.hp_tune() +ray.shutdown() diff --git a/gradsflow/autotasks/autoclassification/text/text.py b/gradsflow/autotasks/autoclassification/text/text.py index 6430dc57..bc2805c8 100644 --- a/gradsflow/autotasks/autoclassification/text/text.py +++ b/gradsflow/autotasks/autoclassification/text/text.py @@ -84,6 +84,7 @@ def build_model(self, config: dict) -> torch.nn.Module: learning_rate [float]: Learning rate for the model. """ from flash.text.classification import TextClassifier + from torchmetrics import Accuracy backbone = config["backbone"] optimizer = config["optimizer"] @@ -94,4 +95,5 @@ def build_model(self, config: dict) -> torch.nn.Module: backbone=backbone, optimizer=self._OPTIMIZER_INDEX[optimizer], learning_rate=learning_rate, + metrics=Accuracy(num_classes=self.num_classes), ) diff --git a/gradsflow/autotasks/engine/backend.py b/gradsflow/autotasks/engine/backend.py index 84332287..755c88b3 100644 --- a/gradsflow/autotasks/engine/backend.py +++ b/gradsflow/autotasks/engine/backend.py @@ -14,6 +14,7 @@ import logging import math +import typing from enum import Enum from typing import Callable, Dict, Optional @@ -24,10 +25,18 @@ from gradsflow.utility.common import module_to_cls_index from gradsflow.utility.imports import is_installed -pl = None -if is_installed("pytorch_lightning"): +if typing.TYPE_CHECKING: import pytorch_lightning as pl + +if is_installed("pytorch_lightning"): + from flash import Task + from flash import Trainer as FlashTrainer + from pytorch_lightning import Trainer as PLTrainer +else: + FlashTrainer = None + PLTrainer = None + logger = logging.getLogger("core.backend") @@ -83,10 +92,12 @@ def _lightning_objective( val_check_interval = max(self.max_steps - 1, 1.0) datamodule = self.autodataset.datamodule + model = self.model_builder(config) + + trainer_cls = FlashTrainer if isinstance(model, Task) else PLTrainer - trainer = pl.Trainer( + trainer: "pl.Trainer" = trainer_cls( logger=True, - checkpoint_callback=False, gpus=math.ceil(gpu), max_epochs=self.max_epochs, max_steps=self.max_steps, @@ -95,7 +106,6 @@ def _lightning_objective( **trainer_config, ) - model = self.model_builder(config) hparams = dict(model=model.hparams) trainer.logger.log_hyperparams(hparams) trainer.fit(model, datamodule=datamodule) diff --git a/tests/autotasks/test_autotrainer.py b/tests/autotasks/test_autotrainer.py index c561340e..eb7f0b0b 100644 --- a/tests/autotasks/test_autotrainer.py +++ b/tests/autotasks/test_autotrainer.py @@ -20,15 +20,16 @@ trainer_config = {"show_progress": False} -@patch("gradsflow.autotasks.engine.backend.pl") -def test_optimization_objective(mock_pl: Mock): +@patch("gradsflow.autotasks.engine.backend.FlashTrainer") +@patch("gradsflow.autotasks.engine.backend.PLTrainer") +def test_optimization_objective(mock_pl_trainer: Mock, mock_fl_trainer: Mock): dm = MagicMock() model_builder = MagicMock() # backend_type is pl autotrainer = Backend(dm, model_builder, optimization_metric="val_accuracy", backend="pl") autotrainer.optimization_objective({}, trainer_config) - mock_pl.Trainer.assert_called() + assert mock_pl_trainer.called or mock_fl_trainer.called # wrong backend_type is passed with pytest.raises(NotImplementedError): diff --git a/tests/autotasks/test_core_automodel.py b/tests/autotasks/test_core_automodel.py index 3741fc29..9ebe231e 100644 --- a/tests/autotasks/test_core_automodel.py +++ b/tests/autotasks/test_core_automodel.py @@ -48,8 +48,9 @@ def test_create_search_space(): @patch.multiple(AutoModel, __abstractmethods__=set()) -@patch("gradsflow.autotasks.engine.backend.pl") -def test_objective(mock_pl): +@patch("gradsflow.autotasks.engine.backend.FlashTrainer") +@patch("gradsflow.autotasks.engine.backend.PLTrainer") +def test_objective(mock_pl_trainer, mock_fl_trainer): optimization_metric = "val_accuracy" model = AutoModel( datamodule, @@ -58,8 +59,8 @@ def test_objective(mock_pl): ) model.backend.model_builder = MagicMock() - trainer = mock_pl.Trainer = MagicMock() - trainer.callback_metrics = {optimization_metric: torch.as_tensor([1])} + + mock_pl_trainer.callback_metrics = mock_fl_trainer.callback_metrics = {optimization_metric: torch.as_tensor([1])} model.backend.optimization_objective({}, {})