pickle the AutoML object (#37)

* pickle the AutoML object * get best model per estimator * test deberta * stateless API * Add Gitter badge (#41) * prevent divide by zero * test roberta * BlendSearchTuner Co-authored-by: Chi Wang (MSR) <chiw@microsoft.com> Co-authored-by: The Gitter Badger <badger@gitter.im>
2021-03-16 22:13:35 -07:00 · 2021-03-16 22:13:35 -07:00 · 4a8110c87b
parent ec37ae8f8f
commit 4a8110c87b
19 changed files with 685 additions and 42 deletions
--- a/.gitignore
+++ b/.gitignore
@ -151,4 +151,5 @@ catboost_info
 notebook/*.pkl
 notebook/.azureml
 mlruns
-logs
+logs
+automl.pkl
--- a/README.md
+++ b/README.md
@ -1,7 +1,8 @@
 [![PyPI version](https://badge.fury.io/py/FLAML.svg)](https://badge.fury.io/py/FLAML)
 [![Build](https://github.com/microsoft/FLAML/actions/workflows/python-package.yml/badge.svg)](https://github.com/microsoft/FLAML/actions/workflows/python-package.yml)
 ![Python Version](https://img.shields.io/badge/3.6%20%7C%203.7%20%7C%203.8-blue)
-[![Downloads](https://pepy.tech/badge/flaml/month)](https://pepy.tech/project/flaml) [![Join the chat at https://gitter.im/FLAMLer/community](https://badges.gitter.im/FLAMLer/community.svg)](https://gitter.im/FLAMLer/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge)
+[![Downloads](https://pepy.tech/badge/flaml/month)](https://pepy.tech/project/flaml)
+[![Join the chat at https://gitter.im/FLAMLer/community](https://badges.gitter.im/FLAMLer/community.svg)](https://gitter.im/FLAMLer/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge)

 # FLAML - Fast and Lightweight AutoML

@ -12,11 +13,14 @@

 FLAML is a lightweight Python library that finds accurate machine
 learning models automatically, efficiently and economically. It frees users from selecting
-learners and hyperparameters for each learner. It is fast and cheap.
+learners and hyperparameters for each learner. It is fast and economical. 
 The simple and lightweight design makes it easy to extend, such as
 adding customized learners or metrics. FLAML is powered by a new, [cost-effective
 hyperparameter optimization](https://github.com/microsoft/FLAML/tree/main/flaml/tune)
 and learner selection method invented by Microsoft Research.
+FLAML leverages the structure of the search space to choose a search order optimized for both cost and error. For example, the system tends to propose cheap configurations at the beginning stage of the search,
+but quickly moves to configurations with high model complexity and large sample size when needed in the later stage of the search. For another example, it favors cheap learners in the beginning but penalizes them later if the error improvement is slow. The cost-bounded search and cost-based prioritization make a big difference in the the search efficiency under budget constraints.
+
 FLAML is easy to use:

 * With three lines of code, you can start using this economical and fast
@ -117,7 +121,7 @@ And they can be used in distributed HPO frameworks such as ray tune or nni.

 For more technical details, please check our papers.

-* [FLAML: A Fast and Lightweight AutoML Library](https://arxiv.org/abs/1911.04706). Chi Wang, Qingyun Wu, Markus Weimer, Erkang Zhu. To appear in MLSys, 2021.
+* [FLAML: A Fast and Lightweight AutoML Library](https://www.microsoft.com/en-us/research/publication/flaml-a-fast-and-lightweight-automl-library/). Chi Wang, Qingyun Wu, Markus Weimer, Erkang Zhu. To appear in MLSys, 2021.
 ```
@inproceedings{wang2021flaml,
    title={FLAML: A Fast and Lightweight AutoML Library},
@ -127,7 +131,7 @@ For more technical details, please check our papers.
 }
 ```
 * [Frugal Optimization for Cost-related Hyperparameters](https://arxiv.org/abs/2005.01571). Qingyun Wu, Chi Wang, Silu Huang. AAAI 2021.
-* Economical Hyperparameter Optimization With Blended Search Strategy. Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. To appear in ICLR 2021.
+* [Economical Hyperparameter Optimization With Blended Search Strategy](https://www.microsoft.com/en-us/research/publication/economical-hyperparameter-optimization-with-blended-search-strategy/). Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. To appear in ICLR 2021.

 ## Contributing

@ -135,6 +139,8 @@ This project welcomes contributions and suggestions. Most contributions require
 Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us
 the rights to use your contribution. For details, visit <https://cla.opensource.microsoft.com>.

+If you are new to GitHub [here](https://help.github.com/categories/collaborating-with-issues-and-pull-requests/) is a detailed help source on getting involved with development on GitHub.
+
 When you submit a pull request, a CLA bot will automatically determine whether you need to provide
 a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions
 provided by the bot. You will only need to do this once across all repos using our CLA.
@ -143,6 +149,23 @@ This project has adopted the [Microsoft Open Source Code of Conduct](https://ope
 For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or
 contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.

+## Developing
+
+### Setup:
+
+```
+git clone https://github.com/microsoft/FLAML.git
+pip install -e .[test,notebook]
+```
+
+### Coverage
+Any code you commit should generally not significantly impact coverage. To run all unit tests:
+```
+coverage run -m pytest test
+```
+
+If all the tests are passed, please also test run notebook/flaml_automl to make sure your commit does not break the notebook example.
+
 ## Authors

 * Chi Wang
--- a/flaml/init.py
+++ b/flaml/init.py
@ -1,4 +1,4 @@
-from flaml.searcher import CFO, BlendSearch, FLOW2
+from flaml.searcher import CFO, BlendSearch, FLOW2, BlendSearchTuner
 from flaml.automl import AutoML, logger_formatter
 from flaml.version import __version__
 import logging
--- a/flaml/automl.py
+++ b/flaml/automl.py
@ -253,6 +253,8 @@ class AutoML:

    '''

+    from .version import __version__
+
    def __init__(self):
        self._track_iter = 0
        self._state = AutoMLState()
@ -283,6 +285,22 @@ class AutoML:
        else:
            return None

+    def best_model_for_estimator(self, estimator_name):
+        '''Return the best model found for a particular estimator
+
+        Args:
+            estimator_name: a str of the estimator's name
+        
+        Returns:
+            An object with `predict()` and `predict_proba()` method (for
+        classification), storing the best trained model for estimator_name.
+        '''
+        if estimator_name in self._search_states:
+            state = self._search_states[estimator_name]
+            if hasattr(state, 'trained_estimator'):
+                return state.trained_estimator.model
+        return None
+
    @property
    def best_estimator(self):
        '''A string indicating the best estimator found.'''        
@ -1208,9 +1226,10 @@ class AutoML:
                gap = search_state.best_loss - self._state.best_loss
                if gap > 0 and not self._ensemble:
                    delta_loss = (search_state.best_loss_old - 
-                                 search_state.best_loss)
+                                 search_state.best_loss) or \
+                                     search_state.best_loss
                    delta_time = (search_state.total_time_used - 
-                                 search_state.time_best_found_old)
+                                 search_state.time_best_found_old) or 1e-10
                    speed = delta_loss / delta_time
                    try:
                        estimated_cost = 2*gap/speed
--- a/flaml/searcher/init.py
+++ b/flaml/searcher/init.py
@ -1,2 +1,2 @@
-from .blendsearch import CFO, BlendSearch
+from .blendsearch import CFO, BlendSearch, BlendSearchTuner
 from .flow2 import FLOW2
--- a/flaml/searcher/blendsearch.py
+++ b/flaml/searcher/blendsearch.py
@ -7,6 +7,7 @@ from typing import Dict, Optional, List, Tuple
 import numpy as np
 import time
 import pickle
+
 try:
    from ray.tune.suggest import Searcher
    from ray.tune.suggest.optuna import OptunaSearch as GlobalSearch
@ -143,20 +144,31 @@ class BlendSearch(Searcher):
        self._deadline = np.inf

    def save(self, checkpoint_path: str):
-        save_object = (self._metric_target, self._search_thread_pool,
-            self._thread_count, self._init_used, self._trial_proposed_by,
-            self._ls_bound_min, self._ls_bound_max, self._result,
-            self._deadline)
+        save_object = self
        with open(checkpoint_path, "wb") as outputFile:
            pickle.dump(save_object, outputFile)
            
    def restore(self, checkpoint_path: str):
        with open(checkpoint_path, "rb") as inputFile:
-            save_object = pickle.load(inputFile)
-        self._metric_target, self._search_thread_pool, \
-            self._thread_count, self._init_used, self._trial_proposed_by, \
-            self._ls_bound_min, self._ls_bound_max, self._result, \
-            self._deadline = save_object
+            state = pickle.load(inputFile)
+        self._metric_target = state._metric_target
+        self._search_thread_pool = state._search_thread_pool
+        self._thread_count = state._thread_count
+        self._init_used = state._init_used
+        self._trial_proposed_by = state._trial_proposed_by
+        self._ls_bound_min = state._ls_bound_min
+        self._ls_bound_max = state._ls_bound_max
+        self._gs_admissible_min = state._gs_admissible_min
+        self._gs_admissible_max = state._gs_admissible_max
+        self._result = state._result
+        self._deadline = state._deadline
+        self._metric, self._mode = state._metric, state._mode
+        self._points_to_evaluate = state._points_to_evaluate
+        self._gs = state._gs
+        self._ls = state._ls
+        self._resources_per_trial = state._resources_per_trial
+        self._mem_size = state._mem_size
+        self._mem_threshold = state._mem_threshold

    def restore_from_dir(self, checkpoint_dir: str):
        super.restore_from_dir(checkpoint_dir)
@ -526,3 +538,87 @@ class CFO(BlendSearchTuner):
        return len(self._search_thread_pool) < 2


+def create_next(client):
+    '''A stateless API for HPO
+    '''
+    state = client.get_state()
+    setting = client.get_settings_dict()
+    if state is None:
+        # first time call
+        try:
+            from ray.tune import (uniform, quniform, choice, randint, qrandint, randn,
+        qrandn, loguniform, qloguniform)
+            from ray.tune.trial import Trial
+        except:
+            from ..tune.sample import (uniform, quniform, choice, randint, qrandint, randn,
+        qrandn, loguniform, qloguniform)
+            from ..tune.trial import Trial
+        method = setting.get('method', 'BlendSearch')
+        mode = client.get_optimization_mode()
+        if mode == 'minimize':
+            mode = 'min'
+        elif mode == 'maximize':
+            mode = 'max'
+        metric = client.get_primary_metric()
+        hp_space = client.get_hyperparameter_space_dict()
+        space = {}
+        for key, value in hp_space.items():
+            t = value["type"]
+            if t == 'continuous':
+                space[key] = uniform(value["min_val"], value["max_val"])
+            elif t == 'discrete':
+                space[key] = choice(value["values"])
+            elif t == 'integral':
+                space[key] = randint(value["min_val"], value["max_val"])
+            elif t == 'quantized_continuous':
+                space[key] = quniform(value["min_val"], value["max_val"],
+                 value["step"])
+        init_config = setting.get('init_config', None)
+        if init_config:
+            points_to_evaluate = [init_config]
+        else:
+            points_to_evaluate = None
+        cat_hp_cost = setting.get('cat_hp_cost', None)
+
+        if method == 'BlendSearch':
+            Algo = BlendSearch
+        elif method == 'CFO':
+            Algo = CFO
+        algo = Algo(
+            mode=mode, 
+            metric=metric, 
+            space=space,
+            points_to_evaluate=points_to_evaluate,
+            cat_hp_cost=cat_hp_cost,
+            )
+        time_budget_s = setting.get('time_budget_s', None)
+        if time_budget_s:
+            algo._deadline = time_budget_s + time.time()
+        config2trialid = {}
+    else:
+        algo = state['algo']
+        config2trialid = state['config2trialid']
+    # update finished trials
+    trials_completed = []
+    for trial in client.get_trials():
+        if trial.end_time is not None:
+            signature = algo._ls.config_signature(trial.hp_sample)
+            if not algo._result[signature]:
+                trials_completed.append((trial.end_time, trial))
+    trials_completed.sort()
+    for t in trials_completed:
+        end_time, trial = t
+        trial_id = config2trialid[trial.hp_sample]
+        result = {}
+        result[algo.metric] = trial.metrics[algo.metric].values[-1]
+        result[algo.cost_attr] = (end_time - trial.start_time).total_seconds()
+        for key, value in trial.hp_sample.items():
+            result['config/'+key] = value
+        algo.on_trial_complete(trial_id, result=result)
+    # propose new trial
+    trial_id = Trial.generate_id()
+    config = algo.suggest(trial_id)
+    if config:
+        config2trialid[config] = trial_id
+        client.launch_trial(config)
+    client.update_state({'algo': algo, 'config2trialid': config2trialid})
--- a/flaml/training_log.py
+++ b/flaml/training_log.py
@ -118,6 +118,7 @@ class TrainingLogWriter(object):

    def close(self):
        self.file.close()
+        self.file = None # for pickle


 class TrainingLogReader(object):
@ -141,6 +142,7 @@ class TrainingLogReader(object):

    def close(self):
        self.file.close()
+        self.file = None # for pickle

    def get_record(self, record_id) -> TrainingLogRecord:
        if self.file is None:
--- a/flaml/tune/README.md
+++ b/flaml/tune/README.md
@ -172,7 +172,7 @@ For more technical details, please check our papers.
 }
 ```

-* Economical Hyperparameter Optimization With Blended Search Strategy. Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. To appear in ICLR 2021.
+* [Economical Hyperparameter Optimization With Blended Search Strategy](https://www.microsoft.com/en-us/research/publication/economical-hyperparameter-optimization-with-blended-search-strategy/). Chi Wang, Qingyun Wu, Silu Huang, Amin Saied. To appear in ICLR 2021.

 ```
@inproceedings{wang2021blendsearch,
--- a/flaml/version.py
+++ b/flaml/version.py
@ -1 +1 @@
-__version__ = "0.2.8"
+__version__ = "0.2.9"
--- a/notebook/flaml_automl.ipynb
+++ b/notebook/flaml_automl.ipynb
@ -385,10 +385,10 @@
   },
   "outputs": [],
   "source": [
-    "''' pickle and save the best model '''\n",
+    "''' pickle and save the automl object '''\n",
    "import pickle\n",
-    "with open('best_model.pkl', 'wb') as f:\n",
-    "    pickle.dump(automl.model, f, pickle.HIGHEST_PROTOCOL)"
+    "with open('automl.pkl', 'wb') as f:\n",
+    "    pickle.dump(automl, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
--- a/notebook/flaml_azureml.ipynb
+++ b/notebook/flaml_azureml.ipynb
@ -302,10 +302,10 @@
   },
   "outputs": [],
   "source": [
-    "''' pickle and save the best model '''\n",
+    "''' pickle and save the automl object '''\n",
    "import pickle\n",
-    "with open('best_model.pkl', 'wb') as f:\n",
-    "    pickle.dump(automl.model, f, pickle.HIGHEST_PROTOCOL)"
+    "with open('automl.pkl', 'wb') as f:\n",
+    "    pickle.dump(automl, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
--- a/notebook/flaml_lightgbm.ipynb
+++ b/notebook/flaml_lightgbm.ipynb
@ -273,10 +273,10 @@
   },
   "outputs": [],
   "source": [
-    "''' pickle and save the best model '''\n",
+    "''' pickle and save the automl object '''\n",
    "import pickle\n",
-    "with open('best_model.pkl', 'wb') as f:\n",
-    "    pickle.dump(automl.model, f, pickle.HIGHEST_PROTOCOL)"
+    "with open('automl.pkl', 'wb') as f:\n",
+    "    pickle.dump(automl, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
--- a/notebook/flaml_xgboost.ipynb
+++ b/notebook/flaml_xgboost.ipynb
@ -321,10 +321,10 @@
   },
   "outputs": [],
   "source": [
-    "''' pickle and save the best model '''\n",
+    "''' pickle and save the automl object '''\n",
    "import pickle\n",
-    "with open('best_model.pkl', 'wb') as f:\n",
-    "    pickle.dump(automl.model, f, pickle.HIGHEST_PROTOCOL)"
+    "with open('automl.pkl', 'wb') as f:\n",
+    "    pickle.dump(automl, f, pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
--- a/setup.py
+++ b/setup.py
@ -53,7 +53,7 @@ setuptools.setup(
            "optuna==2.3.0"
        ],
        "ray": [
-            "ray[tune]==1.1.0",
+            "ray[tune]==1.2.0",
            "pyyaml<5.3.1",            
        ],
        "azureml": [
--- a/test/hf/test_deberta.py
+++ b/test/hf/test_deberta.py
@ -0,0 +1,250 @@
+'''Require: pip install torch transformers datasets flaml[blendsearch,ray]
+'''
+import time
+import numpy as np
+
+try:
+    import ray
+    from datasets import (
+        load_dataset,
+        load_metric,
+    )
+    from transformers import (
+        AutoModelForSequenceClassification,
+        AutoTokenizer,
+        Trainer,
+        TrainingArguments,
+    )
+    MODEL_CHECKPOINT = "microsoft/deberta-base"
+    task_to_keys = {
+        "cola": ("sentence", None),
+        "mnli": ("premise", "hypothesis"),
+        "mrpc": ("sentence1", "sentence2"),
+        "qnli": ("question", "sentence"),
+        "qqp": ("question1", "question2"),
+        "rte": ("sentence1", "sentence2"),
+        "sst2": ("sentence", None),
+        "stsb": ("sentence1", "sentence2"),
+        "wnli": ("sentence1", "sentence2"),
+    }
+    max_seq_length=128
+    overwrite_cache=False
+    pad_to_max_length=True
+    padding = "max_length"
+
+    TASK = "qnli"
+    # HP_METRIC, MODE = "loss", "min"
+    HP_METRIC, MODE = "accuracy", "max"
+
+    sentence1_key, sentence2_key = task_to_keys[TASK]
+    # Define tokenize method
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, use_fast=True)
+
+    def tokenize(examples):
+        args = (
+            (examples[sentence1_key],) if sentence2_key is None else (
+                examples[sentence1_key], examples[sentence2_key])
+        )
+        return tokenizer(*args, padding=padding, max_length=max_seq_length,
+         truncation=True)
+
+except:
+    print("pip install torch transformers datasets flaml[blendsearch,ray]")
+    
+import logging
+logger = logging.getLogger(__name__)
+import os
+os.makedirs('logs', exist_ok=True)
+logger.addHandler(logging.FileHandler('logs/tune_deberta.log'))
+logger.setLevel(logging.INFO)
+
+import flaml
+
+def train_deberta(config: dict):
+
+    # Load dataset and apply tokenizer
+    data_raw = load_dataset("glue", TASK)
+    data_encoded = data_raw.map(tokenize, batched=True)
+    train_dataset, eval_dataset = data_encoded["train"], data_encoded["validation"]
+
+    NUM_LABELS = len(train_dataset.features["label"].names)
+
+    metric = load_metric("glue", TASK)
+
+    def compute_metrics(eval_pred):
+        predictions, labels = eval_pred
+        predictions = np.argmax(predictions, axis=1)
+        return metric.compute(predictions=predictions, references=labels)
+
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_CHECKPOINT, num_labels=NUM_LABELS
+    )
+
+    training_args = TrainingArguments(
+        output_dir='.',
+        do_eval=False,
+        disable_tqdm=True,
+        logging_steps=20000,
+        save_total_limit=0,
+        fp16=True,
+        **config,
+    )
+
+    trainer = Trainer(
+        model,
+        training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+    )
+
+    # train model
+    trainer.train()
+
+    # evaluate model
+    eval_output = trainer.evaluate()
+
+    flaml.tune.report(
+        loss=eval_output["eval_loss"],
+        accuracy=eval_output["eval_accuracy"],
+        )
+
+    try:
+        from azureml.core import Run
+        run = Run.get_context()
+        run.log('accuracy', eval_output["eval_accuracy"])
+        run.log('loss', eval_output["eval_loss"])
+        run.log('config', config)
+    except: pass
+
+def _test_deberta(method='BlendSearch'):
+ 
+    max_num_epoch = 100
+    num_samples = -1
+    time_budget_s = 3600
+
+    search_space = {
+        # You can mix constants with search space objects.
+        "num_train_epochs": flaml.tune.loguniform(1, max_num_epoch),
+        "learning_rate": flaml.tune.loguniform(3e-5, 1.5e-4),
+        "weight_decay": flaml.tune.uniform(0, 0.3),
+        "per_device_train_batch_size": flaml.tune.choice([16, 32, 64, 128]),
+        "seed": flaml.tune.choice([12, 22, 33, 42]),
+    }
+
+    start_time = time.time()
+    ray.init(num_cpus=4, num_gpus=4)
+    if 'ASHA' == method:
+        algo = None
+    elif 'BOHB' == method:
+        from ray.tune.schedulers import HyperBandForBOHB
+        from ray.tune.suggest.bohb import tuneBOHB
+        algo = tuneBOHB(max_concurrent=4)
+        scheduler = HyperBandForBOHB(max_t=max_num_epoch)
+    elif 'Optuna' == method:
+        from ray.tune.suggest.optuna import OptunaSearch
+        algo = OptunaSearch()
+    elif 'CFO' == method:
+        from flaml import CFO
+        algo = CFO(points_to_evaluate=[{
+            "num_train_epochs": 1,
+            "per_device_train_batch_size": 128,
+        }])
+    elif 'BlendSearch' == method:
+        from flaml import BlendSearch
+        algo = BlendSearch(points_to_evaluate=[{
+            "num_train_epochs": 1,
+            "per_device_train_batch_size": 128,
+        }])
+    elif 'Dragonfly' == method:
+        from ray.tune.suggest.dragonfly import DragonflySearch
+        algo = DragonflySearch()
+    elif 'SkOpt' == method:
+        from ray.tune.suggest.skopt import SkOptSearch
+        algo = SkOptSearch()
+    elif 'Nevergrad' == method:
+        from ray.tune.suggest.nevergrad import NevergradSearch
+        import nevergrad as ng
+        algo = NevergradSearch(optimizer=ng.optimizers.OnePlusOne)
+    elif 'ZOOpt' == method:
+        from ray.tune.suggest.zoopt import ZOOptSearch
+        algo = ZOOptSearch(budget=num_samples)
+    elif 'Ax' == method:
+        from ray.tune.suggest.ax import AxSearch
+        algo = AxSearch(max_concurrent=3)
+    elif 'HyperOpt' == method:
+        from ray.tune.suggest.hyperopt import HyperOptSearch
+        algo = HyperOptSearch()
+        scheduler = None
+    if method != 'BOHB':
+        from ray.tune.schedulers import ASHAScheduler
+        scheduler = ASHAScheduler(
+            max_t=max_num_epoch,
+            grace_period=1)
+    scheduler = None
+    analysis = ray.tune.run(
+        train_deberta,
+        metric=HP_METRIC,
+        mode=MODE,
+        resources_per_trial={"gpu": 4, "cpu": 4},
+        config=search_space, local_dir='logs/',
+        num_samples=num_samples, time_budget_s=time_budget_s,
+        keep_checkpoints_num=1, checkpoint_score_attr=HP_METRIC,
+        scheduler=scheduler, search_alg=algo)
+
+    ray.shutdown()
+
+    best_trial = analysis.get_best_trial(HP_METRIC, MODE, "all")
+    metric = best_trial.metric_analysis[HP_METRIC][MODE]
+
+    logger.info(f"method={method}")
+    logger.info(f"n_trials={len(analysis.trials)}")
+    logger.info(f"time={time.time()-start_time}")
+    logger.info(f"Best model eval {HP_METRIC}: {metric:.4f}")
+    logger.info(f"Best model parameters: {best_trial.config}")
+
+
+def _test_deberta_cfo():
+    _test_deberta('CFO')
+
+
+def _test_deberta_dragonfly():
+    _test_deberta('Dragonfly')
+
+
+def _test_deberta_skopt():
+    _test_deberta('SkOpt')
+
+
+def _test_deberta_nevergrad():
+    _test_deberta('Nevergrad')
+
+
+def _test_deberta_zoopt():
+    _test_deberta('ZOOpt')
+
+
+def _test_deberta_ax():
+    _test_deberta('Ax')
+
+
+def __test_deberta_hyperopt():
+    _test_deberta('HyperOpt')
+
+
+def _test_deberta_optuna():
+    _test_deberta('Optuna')
+
+
+def _test_deberta_asha():
+    _test_deberta('ASHA')
+
+
+def _test_deberta_bohb():
+    _test_deberta('BOHB')
+
+
+if __name__ == "__main__":
+    _test_deberta()
--- a/test/hf/test_electra.py
+++ b/test/hf/test_electra.py
@ -130,14 +130,8 @@ def _test_electra(method='BlendSearch'):
        "num_train_epochs": flaml.tune.loguniform(1, max_num_epoch),
        "learning_rate": flaml.tune.loguniform(3e-5, 1.5e-4),
        "weight_decay": flaml.tune.uniform(0, 0.3),
-        # "warmup_ratio": flaml.tune.uniform(0, 0.2),
-        # "hidden_dropout_prob": flaml.tune.uniform(0, 0.2),
-        # "attention_probs_dropout_prob": flaml.tune.uniform(0, 0.2),
        "per_device_train_batch_size": flaml.tune.choice([16, 32, 64, 128]),
        "seed": flaml.tune.choice([12, 22, 33, 42]),
-        # "adam_beta1": flaml.tune.uniform(0.8, 0.99),
-        # "adam_beta2": flaml.tune.loguniform(98e-2, 9999e-4),
-        # "adam_epsilon": flaml.tune.loguniform(1e-9, 1e-7),
    }

    start_time = time.time()
--- a/test/hf/test_roberta.py
+++ b/test/hf/test_roberta.py
@ -0,0 +1,251 @@
+'''Require: pip install torch transformers datasets flaml[blendsearch,ray]
+'''
+import time
+import numpy as np
+
+try:
+    import ray
+    from datasets import (
+        load_dataset,
+        load_metric,
+    )
+    from transformers import (
+        AutoModelForSequenceClassification,
+        AutoTokenizer,
+        Trainer,
+        TrainingArguments,
+    )
+    MODEL_CHECKPOINT = "roberta-base"
+    task_to_keys = {
+        "cola": ("sentence", None),
+        "mnli": ("premise", "hypothesis"),
+        "mrpc": ("sentence1", "sentence2"),
+        "qnli": ("question", "sentence"),
+        "qqp": ("question1", "question2"),
+        "rte": ("sentence1", "sentence2"),
+        "sst2": ("sentence", None),
+        "stsb": ("sentence1", "sentence2"),
+        "wnli": ("sentence1", "sentence2"),
+    }
+    max_seq_length=128
+    overwrite_cache=False
+    pad_to_max_length=True
+    padding = "max_length"
+
+    TASK = "qnli"
+    # HP_METRIC, MODE = "loss", "min"
+    HP_METRIC, MODE = "accuracy", "max"
+
+    sentence1_key, sentence2_key = task_to_keys[TASK]
+    # Define tokenize method
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, use_fast=True)
+
+    def tokenize(examples):
+        args = (
+            (examples[sentence1_key],) if sentence2_key is None else (
+                examples[sentence1_key], examples[sentence2_key])
+        )
+        return tokenizer(*args, padding=padding, max_length=max_seq_length,
+         truncation=True)
+
+except:
+    print("pip install torch transformers datasets flaml[blendsearch,ray]")
+    
+import logging
+logger = logging.getLogger(__name__)
+import os
+os.makedirs('logs', exist_ok=True)
+logger.addHandler(logging.FileHandler('logs/tune_roberta.log'))
+logger.setLevel(logging.INFO)
+
+import flaml
+
+def train_roberta(config: dict):
+
+    # Load dataset and apply tokenizer
+    data_raw = load_dataset("glue", TASK)
+    data_encoded = data_raw.map(tokenize, batched=True)
+    train_dataset, eval_dataset = data_encoded["train"], data_encoded["validation"]
+
+    NUM_LABELS = len(train_dataset.features["label"].names)
+
+    metric = load_metric("glue", TASK)
+
+    def compute_metrics(eval_pred):
+        predictions, labels = eval_pred
+        predictions = np.argmax(predictions, axis=1)
+        return metric.compute(predictions=predictions, references=labels)
+
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_CHECKPOINT, num_labels=NUM_LABELS
+    )
+
+    training_args = TrainingArguments(
+        output_dir='.',
+        do_eval=False,
+        disable_tqdm=True,
+        logging_steps=20000,
+        save_total_limit=0,
+        fp16=True,
+        **config,
+    )
+
+    trainer = Trainer(
+        model,
+        training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+    )
+
+    # train model
+    trainer.train()
+
+    # evaluate model
+    eval_output = trainer.evaluate()
+
+    flaml.tune.report(
+        loss=eval_output["eval_loss"],
+        accuracy=eval_output["eval_accuracy"],
+        )
+
+    try:
+        from azureml.core import Run
+        run = Run.get_context()
+        run.log('accuracy', eval_output["eval_accuracy"])
+        run.log('loss', eval_output["eval_loss"])
+        run.log('config', config)
+    except: pass
+
+def _test_roberta(method='BlendSearch'):
+ 
+    max_num_epoch = 100
+    num_samples = -1
+    time_budget_s = 3600
+
+    search_space = {
+        # You can mix constants with search space objects.
+        "num_train_epochs": flaml.tune.loguniform(1, max_num_epoch),
+        "learning_rate": flaml.tune.loguniform(1e-5, 3e-5),
+        "weight_decay": flaml.tune.uniform(0, 0.3),
+        "per_device_train_batch_size": flaml.tune.choice([16, 32, 64, 128]),
+        "seed": flaml.tune.choice([12, 22, 33, 42]),
+    }
+
+    start_time = time.time()
+    ray.init(num_cpus=4, num_gpus=4)
+    if 'ASHA' == method:
+        algo = None
+    elif 'BOHB' == method:
+        from ray.tune.schedulers import HyperBandForBOHB
+        from ray.tune.suggest.bohb import tuneBOHB
+        algo = tuneBOHB(max_concurrent=4)
+        scheduler = HyperBandForBOHB(max_t=max_num_epoch)
+    elif 'Optuna' == method:
+        from ray.tune.suggest.optuna import OptunaSearch
+        algo = OptunaSearch()
+    elif 'CFO' == method:
+        from flaml import CFO
+        algo = CFO(points_to_evaluate=[{
+            "num_train_epochs": 1,
+            "per_device_train_batch_size": 128,
+        }])
+    elif 'BlendSearch' == method:
+        from flaml import BlendSearch
+        algo = BlendSearch(points_to_evaluate=[{
+            "num_train_epochs": 1,
+            "per_device_train_batch_size": 128,
+        }])
+    elif 'Dragonfly' == method:
+        from ray.tune.suggest.dragonfly import DragonflySearch
+        algo = DragonflySearch()
+    elif 'SkOpt' == method:
+        from ray.tune.suggest.skopt import SkOptSearch
+        algo = SkOptSearch()
+    elif 'Nevergrad' == method:
+        from ray.tune.suggest.nevergrad import NevergradSearch
+        import nevergrad as ng
+        algo = NevergradSearch(optimizer=ng.optimizers.OnePlusOne)
+    elif 'ZOOpt' == method:
+        from ray.tune.suggest.zoopt import ZOOptSearch
+        algo = ZOOptSearch(budget=num_samples)
+    elif 'Ax' == method:
+        from ray.tune.suggest.ax import AxSearch
+        algo = AxSearch(max_concurrent=3)
+    elif 'HyperOpt' == method:
+        from ray.tune.suggest.hyperopt import HyperOptSearch
+        algo = HyperOptSearch()
+        scheduler = None
+    if method != 'BOHB':
+        from ray.tune.schedulers import ASHAScheduler
+        scheduler = ASHAScheduler(
+            max_t=max_num_epoch,
+            grace_period=1)
+    scheduler = None
+    analysis = ray.tune.run(
+        train_roberta,
+        metric=HP_METRIC,
+        mode=MODE,
+        resources_per_trial={"gpu": 4, "cpu": 4},
+        config=search_space, local_dir='logs/',
+        num_samples=num_samples, time_budget_s=time_budget_s,
+        keep_checkpoints_num=1, checkpoint_score_attr=HP_METRIC,
+        scheduler=scheduler, search_alg=algo)
+
+    ray.shutdown()
+
+    best_trial = analysis.get_best_trial(HP_METRIC, MODE, "all")
+    metric = best_trial.metric_analysis[HP_METRIC][MODE]
+
+    logger.info(f"method={method}")
+    logger.info(f"n_trials={len(analysis.trials)}")
+    logger.info(f"time={time.time()-start_time}")
+    logger.info(f"Best model eval {HP_METRIC}: {metric:.4f}")
+    logger.info(f"Best model parameters: {best_trial.config}")
+
+
+def _test_roberta_cfo():
+    _test_roberta('CFO')
+
+
+def _test_roberta_dragonfly():
+    _test_roberta('Dragonfly')
+
+
+def _test_roberta_skopt():
+    _test_roberta('SkOpt')
+
+
+def _test_roberta_nevergrad():
+    _test_roberta('Nevergrad')
+
+
+def _test_roberta_zoopt():
+    _test_roberta('ZOOpt')
+
+
+def _test_roberta_ax():
+    _test_roberta('Ax')
+
+
+def __test_roberta_hyperopt():
+    _test_roberta('HyperOpt')
+
+
+def _test_roberta_optuna():
+    _test_roberta('Optuna')
+
+
+def _test_roberta_asha():
+    _test_roberta('ASHA')
+
+
+def _test_roberta_bohb():
+    _test_roberta('BOHB')
+
+
+if __name__ == "__main__":
+    _test_roberta()
+
--- a/test/test_automl.py
+++ b/test/test_automl.py
@ -98,6 +98,8 @@ class TestAutoML(unittest.TestCase):

        '''The main flaml automl API'''
        automl.fit(X_train = X_train, y_train = y_train, **settings)
+        # print the best model found for RGF
+        print(automl.best_model_for_estimator("RGF"))

    def test_ensemble(self):
        automl = AutoML()
--- a/test/test_python_log.py
+++ b/test/test_python_log.py
@ -26,7 +26,7 @@ class TestLogging(unittest.TestCase):
            logger.addHandler(ch)

            # Run a simple job.
-            automl_experiment = AutoML()
+            automl = AutoML()
            automl_settings = {
                "time_budget": 1,
                "metric": 'mse',
@ -34,13 +34,18 @@ class TestLogging(unittest.TestCase):
                "log_file_name": training_log,
                "log_training_metric": True,
                "n_jobs": 1,
-                "model_history": True
+                "model_history": True,
            }
            X_train, y_train = load_boston(return_X_y=True)
            n = len(y_train) >> 1
-            automl_experiment.fit(X_train=X_train[:n], y_train=y_train[:n],
+            automl.fit(X_train=X_train[:n], y_train=y_train[:n],
                                  X_val=X_train[n:], y_val=y_train[n:],
                                  **automl_settings)

            # Check if the log buffer is populated.
            self.assertTrue(len(buf.getvalue()) > 0)
+
+        import pickle
+        with open('automl.pkl', 'wb') as f:
+            pickle.dump(automl, f, pickle.HIGHEST_PROTOCOL)
+        print(automl.__version__)