data validation (#45)

* pickle the AutoML object

* get best model per estimator

* test deberta

* stateless API

* prevent divide by zero

* test roberta

* BlendSearchTuner

* delta time

* reindex columns when dropping int-indexed columns

* test drop columns and small training data

* param set for ensemble builder

* fillna on copy

Co-authored-by: Chi Wang (MSR) <chiw@microsoft.com>

.gitignore

@@ -151,3 +151,5 @@ catboost_info
 notebook/*.pkl
 notebook/.azureml
 mlruns
+logs
+automl.pkl

README.md

@@ -2,6 +2,7 @@
 [![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)
 
 # 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

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

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.'''        
@@ -571,6 +589,12 @@ class AutoML:
             self._state.y_val = (X_train, y_train, X_val, y_val)
         if self._split_type == "stratified":
             logger.info("Using StratifiedKFold")
+            assert y_train_all.size >= n_splits, (
+                f"{n_splits}-fold cross validation"
+                f" requires input data with at least {n_splits} examples.")
+            assert y_train_all.size >= 2*n_splits, (
+                f"{n_splits}-fold cross validation with metric=r2 "
+                f"requires input data with at least {n_splits*2} examples.")
             self._state.kf = RepeatedStratifiedKFold(n_splits=n_splits,
              n_repeats=1, random_state=RANDOM_SEED)
         else:
@@ -765,6 +789,7 @@ class AutoML:
             split_type="stratified",
             learner_selector='sample',
             hpo_method=None,
+            verbose=1,
             **fit_kwargs):
         '''Find a model for a given task
 
@@ -823,6 +848,8 @@ class AutoML:
             y_val: None | a numpy array or a pandas series of validation labels
             sample_weight_val: None | a numpy array of the sample weight of
                 validation data
+            verbose: int, default=1 | Controls the verbosity, higher means more
+                messages
             **fit_kwargs: Other key word arguments to pass to fit() function of
                 the searched learners, such sample_weight
         '''
@@ -835,6 +862,10 @@ class AutoML:
         self._search_states = {}  #key: estimator name; value: SearchState
         self._random = np.random.RandomState(RANDOM_SEED)
         self._learner_selector = learner_selector
+        old_level = logger.getEffectiveLevel()
+        self.verbose = verbose
+        if verbose==0:
+            logger.setLevel(logging.WARNING)
         if self._state.task == 'classification':
             self._state.task = get_classification_objective(
                 len(np.unique(self._y_train_all)))
@@ -845,7 +876,7 @@ class AutoML:
         if eval_method == 'auto' or self._state.X_val is not None:
             eval_method = self._decide_eval_method(time_budget)
         self._state.eval_method = eval_method
-        if not mlflow or not mlflow.active_run() and not logger.handler:
+        if (not mlflow or not mlflow.active_run()) and not logger.handlers:
             # Add the console handler.
             _ch = logging.StreamHandler()
             _ch.setFormatter(logger_formatter)
@@ -906,6 +937,8 @@ class AutoML:
             self._state.n_jobs = n_jobs
             self._search()
             logger.info("fit succeeded")
+        if verbose==0:
+            logger.setLevel(old_level)
 
     def _search(self):
         # initialize the search_states
@@ -1018,7 +1051,7 @@ class AutoML:
                 init_config=None, 
                 search_alg=search_state.search_alg,
                 time_budget_s=budget_left,
-                verbose=0, local_dir='logs/tune_results',
+                verbose=max(self.verbose-1,0), #local_dir='logs/tune_results',
                 use_ray=False,
                 )
             # warnings.resetwarnings()
@@ -1074,7 +1107,7 @@ class AutoML:
                                         search_state.best_config,
                                         estimator,
                                         search_state.sample_size)                
-                    if mlflow is not None:
+                    if mlflow is not None and mlflow.active_run():
                         with mlflow.start_run(nested=True) as run:
                             mlflow.log_metric('iter_counter',
                                 self._iter_per_learner[estimator])
@@ -1208,9 +1241,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

flaml/data.py

@@ -192,11 +192,13 @@ class DataTransformer:
             X = X.copy()
             n = X.shape[0]
             cat_columns, num_columns = [], []
+            drop = False
             for column in X.columns:
                 if X[column].dtype.name in ('object', 'category'):
                     if X[column].nunique() == 1 or X[column].nunique(
                             dropna=True) == n - X[column].isnull().sum():
                         X.drop(columns=column, inplace=True)
+                        drop = True
                     elif X[column].dtype.name == 'category':
                         current_categories = X[column].cat.categories
                         if '__NAN__' not in current_categories:
@@ -204,27 +206,33 @@ class DataTransformer:
                                 '__NAN__').fillna('__NAN__')
                         cat_columns.append(column)
                     else:
-                        X[column].fillna('__NAN__', inplace=True)
+                        X[column] = X[column].fillna('__NAN__')
                         cat_columns.append(column)
                 else:
                     # print(X[column].dtype.name)
                     if X[column].nunique(dropna=True) < 2:
                         X.drop(columns=column, inplace=True)
+                        drop = True
                     else:
-                        X[column].fillna(np.nan, inplace=True)
+                        X[column] = X[column].fillna(np.nan)
                         num_columns.append(column)
             X = X[cat_columns + num_columns]
             if cat_columns:
                 X[cat_columns] = X[cat_columns].astype('category')
             if num_columns:
+                X_num = X[num_columns]
+                if drop and np.issubdtype(X_num.columns.dtype, np.integer):
+                    X_num.columns = range(X_num.shape[1])
+                else: drop = False
                 from sklearn.impute import SimpleImputer
                 from sklearn.compose import ColumnTransformer
                 self.transformer = ColumnTransformer([(
                     'continuous',
                     SimpleImputer(missing_values=np.nan, strategy='median'),
-                    num_columns)])
-                X[num_columns] = self.transformer.fit_transform(X)
+                        X_num.columns)])
+                X[num_columns] = self.transformer.fit_transform(X_num)
             self._cat_columns, self._num_columns = cat_columns, num_columns
+            self._drop = drop
             
         if task == 'regression':
             self.label_transformer = None
@@ -241,7 +249,7 @@ class DataTransformer:
             for column in cat_columns:
                 # print(column, X[column].dtype.name)
                 if X[column].dtype.name == 'object':
-                    X[column].fillna('__NAN__', inplace=True)
+                    X[column] = X[column].fillna('__NAN__')
                 elif X[column].dtype.name == 'category':
                     current_categories = X[column].cat.categories
                     if '__NAN__' not in current_categories:
@@ -250,6 +258,8 @@ class DataTransformer:
             if cat_columns:
                 X[cat_columns] = X[cat_columns].astype('category')
             if num_columns:
-                X[num_columns].fillna(np.nan, inplace=True)
-                X[num_columns] = self.transformer.transform(X)
+                X_num = X[num_columns].fillna(np.nan)
+                if self._drop:
+                    X_num.columns = range(X_num.shape[1])
+                X[num_columns] = self.transformer.transform(X_num)
         return X

flaml/model.py

@@ -239,11 +239,8 @@ class LGBMEstimator(BaseEstimator):
         else: objective = 'regression'
         self.params = {
             "n_estimators": int(round(n_estimators)),
-            "num_leaves":  params[
-                'num_leaves'] if 'num_leaves' in params else int(
-                    round(max_leaves)),
-            'objective': params[
-                "objective"] if "objective" in params else objective,
+            "num_leaves":  params.get('num_leaves', int(round(max_leaves))),
+            'objective': params.get("objective", objective),
             'n_jobs': n_jobs,
             'learning_rate': float(learning_rate),
             'reg_alpha': float(reg_alpha),
@@ -359,18 +356,17 @@ class XGBoostEstimator(SKLearnEstimator):
         self._max_leaves = int(round(max_leaves))
         self.params = {
             'max_leaves': int(round(max_leaves)),
-            'max_depth': 0,
-            'grow_policy': params[
-                "grow_policy"] if "grow_policy" in params else 'lossguide',
-            'tree_method':tree_method,
-            'verbosity': 0,
-            'nthread':n_jobs,
+            'max_depth': params.get('max_depth', 0),
+            'grow_policy': params.get("grow_policy", 'lossguide'),
+            'tree_method': tree_method,
+            'verbosity': params.get('verbosity', 0),
+            'nthread': n_jobs,
             'learning_rate': float(learning_rate),
             'subsample': float(subsample),
             'reg_alpha': float(reg_alpha),
             'reg_lambda': float(reg_lambda),
             'min_child_weight': float(min_child_weight),
-            'booster': params['booster'] if 'booster' in params else 'gbtree',
+            'booster': params.get('booster', 'gbtree'),
             'colsample_bylevel': float(colsample_bylevel),
             'colsample_bytree':float(colsample_bytree),
             }
@@ -429,9 +425,8 @@ class XGBoostSklearnEstimator(SKLearnEstimator, LGBMEstimator):
         "n_estimators": int(round(n_estimators)),
         'max_leaves': int(round(max_leaves)),
         'max_depth': 0,
-        'grow_policy': params[
-                "grow_policy"] if "grow_policy" in params else 'lossguide',
-        'tree_method':tree_method,
+        'grow_policy': params.get("grow_policy", 'lossguide'),
+        'tree_method': tree_method,
         'verbosity': 0,
         'n_jobs': n_jobs,
         'learning_rate': float(learning_rate),
@@ -439,7 +434,7 @@ class XGBoostSklearnEstimator(SKLearnEstimator, LGBMEstimator):
         'reg_alpha': float(reg_alpha),
         'reg_lambda': float(reg_lambda),
         'min_child_weight': float(min_child_weight),
-        'booster': params['booster'] if 'booster' in params else 'gbtree',
+        'booster': params.get('booster', 'gbtree'),
         'colsample_bylevel': float(colsample_bylevel),
         'colsample_bytree': float(colsample_bytree),
         }
@@ -544,10 +539,10 @@ class LRL1Classifier(SKLearnEstimator):
         **params):
         super().__init__(task, **params)
         self.params = {
-            'penalty': 'l1',
+            'penalty': params.get("penalty", 'l1'),
             'tol': float(tol),
             'C': float(C),
-            'solver': 'saga',
+            'solver': params.get("solver", 'saga'),
             'n_jobs': n_jobs,
         }
         if 'regression' in task:
@@ -573,10 +568,10 @@ class LRL2Classifier(SKLearnEstimator):
         **params):
         super().__init__(task, **params)
         self.params = {
-            'penalty': 'l2',
+            'penalty': params.get("penalty", 'l2'),
             'tol': float(tol),
             'C': float(C),
-            'solver': 'lbfgs',
+            'solver': params.get("solver", 'lbfgs'),
             'n_jobs': n_jobs,
         }
         if 'regression' in task:
@@ -625,9 +620,8 @@ class CatBoostEstimator(BaseEstimator):
             "n_estimators": n_estimators, 
             'learning_rate': learning_rate,
             'thread_count': n_jobs,
-            'verbose': False,
-            'random_seed': params[
-                "random_seed"] if "random_seed" in params else 10242048,
+            'verbose': params.get('verbose', False),
+            'random_seed': params.get("random_seed", 10242048),
         }
         if 'regression' in task:
             from catboost import CatBoostRegressor
@@ -724,7 +718,7 @@ class KNeighborsEstimator(BaseEstimator):
         super().__init__(task, **params)
         self.params= {
             'n_neighbors': int(round(n_neighbors)),
-            'weights': 'distance',
+            'weights': params.get('weights', 'distance'),
             'n_jobs': n_jobs,
         }
         if 'regression' in task:

flaml/searcher/__init__.py

@@ -1,2 +1,2 @@
-from .blendsearch import CFO, BlendSearch
+from .blendsearch import CFO, BlendSearch, BlendSearchTuner
 from .flow2 import FLOW2

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
@@ -25,6 +26,8 @@ class BlendSearch(Searcher):
     '''class for BlendSearch algorithm
     '''
 
+    cost_attr = "time_total_s" # cost attribute in result
+
     def __init__(self,
                  metric: Optional[str] = None,
                  mode: Optional[str] = None,
@@ -133,26 +136,39 @@ class BlendSearch(Searcher):
         self._thread_count = 1 # total # threads created
         self._init_used = self._ls.init_config is None
         self._trial_proposed_by = {} # trial_id: str -> thread_id: int
-        self._admissible_min = self._ls.normalize(self._ls.init_config)
-        self._admissible_max = self._admissible_min.copy()
+        self._ls_bound_min = self._ls.normalize(self._ls.init_config)
+        self._ls_bound_max = self._ls_bound_min.copy()
+        self._gs_admissible_min = self._ls_bound_min.copy()
+        self._gs_admissible_max = self._ls_bound_max.copy()
         self._result = {} # config_signature: tuple -> result: Dict
         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._admissible_min, self._admissible_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._admissible_min, self._admissible_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)
@@ -179,25 +195,20 @@ class BlendSearch(Searcher):
             # update target metric if improved
             if (result[self._metric]-self._metric_target)*self._ls.metric_op<0:
                 self._metric_target = result[self._metric]
-            if thread_id: # from local search
-                # update admissible region
-                normalized_config = self._ls.normalize(config)
-                for key in self._admissible_min:
-                    value = normalized_config[key]
-                    if value > self._admissible_max[key]:
-                        self._admissible_max[key] = value
-                    elif value < self._admissible_min[key]:
-                        self._admissible_min[key] = value
-            elif self._create_condition(result):
+            if not thread_id and self._create_condition(result): 
                 # thread creator
                 self._search_thread_pool[self._thread_count] = SearchThread(
                     self._ls.mode,
                     self._ls.create(config, result[self._metric], cost=result[
-                        "time_total_s"])
+                        self.cost_attr])
                 )
                 thread_id = self._thread_count
                 self._thread_count += 1
-                
+                self._update_admissible_region(config, self._ls_bound_min,
+                    self._ls_bound_max)
+            # reset admissible region to ls bounding box
+            self._gs_admissible_min.update(self._ls_bound_min)
+            self._gs_admissible_max.update(self._ls_bound_max)
         # cleaner
         # logger.info(f"thread {thread_id} in search thread pool="
         #     f"{thread_id in self._search_thread_pool}")
@@ -205,6 +216,16 @@ class BlendSearch(Searcher):
             # local search thread
             self._clean(thread_id)
 
+    def _update_admissible_region(self, config, admissible_min, admissible_max):
+        # update admissible region
+        normalized_config = self._ls.normalize(config)
+        for key in admissible_min:
+            value = normalized_config[key]
+            if value > admissible_max[key]:
+                admissible_max[key] = value
+            elif value < admissible_min[key]:
+                admissible_min[key] = value
+
     def _create_condition(self, result: Dict) -> bool:
         ''' create thread condition
         '''
@@ -232,9 +253,9 @@ class BlendSearch(Searcher):
         #     f"{self._search_thread_pool[thread_id].converged}")
         if self._search_thread_pool[thread_id].converged:
             todelete.add(thread_id)
-            for key in self._admissible_min:
-                self._admissible_max[key] += self._ls.STEPSIZE
-                self._admissible_min[key] -= self._ls.STEPSIZE            
+            for key in self._ls_bound_max:
+                self._ls_bound_max[key] += self._ls.STEPSIZE
+                self._ls_bound_min[key] -= self._ls.STEPSIZE            
         for id in todelete:
             del self._search_thread_pool[id]
 
@@ -259,50 +280,66 @@ class BlendSearch(Searcher):
         '''
         if self._init_used and not self._points_to_evaluate:
             choice, backup = self._select_thread()
-            # logger.debug(f"choice={choice}, backup={backup}")
+            # print(f"choice={choice}, backup={backup}")
             if choice < 0: return None # timeout
             self._use_rs = False
             config = self._search_thread_pool[choice].suggest(trial_id)
+            # preliminary check; not checking config validation
             skip = self._should_skip(choice, trial_id, config)
             if skip:
                 if choice: 
-                    # logger.info(f"skipping choice={choice}, config={config}")
+                    # print(f"skipping choice={choice}, config={config}")
                     return None
-                # use rs
+                # use rs when BO fails to suggest a config
                 self._use_rs = True
                 for _, generated in generate_variants(
                     {'config': self._ls.space}):
                     config = generated['config']
-                    break
+                    break # get one random config
                 # logger.debug(f"random config {config}")
                 skip = self._should_skip(choice, trial_id, config)
                 if skip: return None
-            # if not choice: logger.info(config)
-            if choice or backup == choice or self._valid(config): 
+            # if not choice: print(config)
+            if choice or self._valid(config): 
                 # LS or valid or no backup choice
                 self._trial_proposed_by[trial_id] = choice
             else: # invalid config proposed by GS
-                if not self._use_rs:
-                    self._search_thread_pool[choice].on_trial_complete(
-                        trial_id, {}, error=True) # tell GS there is an error
+                # if not self._use_rs:
+                #     self._search_thread_pool[choice].on_trial_complete(
+                #         trial_id, {}, error=True) # tell GS there is an error
                 self._use_rs = False
-                config = self._search_thread_pool[backup].suggest(trial_id)
-                skip = self._should_skip(backup, trial_id, config)
-                if skip: 
-                    return None
-                self._trial_proposed_by[trial_id] = backup
-                choice = backup
-            # if choice: self._pending.add(choice) # local search thread pending
-            if not choice:
+                if choice == backup:
+                    # use CFO's init point
+                    init_config = self._ls.init_config
+                    config = self._ls.complete_config(init_config,
+                        self._ls_bound_min, self._ls_bound_max)
+                    self._trial_proposed_by[trial_id] = choice
+                else:
+                    config = self._search_thread_pool[backup].suggest(trial_id)
+                    skip = self._should_skip(backup, trial_id, config)
+                    if skip: 
+                        return None
+                    self._trial_proposed_by[trial_id] = backup
+                    choice = backup
+            if not choice: # global search
                 if self._ls._resource: 
                 # TODO: add resource to config proposed by GS, min or median?
                     config[self._ls.prune_attr] = self._ls.min_resource
+                # temporarily relax admissible region for parallel proposals
+                self._update_admissible_region(config, self._gs_admissible_min,
+                    self._gs_admissible_max)
+            else:
+                self._update_admissible_region(config, self._ls_bound_min,
+                    self._ls_bound_max)
+                self._gs_admissible_min.update(self._ls_bound_min)
+                self._gs_admissible_max.update(self._ls_bound_max)
             self._result[self._ls.config_signature(config)] = {}
         else: # use init config
+            # print("use init config")
             init_config = self._points_to_evaluate.pop(
                 0) if self._points_to_evaluate else self._ls.init_config
             config = self._ls.complete_config(init_config,
-             self._admissible_min, self._admissible_max)
+             self._ls_bound_min, self._ls_bound_max)
                 # logger.info(f"reset config to {config}")
             config_signature = self._ls.config_signature(config)
             result = self._result.get(config_signature)
@@ -313,6 +350,7 @@ class BlendSearch(Searcher):
                 self._result[config_signature] = {}
             else: return None # running but no result yet
             self._init_used = True
+            self._trial_proposed_by[trial_id] = 0
         # logger.info(f"config={config}")
         return config
 
@@ -338,10 +376,10 @@ class BlendSearch(Searcher):
                     if choice:
                         # local search thread
                         self._clean(choice)
-                else:
-                    # tell the thread there is an error
-                    self._search_thread_pool[choice].on_trial_complete(
-                        trial_id, {}, error=True) 
+                # else:
+                #     # tell the thread there is an error
+                #     self._search_thread_pool[choice].on_trial_complete(
+                #         trial_id, {}, error=True) 
             return True
         return False
 
@@ -362,10 +400,10 @@ class BlendSearch(Searcher):
 
         top_thread_id = backup_thread_id = 0
         priority1 = priority2 = self._search_thread_pool[0].priority
-        # logger.debug(f"priority of thread 0={priority1}")
+        # print(f"priority of thread 0={priority1}, obj_best1={self._search_thread_pool[0].obj_best1}")
         for thread_id, thread in self._search_thread_pool.items():
             # if thread_id:
-            #     logger.debug(
+            #     print(
             #         f"priority of thread {thread_id}={thread.priority}")
             #     logger.debug(
             #         f"thread {thread_id}.can_suggest={thread.can_suggest}")
@@ -382,18 +420,101 @@ class BlendSearch(Searcher):
     def _valid(self, config: Dict) -> bool:
         ''' config validator
         '''
-        for key in self._admissible_min:
+        normalized_config = self._ls.normalize(config)
+        for key in self._gs_admissible_min:
             if key in config:
-                value = config[key]
+                value = normalized_config[key]
                 # logger.info(
                 #     f"{key},{value},{self._admissible_min[key]},{self._admissible_max[key]}")
-                if value<self._admissible_min[
-                    key] or value>self._admissible_max[key]:
+                if value+self._ls.STEPSIZE<self._gs_admissible_min[
+                    key] or value>self._gs_admissible_max[key]+self._ls.STEPSIZE:
                     return False
         return True
 
 
-class CFO(BlendSearch):
+try:
+    from nni.tuner import Tuner as NNITuner
+    from nni.utils import extract_scalar_reward
+    try:
+        from ray.tune import (uniform, quniform, choice, randint, qrandint, randn,
+    qrandn, loguniform, qloguniform)
+    except:
+        from ..tune.sample import (uniform, quniform, choice, randint, qrandint, randn,
+    qrandn, loguniform, qloguniform)
+
+    class BlendSearchTuner(BlendSearch, NNITuner):
+        '''Tuner class for NNI
+        '''
+
+        def receive_trial_result(self, parameter_id, parameters, value,
+         **kwargs):
+            '''
+            Receive trial's final result.
+            parameter_id: int
+            parameters: object created by 'generate_parameters()'
+            value: final metrics of the trial, including default metric
+            '''
+            result = {}
+            for key, value in parameters.items():
+                result['config/'+key] = value
+            reward = extract_scalar_reward(value)
+            result[self._metric] = reward
+            # if nni does not report training cost, 
+            # using sequence as an approximation.
+            # if no sequence, using a constant 1
+            result[self.cost_attr] = value.get(self.cost_attr, value.get(
+                'sequence', 1))
+            self.on_trial_complete(str(parameter_id), result)
+        ...
+
+        def generate_parameters(self, parameter_id, **kwargs) -> Dict:
+            '''
+            Returns a set of trial (hyper-)parameters, as a serializable object
+            parameter_id: int
+            '''            
+            return self.suggest(str(parameter_id))
+        ...
+
+        def update_search_space(self, search_space):
+            '''
+            Tuners are advised to support updating search space at run-time.
+            If a tuner can only set search space once before generating first hyper-parameters,
+            it should explicitly document this behaviour.
+            search_space: JSON object created by experiment owner
+            '''
+            config = {}
+            for key, value in search_space.items():
+                v = value.get("_value")
+                _type = value['_type']
+                if _type == 'choice':
+                    config[key] = choice(v)
+                elif _type == 'randint':
+                    config[key] = randint(v[0], v[1]-1)
+                elif _type == 'uniform':
+                    config[key] = uniform(v[0], v[1])
+                elif _type == 'quniform':
+                    config[key] = quniform(v[0], v[1], v[2])
+                elif _type == 'loguniform':
+                    config[key] = loguniform(v[0], v[1])
+                elif _type == 'qloguniform':
+                    config[key] = qloguniform(v[0], v[1], v[2])
+                elif _type == 'normal':
+                    config[key] = randn(v[1], v[2])
+                elif _type == 'qnormal':
+                    config[key] = qrandn(v[1], v[2], v[3])
+                else:
+                    raise ValueError(
+                    f'unsupported type in search_space {_type}')
+            self._ls.set_search_properties(None, None, config)
+            if self._gs is not None:
+                self._gs.set_search_properties(None, None, config)
+            self._init_search()
+
+except:
+    class BlendSearchTuner(BlendSearch): pass
+
+
+class CFO(BlendSearchTuner):
     ''' class for CFO algorithm
     '''
 
@@ -416,3 +537,89 @@ class CFO(BlendSearch):
         ''' create thread condition
         '''
         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})

flaml/searcher/flow2.py

@@ -9,9 +9,10 @@ try:
     from ray.tune.suggest import Searcher
     from ray.tune.suggest.variant_generator import generate_variants
     from ray.tune import sample
+    from ray.tune.utils.util import flatten_dict, unflatten_dict
 except ImportError:
     from .suggestion import Searcher
-    from .variant_generator import generate_variants
+    from .variant_generator import generate_variants, flatten_dict, unflatten_dict
     from ..tune import sample
 
 
@@ -86,6 +87,7 @@ class FLOW2(Searcher):
         elif mode == "min":
             self.metric_op = 1.
         self.space = space or {}
+        self.space = flatten_dict(self.space, prevent_delimiter=True)
         self._random = np.random.RandomState(seed)
         self._seed = seed
         if not init_config:
@@ -95,7 +97,8 @@ class FLOW2(Searcher):
                 "consider providing init values for cost-related hps via "
                 "'init_config'."
                 )
-        self.init_config = self.best_config = init_config
+        self.init_config = init_config
+        self.best_config = flatten_dict(init_config)
         self.cat_hp_cost = cat_hp_cost
         self.prune_attr = prune_attr
         self.min_resource = min_resource
@@ -126,16 +129,16 @@ class FLOW2(Searcher):
             if callable(getattr(domain, 'get_sampler', None)):
                 self._tunable_keys.append(key)
                 sampler = domain.get_sampler()
-                if isinstance(sampler, sample.Quantized):
-                    sampler_inner = sampler.get_sampler()
-                    if str(sampler_inner) == 'Uniform':
-                        self._step_lb = min(
-                            self._step_lb, sampler.q/(domain.upper-domain.lower))
-                elif isinstance(domain, sample.Integer) and str(
-                    sampler) == 'Uniform':
-                    self._step_lb = min(
-                        self._step_lb, 1.0/(domain.upper-domain.lower))
-                elif isinstance(domain, sample.Categorical):
+                # if isinstance(sampler, sample.Quantized):
+                #     sampler_inner = sampler.get_sampler()
+                #     if str(sampler_inner) == 'Uniform':
+                #         self._step_lb = min(
+                #             self._step_lb, sampler.q/(domain.upper-domain.lower))
+                # elif isinstance(domain, sample.Integer) and str(
+                #     sampler) == 'Uniform':
+                #     self._step_lb = min(
+                #         self._step_lb, 1.0/(domain.upper-domain.lower))
+                if isinstance(domain, sample.Categorical):
                     cat_hp_cost = self.cat_hp_cost
                     if cat_hp_cost and key in cat_hp_cost:
                         cost = np.array(cat_hp_cost[key])
@@ -146,7 +149,7 @@ class FLOW2(Searcher):
                         for i, choice in enumerate(l):
                             d[choice] = i
                         self._ordered_cat_hp[key] = (l, d)
-                        self._step_lb = min(self._step_lb, 1.0/len(l))
+                        # self._step_lb = min(self._step_lb, 1.0/len(l))
                     elif all(isinstance(x, int) or isinstance(x, float)
                      for x in domain.categories):
                         l = sorted(domain.categories)
@@ -154,10 +157,10 @@ class FLOW2(Searcher):
                         for i, choice in enumerate(l):
                             d[choice] = i
                         self._ordered_choice_hp[key] = (l, d) 
-                        self._step_lb = min(self._step_lb, 1.0/len(l))
+                        # self._step_lb = min(self._step_lb, 1.0/len(l))
                     else:
                         self._unordered_cat_hp[key] = l = len(domain.categories)
-                        self._step_lb = min(self._step_lb, 1.0/l)
+                        # self._step_lb = min(self._step_lb, 1.0/l)
                 if str(sampler) != 'Normal':
                     self._bounded_keys.append(key)
         self._space_keys = list(self.space.keys())
@@ -171,7 +174,7 @@ class FLOW2(Searcher):
             # logger.info(self._resource)
         else: self._resource = None
         self.incumbent = {}
-        self.incumbent = self.normalize(self.init_config)
+        self.incumbent = self.normalize(self.best_config) # flattened
         self.best_obj = self.cost_incumbent = None
         self.dim = len(self._tunable_keys)  # total # tunable dimensions
         self._direction_tried = None        
@@ -247,7 +250,7 @@ class FLOW2(Searcher):
                 if key not in self._unordered_cat_hp:
                     if upper and lower:
                         u, l = upper[key], lower[key]
-                        gauss_std = u-l
+                        gauss_std = u-l or self.STEPSIZE
                         # allowed bound
                         u += self.STEPSIZE
                         l -= self.STEPSIZE
@@ -261,11 +264,11 @@ class FLOW2(Searcher):
                     normalized[key] = max(l, min(u, normalized[key] + delta))
             # use best config for unordered cat choice
             config = self.denormalize(normalized)
-            self._reset_times += 1
         else:
             # first time init_config, or other configs, take as is
             config = partial_config.copy()
-
+        if partial_config == self.init_config: self._reset_times += 1
+        config = flatten_dict(config)
         for key, value in self.space.items():
             if key not in config:
                 config[key] = value
@@ -277,13 +280,13 @@ class FLOW2(Searcher):
 
         if self._resource:
             config[self.prune_attr] = self.min_resource
-        return config
+        return unflatten_dict(config)
 
     def create(self, init_config: Dict, obj: float, cost: float) -> Searcher:
         flow2 = FLOW2(init_config, self.metric, self.mode, self._cat_hp_cost,
-                      self.space, self.prune_attr, self.min_resource,
-                      self.max_resource, self.resource_multiple_factor,
-                      self._seed+1)
+                      unflatten_dict(self.space), self.prune_attr, 
+                      self.min_resource, self.max_resource, 
+                      self.resource_multiple_factor, self._seed+1)
         flow2.best_obj = obj * self.metric_op  # minimize internally
         flow2.cost_incumbent = cost
         return flow2
@@ -292,7 +295,7 @@ class FLOW2(Searcher):
         ''' normalize each dimension in config to [0,1]
         '''
         config_norm = {}
-        for key, value in config.items():
+        for key, value in flatten_dict(config).items():
             if key in self.space:
                 # domain: sample.Categorical/Integer/Float/Function
                 domain = self.space[key]
@@ -303,10 +306,10 @@ class FLOW2(Searcher):
                         # normalize categorical
                         if key in self._ordered_cat_hp:
                             l, d = self._ordered_cat_hp[key]
-                            config_norm[key] = d[value]/len(l)
+                            config_norm[key] = (d[value]+0.5)/len(l) # center
                         elif key in self._ordered_choice_hp:
                             l, d = self._ordered_choice_hp[key]
-                            config_norm[key] = d[value]/len(l)
+                            config_norm[key] = (d[value]+0.5)/len(l) # center
                         elif key in self.incumbent:
                             config_norm[key] = self.incumbent[
                                 key] if value == self.best_config[
@@ -406,6 +409,7 @@ class FLOW2(Searcher):
             self._metric = metric
         if mode:
             assert mode in ["min", "max"], "`mode` must be 'min' or 'max'."
+            self._mode = mode
             if mode == "max":
                 self.metric_op = -1.
             elif mode == "min":
@@ -426,7 +430,7 @@ class FLOW2(Searcher):
             obj = result.get(self._metric)
             if obj: 
                 obj *= self.metric_op
-                if obj < self.best_obj:
+                if self.best_obj is None or obj < self.best_obj:
                     self.best_obj, self.best_config = obj, self._configs[
                         trial_id]
                     self.incumbent = self.normalize(self.best_config)
@@ -437,7 +441,8 @@ class FLOW2(Searcher):
                     self._cost_complete4incumbent = 0
                     self._num_allowed4incumbent = 2 * self.dim
                     self._proposed_by.clear()
-                    if self._K > 0:
+                    if self._K > 0: 
+                        # self._oldK must have been set when self._K>0
                         self.step *= np.sqrt(self._K/self._oldK)
                     if self.step > self.step_ub: self.step = self.step_ub
                     self._iter_best_config = self.trial_count
@@ -474,7 +479,7 @@ class FLOW2(Searcher):
             obj = result.get(self._metric)
             if obj: 
                 obj *= self.metric_op
-                if obj < self.best_obj:
+                if self.best_obj is None or obj < self.best_obj:
                     self.best_obj = obj
                     config = self._configs[trial_id]
                     if self.best_config != config:
@@ -528,12 +533,12 @@ class FLOW2(Searcher):
         self._direction_tried = self.rand_vector_unit_sphere(
             self.dim) * self.step
         for i, key in enumerate(self._tunable_keys):
-            move[key] += self._direction_tried[i]            
+            move[key] += self._direction_tried[i]
         self._project(move)
         config = self.denormalize(move)
         self._proposed_by[trial_id] = self.incumbent
         self._configs[trial_id] = config
-        return config
+        return unflatten_dict(config)
 
     def _project(self, config):
         ''' project normalized config in the feasible region and set prune_attr
@@ -553,6 +558,7 @@ class FLOW2(Searcher):
     def config_signature(self, config) -> tuple:
         ''' return the signature tuple of a config
         '''
+        config = flatten_dict(config)
         value_list = []
         for key in self._space_keys:
             if key in config:

flaml/searcher/search_thread.py

@@ -20,12 +20,14 @@ class SearchThread:
     '''
 
     cost_attr = 'time_total_s'
+    eps = 1e-10
 
     def __init__(self, mode: str = "min", 
                  search_alg: Optional[Searcher] = None):
         ''' When search_alg is omitted, use local search FLOW2
         '''
         self._search_alg = search_alg
+        self._is_ls = isinstance(search_alg, FLOW2)
         self._mode = mode
         self._metric_op = 1 if mode=='min' else -1
         self.cost_best = self.cost_last = self.cost_total = self.cost_best1 = \
@@ -36,6 +38,7 @@ class SearchThread:
         # eci: expected cost for improvement
         self.eci = self.cost_best
         self.priority = self.speed = 0
+        self._init_config = True 
         
     def suggest(self, trial_id: str) -> Optional[Dict]:
         ''' use the suggest() of the underlying search algorithm
@@ -70,7 +73,7 @@ class SearchThread:
         # calculate speed; use 0 for invalid speed temporarily
         if self.obj_best2 > self.obj_best1: 
             self.speed = (self.obj_best2 - self.obj_best1) / (
-                self.cost_total - self.cost_best2)
+                self.cost_total - self.cost_best2 + self.eps)
         else: self.speed = 0
 
     def on_trial_complete(self, trial_id: str, result: Optional[Dict] = None,
@@ -81,7 +84,12 @@ class SearchThread:
         if not hasattr(self._search_alg, '_ot_trials') or (not error and
             trial_id in self._search_alg._ot_trials):
             # optuna doesn't handle error
-            self._search_alg.on_trial_complete(trial_id, result, error)
+            if self._is_ls or not self._init_config:
+                self._search_alg.on_trial_complete(trial_id, result, error)
+            else: 
+                # init config is not proposed by self._search_alg
+                # under this thread
+                self._init_config = False
         if result:
             if self.cost_attr in result:
                 self.cost_last = result[self.cost_attr]

flaml/searcher/variant_generator.py

@@ -28,6 +28,46 @@ from ..tune.sample import Categorical, Domain, Function
 logger = logging.getLogger(__name__)
 
 
+def flatten_dict(dt, delimiter="/", prevent_delimiter=False):
+    dt = copy.deepcopy(dt)
+    if prevent_delimiter and any(delimiter in key for key in dt):
+        # Raise if delimiter is any of the keys
+        raise ValueError(
+            "Found delimiter `{}` in key when trying to flatten array."
+            "Please avoid using the delimiter in your specification.")
+    while any(isinstance(v, dict) for v in dt.values()):
+        remove = []
+        add = {}
+        for key, value in dt.items():
+            if isinstance(value, dict):
+                for subkey, v in value.items():
+                    if prevent_delimiter and delimiter in subkey:
+                        # Raise  if delimiter is in any of the subkeys
+                        raise ValueError(
+                            "Found delimiter `{}` in key when trying to "
+                            "flatten array. Please avoid using the delimiter "
+                            "in your specification.")
+                    add[delimiter.join([key, str(subkey)])] = v
+                remove.append(key)
+        dt.update(add)
+        for k in remove:
+            del dt[k]
+    return dt
+
+
+def unflatten_dict(dt, delimiter="/"):
+    """Unflatten dict. Does not support unflattening lists."""
+    dict_type = type(dt)
+    out = dict_type()
+    for key, val in dt.items():
+        path = key.split(delimiter)
+        item = out
+        for k in path[:-1]:
+            item = item.setdefault(k, dict_type())
+        item[path[-1]] = val
+    return out
+    
+    
 class TuneError(Exception):
     """General error class raised by ray.tune."""
     pass

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:

flaml/tune/README.md

@@ -159,6 +159,12 @@ Recommended scenario: cost-related hyperparameters exist, a low-cost
 initial point is known, and the search space is complex such that local search
 is prone to be stuck at local optima.
 
+An example of using BlendSearch with NNI can be seen in [test](https://github.com/microsoft/FLAML/tree/main/test/nni), CFO can be used with NNI as well in a similar manner. To run the example, first make sure you have [NNI](https://nni.readthedocs.io/en/stable/) installed, then run:
+
+```shell
+$nnictl create --config ./config.yml
+```
+
 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.
@@ -172,7 +178,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,

flaml/tune/tune.py

@@ -17,6 +17,8 @@ logger = logging.getLogger(__name__)
 _use_ray = True
 _runner = None
 _verbose = 0
+_running_trial = None
+_training_iteration = 0
 
 
 class ExperimentAnalysis(EA):
@@ -68,6 +70,8 @@ def report(_metric=None, **kwargs):
     '''
     global _use_ray
     global _verbose
+    global _running_trial
+    global _training_iteration
     if _use_ray:
         from ray import tune
         return tune.report(_metric, **kwargs)
@@ -77,6 +81,12 @@ def report(_metric=None, **kwargs):
             logger.info(f"result: {kwargs}")
         if _metric: result['_default_anonymous_metric'] = _metric
         trial = _runner.running_trial
+        if _running_trial == trial:
+            _training_iteration += 1
+        else:
+            _training_iteration = 0
+            _running_trial = trial
+        result["training_iteration"] = _training_iteration
         result['config'] = trial.config
         for key, value in trial.config.items():
             result['config/'+key] = value
@@ -195,7 +205,7 @@ def run(training_function,
             0 = silent, 1 = only status updates, 2 = status and brief trial
             results, 3 = status and detailed trial results. Defaults to 2.
         local_dir: A string of the local dir to save ray logs if ray backend is
-            used.
+            used; or a local dir to save the tuning log.
         num_samples: An integer of the number of configs to try. Defaults to 1.
         resources_per_trial: A dictionary of the hardware resources to allocate
             per trial, e.g., `{'mem': 1024**3}`. When not using ray backend,
@@ -211,9 +221,18 @@ def run(training_function,
         _verbose = verbose
         if verbose > 0:
             import os
-            os.makedirs(local_dir, exist_ok=True)
-            logger.addHandler(logging.FileHandler(local_dir+'/tune_'+str(
-                datetime.datetime.now())+'.log'))
+            if local_dir:
+                os.makedirs(local_dir, exist_ok=True)
+                logger.addHandler(logging.FileHandler(local_dir+'/tune_'+str(
+                    datetime.datetime.now()).replace(':', '-')+'.log'))
+            elif not logger.handlers:
+                # Add the console handler.
+                _ch = logging.StreamHandler()
+                logger_formatter = logging.Formatter(
+                    '[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s',
+                    '%m-%d %H:%M:%S')
+                _ch.setFormatter(logger_formatter)
+                logger.addHandler(_ch)        
             if verbose<=2:
                 logger.setLevel(logging.INFO)
             else:

flaml/version.py

@@ -1 +1 @@
-__version__ = "0.2.5"
+__version__ = "0.2.9"

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)"
    ]
   },
   {

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)"
    ]
   },
   {

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)"
    ]
   },
   {

setup.py

@@ -53,12 +53,15 @@ setuptools.setup(
             "optuna==2.3.0"
         ],
         "ray": [
-            "ray[tune]==1.1.0",
+            "ray[tune]==1.2.0",
             "pyyaml<5.3.1",            
         ],
         "azureml": [
-            "azureml-mlflow"
+            "azureml-mlflow",
         ],
+        "nni": [
+            "nni",
+        ]
     },
     classifiers=[
         "Programming Language :: Python :: 3",

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()

test/hf/test_distillbert.py

@@ -15,40 +15,33 @@ try:
         Trainer,
         TrainingArguments,
     )
-except:
-    print("pip install torch transformers datasets flaml[blendsearch,ray]")
-    
-import logging
-logger = logging.getLogger(__name__)
-logger.addHandler(logging.FileHandler('test/tune_distilbert.log'))
-logger.setLevel(logging.INFO)
+    MODEL_CHECKPOINT = "distilbert-base-uncased"
+    TASK = "cola"
+    NUM_LABELS = 2
+    COLUMN_NAME = "sentence"
+    METRIC_NAME = "matthews_correlation"
 
-import flaml
-
-
-MODEL_CHECKPOINT = "distilbert-base-uncased"
-TASK = "cola"
-NUM_LABELS = 2
-COLUMN_NAME = "sentence"
-METRIC_NAME = "matthews_correlation"
-
-# HP_METRIC, MODE = "loss", "min"
-HP_METRIC, MODE = "matthews_correlation", "max"
-
-def train_distilbert(config: dict):
+    # HP_METRIC, MODE = "loss", "min"
+    HP_METRIC, MODE = "matthews_correlation", "max"
 
     # Define tokenize method
     tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, use_fast=True)
     def tokenize(examples):
         return tokenizer(examples[COLUMN_NAME], truncation=True)
-    # Load CoLA dataset and apply tokenizer
-    cola_raw = load_dataset("glue", TASK)
-    cola_encoded = cola_raw.map(tokenize, batched=True)
-    train_dataset, eval_dataset = cola_encoded["train"], cola_encoded["validation"]
 
-    model = AutoModelForSequenceClassification.from_pretrained(
-        MODEL_CHECKPOINT, num_labels=NUM_LABELS
-    )
+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_distilbert.log'))
+logger.setLevel(logging.INFO)
+
+import flaml
+
+def train_distilbert(config: dict):
 
     metric = load_metric("glue", TASK)
 
@@ -57,6 +50,16 @@ def train_distilbert(config: dict):
         predictions = np.argmax(predictions, axis=1)
         return metric.compute(predictions=predictions, references=labels)
 
+    # Load CoLA dataset and apply tokenizer
+    cola_raw = load_dataset("glue", TASK)
+
+    cola_encoded = cola_raw.map(tokenize, batched=True)
+    train_dataset, eval_dataset = cola_encoded["train"], cola_encoded["validation"]
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_CHECKPOINT, num_labels=NUM_LABELS
+    )
+
     training_args = TrainingArguments(
         output_dir='.',
         do_eval=False,
@@ -91,7 +94,7 @@ def _test_distillbert(method='BlendSearch'):
  
     max_num_epoch = 64
     num_samples = -1
-    time_budget_s = 10800
+    time_budget_s = 3600
 
     search_space = {
         # You can mix constants with search space objects.
@@ -123,7 +126,7 @@ def _test_distillbert(method='BlendSearch'):
         from flaml import BlendSearch
         algo = BlendSearch(points_to_evaluate=[{
             "num_train_epochs": 1,
-        }])        
+        }])
     elif 'Dragonfly' == method:
         from ray.tune.suggest.dragonfly import DragonflySearch
         algo = DragonflySearch()
@@ -139,7 +142,7 @@ def _test_distillbert(method='BlendSearch'):
         algo = ZOOptSearch(budget=num_samples)
     elif 'Ax' == method:
         from ray.tune.suggest.ax import AxSearch
-        algo = AxSearch()
+        algo = AxSearch(max_concurrent=3)
     elif 'HyperOpt' == method:
         from ray.tune.suggest.hyperopt import HyperOptSearch
         algo = HyperOptSearch()
@@ -154,9 +157,8 @@ def _test_distillbert(method='BlendSearch'):
         train_distilbert,
         metric=HP_METRIC,
         mode=MODE,
-        # You can add "gpu": 1 to allocate GPUs
-        resources_per_trial={"gpu": 1},
-        config=search_space, local_dir='test/logs/',
+        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)
@@ -214,4 +216,4 @@ def _test_distillbert_bohb():
 
 
 if __name__ == "__main__":
-    _test_distillbert()
+    _test_distillbert()

test/hf/test_electra.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 = "google/electra-base-discriminator"
+    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_electra.log'))
+logger.setLevel(logging.INFO)
+
+import flaml
+
+def train_electra(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_electra(method='BlendSearch'):
+ 
+    max_num_epoch = 9
+    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_electra,
+        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_electra_cfo():
+    _test_electra('CFO')
+
+
+def _test_electra_dragonfly():
+    _test_electra('Dragonfly')
+
+
+def _test_electra_skopt():
+    _test_electra('SkOpt')
+
+
+def _test_electra_nevergrad():
+    _test_electra('Nevergrad')
+
+
+def _test_electra_zoopt():
+    _test_electra('ZOOpt')
+
+
+def _test_electra_ax():
+    _test_electra('Ax')
+
+
+def __test_electra_hyperopt():
+    _test_electra('HyperOpt')
+
+
+def _test_electra_optuna():
+    _test_electra('Optuna')
+
+
+def _test_electra_asha():
+    _test_electra('ASHA')
+
+
+def _test_electra_bohb():
+    _test_electra('BOHB')
+
+
+if __name__ == "__main__":
+    _test_electra()

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()
+

test/nni/config.yml

@@ -0,0 +1,19 @@
+# usage: nnictl create --config ./config.yml
+authorName: default
+experimentName: example_mnist
+trialConcurrency: 1
+maxExecDuration: 1h
+maxTrialNum: 10
+trainingServicePlatform: local
+# The path to Search Space
+searchSpacePath: search_space.json
+useAnnotation: false
+tuner:
+  codeDir: ./
+  classFileName: flaml_nni_wrap.py
+  className: BlendSearchTuner
+# The path and the running command of trial
+trial:
+  command: python3 mnist.py
+  codeDir: .
+  gpuNum: 0

test/nni/flaml_nni_wrap.py

@@ -0,0 +1,7 @@
+from flaml.searcher.blendsearch import BlendSearchTuner as BST
+
+
+class BlendSearchTuner(BST):
+    # for best performance pass low cost initial parameters here 
+    def __init__(self, points_to_evaluate=[{"hidden_size":128}]):
+        super.__init__(self, points_to_evaluate=points_to_evaluate)

test/nni/mnist.py

@@ -0,0 +1,169 @@
+# This file is copied from NNI project
+# https://github.com/microsoft/nni/blob/master/examples/trials/mnist-tfv1/mnist.py
+
+"""
+A deep MNIST classifier using convolutional layers.
+
+This file is a modification of the official pytorch mnist example:
+https://github.com/pytorch/examples/blob/master/mnist/main.py
+"""
+
+import os
+import argparse
+import logging
+import nni
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from nni.utils import merge_parameter
+from torchvision import datasets, transforms
+
+logger = logging.getLogger('mnist_AutoML')
+
+
+class Net(nn.Module):
+    def __init__(self, hidden_size):
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(1, 20, 5, 1)
+        self.conv2 = nn.Conv2d(20, 50, 5, 1)
+        self.fc1 = nn.Linear(4*4*50, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, 10)
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.max_pool2d(x, 2, 2)
+        x = F.relu(self.conv2(x))
+        x = F.max_pool2d(x, 2, 2)
+        x = x.view(-1, 4*4*50)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+def train(args, model, device, train_loader, optimizer, epoch):
+    model.train()
+    for batch_idx, (data, target) in enumerate(train_loader):
+        if (args['batch_num'] is not None) and batch_idx >= args['batch_num']:
+            break
+        data, target = data.to(device), target.to(device)
+        optimizer.zero_grad()
+        output = model(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % args['log_interval'] == 0:
+            logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(train_loader.dataset),
+                100. * batch_idx / len(train_loader), loss.item()))
+
+
+def test(args, model, device, test_loader):
+    model.eval()
+    test_loss = 0
+    correct = 0
+    with torch.no_grad():
+        for data, target in test_loader:
+            data, target = data.to(device), target.to(device)
+            output = model(data)
+            # sum up batch loss
+            test_loss += F.nll_loss(output, target, reduction='sum').item()
+            # get the index of the max log-probability
+            pred = output.argmax(dim=1, keepdim=True)
+            correct += pred.eq(target.view_as(pred)).sum().item()
+
+    test_loss /= len(test_loader.dataset)
+
+    accuracy = 100. * correct / len(test_loader.dataset)
+
+    logger.info('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_loss, correct, len(test_loader.dataset), accuracy))
+
+    return accuracy
+
+
+def main(args):
+    use_cuda = not args['no_cuda'] and torch.cuda.is_available()
+
+    torch.manual_seed(args['seed'])
+
+    device = torch.device("cuda" if use_cuda else "cpu")
+
+    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
+
+    data_dir = args['data_dir']
+
+    train_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(data_dir, train=True, download=True,
+                       transform=transforms.Compose([
+                           transforms.ToTensor(),
+                           transforms.Normalize((0.1307,), (0.3081,))
+                       ])),
+        batch_size=args['batch_size'], shuffle=True, **kwargs)
+    test_loader = torch.utils.data.DataLoader(
+        datasets.MNIST(data_dir, train=False, transform=transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize((0.1307,), (0.3081,))
+        ])),
+        batch_size=1000, shuffle=True, **kwargs)
+
+    hidden_size = args['hidden_size']
+
+    model = Net(hidden_size=hidden_size).to(device)
+    optimizer = optim.SGD(model.parameters(), lr=args['lr'],
+                          momentum=args['momentum'])
+
+    for epoch in range(1, args['epochs'] + 1):
+        train(args, model, device, train_loader, optimizer, epoch)
+        test_acc = test(args, model, device, test_loader)
+
+        # report intermediate result
+        nni.report_intermediate_result(test_acc)
+        logger.debug('test accuracy %g', test_acc)
+        logger.debug('Pipe send intermediate result done.')
+
+    # report final result
+    nni.report_final_result(test_acc)
+    logger.debug('Final result is %g', test_acc)
+    logger.debug('Send final result done.')
+
+
+def get_params():
+    # Training settings
+    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+    parser.add_argument("--data_dir", type=str,
+                        default='./data', help="data directory")
+    parser.add_argument('--batch_size', type=int, default=64, metavar='N',
+                        help='input batch size for training (default: 64)')
+    parser.add_argument("--batch_num", type=int, default=None)
+    parser.add_argument("--hidden_size", type=int, default=512, metavar='N',
+                        help='hidden layer size (default: 512)')
+    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
+                        help='learning rate (default: 0.01)')
+    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
+                        help='SGD momentum (default: 0.5)')
+    parser.add_argument('--epochs', type=int, default=10, metavar='N',
+                        help='number of epochs to train (default: 10)')
+    parser.add_argument('--seed', type=int, default=1, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--no_cuda', action='store_true', default=False,
+                        help='disables CUDA training')
+    parser.add_argument('--log_interval', type=int, default=1000, metavar='N',
+                        help='how many batches to wait before logging training status')
+
+
+    args, _ = parser.parse_known_args()
+    return args
+
+
+if __name__ == '__main__':
+    try:
+        # get parameters form tuner
+        tuner_params = nni.get_next_parameter()
+        logger.debug(tuner_params)
+        params = vars(merge_parameter(get_params(), tuner_params))
+        print(params)
+        main(params)
+    except Exception as exception:
+        logger.exception(exception)
+        raise

test/nni/search_space.json

@@ -0,0 +1,6 @@
+{
+    "batch_size": {"_type":"choice", "_value": [16, 32, 64, 128]},
+    "hidden_size":{"_type":"choice","_value":[128, 256, 512, 1024]},
+    "lr":{"_type":"choice","_value":[0.0001, 0.001, 0.01, 0.1]},
+    "momentum":{"_type":"uniform","_value":[0, 1]}
+}

test/run_electra.py

@@ -0,0 +1,19 @@
+from azureml.core import Workspace, Experiment, ScriptRunConfig
+ws = Workspace.from_config()
+
+compute_target = ws.compute_targets['V100-4']
+# compute_target = ws.compute_targets['K80']
+command = [
+    "pip install torch transformers datasets flaml[blendsearch,ray] ax-platform sqlalchemy && ",
+    "python test_electra.py"]
+
+config = ScriptRunConfig(
+    source_directory='hf/',
+    command=command,
+    compute_target=compute_target,
+)
+
+exp = Experiment(ws, 'test-electra')
+run = exp.submit(config)
+print(run.get_portal_url()) # link to ml.azure.com
+run.wait_for_completion(show_output=True)

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()
@@ -170,6 +172,10 @@ class TestAutoML(unittest.TestCase):
             "model_history": True
         }
         X_train, y_train = load_iris(return_X_y=True, as_frame=as_frame)
+        if as_frame:
+            # test drop column
+            X_train.columns = range(X_train.shape[1])
+            X_train[X_train.shape[1]] = np.zeros(len(y_train))
         automl_experiment.fit(X_train=X_train, y_train=y_train,
                               **automl_settings)
         print(automl_experiment.classes_)
@@ -250,7 +256,8 @@ class TestAutoML(unittest.TestCase):
             "task": 'regression',
             "log_file_name": "test/sparse_regression.log",
             "n_jobs": 1,
-            "model_history": True
+            "model_history": True,
+            "verbose": 0,
         }
         X_train = scipy.sparse.random(300, 900, density=0.0001)
         y_train = np.random.uniform(size=300)
@@ -325,10 +332,11 @@ class TestAutoML(unittest.TestCase):
             "task": 'regression',
             "log_file_name": "test/sparse_regression.log",
             "n_jobs": 1,
-            "model_history": True
+            "model_history": True,
+            "metric": "mse"
         }
-        X_train = scipy.sparse.random(100, 100)
-        y_train = np.random.uniform(size=100)
+        X_train = scipy.sparse.random(8, 100)
+        y_train = np.random.uniform(size=8)
         automl_experiment.fit(X_train=X_train, y_train=y_train,
                               **automl_settings)
         print(automl_experiment.predict(X_train))

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__)

test/test_training_log.py

@@ -25,7 +25,8 @@ class TestTrainingLog(unittest.TestCase):
                 "log_training_metric": True,
                 "mem_thres": 1024*1024,
                 "n_jobs": 1,
-                "model_history": True
+                "model_history": True,
+                "verbose": 2,
             }
             X_train, y_train = load_boston(return_X_y=True)            
             automl_experiment.fit(X_train=X_train, y_train=y_train,

test/test_tune.py

@@ -49,7 +49,6 @@ def _test_xgboost(method='BlendSearch'):
     else:
         from ray import tune
     search_space = {
-        # You can mix constants with search space objects.
         "max_depth": tune.randint(1, 8) if method in [
             "BlendSearch", "BOHB", "Optuna"] else tune.randint(1, 9),
         "min_child_weight": tune.choice([1, 2, 3]),
@@ -57,7 +56,7 @@ def _test_xgboost(method='BlendSearch'):
         "eta": tune.loguniform(1e-4, 1e-1)
     }
     max_iter = 10
-    for num_samples in [256]:
+    for num_samples in [128]:
         time_budget_s = 60 #None
         for n_cpu in [8]:
             start_time = time.time()
@@ -154,6 +153,33 @@ def _test_xgboost(method='BlendSearch'):
             logger.info(f"Best model parameters: {best_trial.config}")
 
 
+def test_nested():
+    from flaml import tune
+    search_space = {
+        # test nested search space
+        "cost_related": {
+            "a": tune.randint(1, 8),
+        },
+        "b": tune.uniform(0.5, 1.0),
+    }
+
+    def simple_func(config):
+        tune.report(
+            metric=(config["cost_related"]["a"]-4)**2 * (config["b"]-0.7)**2)
+
+    analysis = tune.run(
+        simple_func,
+        init_config={
+            "cost_related": {"a": 1,}
+        },
+        metric="metric",
+        mode="min",
+        config=search_space,
+        local_dir='logs/',
+        num_samples=-1,
+        time_budget_s=1)
+
+
 def test_xgboost_bs():
     _test_xgboost()