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>

.github/workflows/python-package.yml

@@ -1,7 +1,7 @@
 # This workflow will install Python dependencies, run tests and lint with a variety of Python versions
 # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
 
-name: Python package
+name: Build
 
 on:
   push:

.gitignore

@@ -146,6 +146,10 @@ dmypy.json
 
 # Cython debug symbols
 cython_debug/
-/catboost_info
+
+catboost_info
 notebook/*.pkl
 notebook/.azureml
+mlruns
+logs
+automl.pkl

README.md

@@ -1,3 +1,9 @@
+[![PyPI version](https://badge.fury.io/py/FLAML.svg)](https://badge.fury.io/py/FLAML)
+[![Build](https://github.com/microsoft/FLAML/actions/workflows/python-package.yml/badge.svg)](https://github.com/microsoft/FLAML/actions/workflows/python-package.yml)
+![Python Version](https://img.shields.io/badge/3.6%20%7C%203.7%20%7C%203.8-blue)
+[![Downloads](https://pepy.tech/badge/flaml/month)](https://pepy.tech/project/flaml)
+[![Join the chat at https://gitter.im/FLAMLer/community](https://badges.gitter.im/FLAMLer/community.svg)](https://gitter.im/FLAMLer/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge)
+
 # FLAML - Fast and Lightweight AutoML
 
 <p align="center">
@@ -5,13 +11,16 @@
     <br>
 </p>
 
-FLAML is a Python library designed to automatically produce accurate machine
-learning models with low computational cost. It frees users from selecting
-learners and hyperparameters for each learner. It is fast and cheap.
+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 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
@@ -112,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},
@@ -122,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
 
@@ -130,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.
@@ -138,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,5 +1,5 @@
-from flaml.searcher import CFO, BlendSearch, FLOW2
-from flaml.automl import AutoML
+from flaml.searcher import CFO, BlendSearch, FLOW2, BlendSearchTuner
+from flaml.automl import AutoML, logger_formatter
 from flaml.version import __version__
 import logging
 
@@ -7,10 +7,3 @@ import logging
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
-# 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)

flaml/automl.py

@@ -25,6 +25,10 @@ from .training_log import training_log_reader, training_log_writer
 
 import logging
 logger = logging.getLogger(__name__)
+logger_formatter = logging.Formatter(
+    '[%(name)s: %(asctime)s] {%(lineno)d} %(levelname)s - %(message)s',
+    '%m-%d %H:%M:%S')
+
 try:
     import mlflow
 except:
@@ -249,6 +253,8 @@ class AutoML:
 
     '''
 
+    from .version import __version__
+
     def __init__(self):
         self._track_iter = 0
         self._state = AutoMLState()
@@ -279,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.'''        
@@ -326,6 +348,10 @@ class AutoML:
             A numpy array of shape n * 1 - - each element is a predicted class
             label for an instance.
         '''
+        if self._trained_estimator is None:
+            warnings.warn(
+                "No estimator is trained. Please run fit with enough budget.")
+            return None
         X_test = self._preprocess(X_test)
         y_pred = self._trained_estimator.predict(X_test)
         if y_pred.ndim > 1: y_pred = y_pred.flatten()
@@ -402,7 +428,7 @@ class AutoML:
             self._X_train_all, self._y_train_all = \
                 self._transformer.fit_transform(X, y, self._state.task)
             self._label_transformer = self._transformer.label_transformer
-
+        self._sample_weight_full = self._state.fit_kwargs.get('sample_weight')      
         if X_val is not None and y_val is not None:
             if not (isinstance(X_val, np.ndarray) or
                 issparse(X_val) or
@@ -446,7 +472,8 @@ class AutoML:
             self._X_train_all, self._y_train_all
         if issparse(X_train_all): 
             X_train_all = X_train_all.tocsr()
-        if self._state.task != 'regression':
+        if self._state.task != 'regression' and self._state.fit_kwargs.get(
+            'sample_weight') is None:
             # logger.info(f"label {pd.unique(y_train_all)}")
             label_set, counts = np.unique(y_train_all, return_counts=True)
             # augment rare classes
@@ -562,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:
@@ -756,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
 
@@ -814,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
         '''
@@ -826,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)))
@@ -836,6 +876,11 @@ 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.handlers:
+            # Add the console handler.
+            _ch = logging.StreamHandler()
+            _ch.setFormatter(logger_formatter)
+            logger.addHandler(_ch)        
         logger.info("Evaluation method: {}".format(eval_method))
         
         self._retrain_full = retrain_full and (eval_method == 'holdout' and
@@ -892,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
@@ -1004,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()
@@ -1060,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])
@@ -1093,8 +1140,9 @@ class AutoML:
                             self._state.best_loss))
             else:
                 logger.info(f"no enough budget for learner {estimator}")
-                self.estimator_list.remove(estimator)
-                self._estimator_index -= 1
+                if self._estimator_index is not None:
+                    self.estimator_list.remove(estimator)
+                    self._estimator_index -= 1
             if self._retrain_full and best_config_sig and not better and (
                 self._search_states[self._best_estimator].sample_size ==
                 self._state.data_size) and (est_retrain_time <=
@@ -1151,7 +1199,11 @@ class AutoML:
                 stacker = Stacker(estimators, best_m,
                     n_jobs=self._state.n_jobs,
                     passthrough=True)
-                stacker.fit(self._X_train_all, self._y_train_all)
+                if self._sample_weight_full is not None:
+                    self._state.fit_kwargs[
+                        'sample_weight'] = self._sample_weight_full
+                stacker.fit(self._X_train_all, self._y_train_all,
+                 **self._state.fit_kwargs)
                 logger.info(f'ensemble: {stacker}')
                 self._trained_estimator = stacker
                 self._trained_estimator.model = stacker
@@ -1189,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,
@@ -113,8 +116,9 @@ class BlendSearch(Searcher):
                 self._deadline = config.get('time_budget_s') + time.time()
             if 'metric_target' in config:
                 self._metric_target = config.get('metric_target')
-        else:    
-            self._metric, self._mode = metric, mode
+        else:
+            if metric: self._metric = metric
+            if mode: self._mode = mode
             self._ls.set_search_properties(metric, mode, config)
             if self._gs is not None:
                 self._gs.set_search_properties(metric, mode, config)
@@ -132,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)
@@ -178,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}")
@@ -204,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
         '''
@@ -231,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]
 
@@ -258,53 +280,67 @@ 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
-            if init_config==self._ls.init_config:
-                config = self._ls.complete_config(init_config,
-             self._admissible_min, self._admissible_max)
+            config = self._ls.complete_config(init_config,
+             self._ls_bound_min, self._ls_bound_max)
                 # logger.info(f"reset config to {config}")
-            else: config = init_config
             config_signature = self._ls.config_signature(config)
             result = self._result.get(config_signature)
             if result: # tried before
@@ -340,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
 
@@ -364,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}")
@@ -384,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
     '''
 
@@ -418,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
@@ -121,21 +124,21 @@ class FLOW2(Searcher):
         self._unordered_cat_hp = {}
         self._cat_hp_cost = {}
         for key, domain in self.space.items():
-            assert not isinstance(domain, dict), \
-                key+"'s domain is grid search which is not supported in FLOW2."
+            assert not (isinstance(domain, dict) and 'grid_search' in domain
+            ), key+"'s domain is grid search which is not supported in FLOW2."
             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        
@@ -190,6 +193,8 @@ class FLOW2(Searcher):
         self._K = 0
         self._iter_best_config = self.trial_count = 1
         self._reset_times = 0
+        # record intermediate trial cost
+        self._trial_cost = {}
 
     @property
     def step_lower_bound(self) -> float:
@@ -237,14 +242,15 @@ class FLOW2(Searcher):
         ''' generate a complete config from the partial config input
         add minimal resource to config if available
         '''
-        if self._reset_times: # not the first time, use random gaussian
+        if self._reset_times and partial_config==self.init_config:
+            # not the first time to complete init_config, use random gaussian
             normalized = self.normalize(partial_config)
             for key in normalized:
                  # don't change unordered cat choice
                 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
@@ -259,27 +265,28 @@ class FLOW2(Searcher):
             # use best config for unordered cat choice
             config = self.denormalize(normalized)
         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
-        logger.debug(f'before random {config}')
+        # logger.debug(f'before random {config}')
         for _, generated in generate_variants({'config': config}):
             config = generated['config']
             break
-        logger.debug(f'after random {config}')
+        # logger.debug(f'after random {config}')
 
         if self._resource:
             config[self.prune_attr] = self.min_resource
-        self._reset_times += 1
-        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
@@ -288,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]
@@ -299,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[
@@ -402,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":
@@ -422,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)
@@ -433,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
@@ -442,7 +451,8 @@ class FLOW2(Searcher):
         if proposed_by == self.incumbent:
             # proposed by current incumbent and no better
             self._num_complete4incumbent += 1
-            cost = result.get(self.cost_attr)
+            cost = result.get(
+                self.cost_attr) if result else self._trial_cost.get(trial_id)
             if cost: self._cost_complete4incumbent += cost
             if self._num_complete4incumbent >= 2*self.dim and \
                 self._num_allowed4incumbent == 0:
@@ -469,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:
@@ -483,6 +493,9 @@ class FLOW2(Searcher):
                         self._num_allowed4incumbent = 2 * self.dim
                         self._proposed_by.clear()
                         self._iter_best_config = self.trial_count
+            cost = result.get(self.cost_attr)
+            # record the cost in case it is pruned and cost info is lost
+            self._trial_cost[trial_id] = cost
 
     def rand_vector_unit_sphere(self, dim) -> np.ndarray:
         vec = self._random.normal(0, 1, dim)
@@ -520,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
@@ -545,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

@@ -6,6 +6,7 @@ The API is compatible with ray tune.
 Example:
 
 ```python
+# require: pip install flaml[blendsearch]
 from flaml import tune
 import time
 
@@ -42,6 +43,7 @@ print(analysis.best_config) # the best config
 
 Or, using ray tune's API:
 ```python
+# require: pip install flaml[blendsearch] ray[tune]
 from ray import tune as raytune
 from flaml import CFO, BlendSearch
 import time
@@ -146,6 +148,7 @@ based on optimism in face of uncertainty.
 Example:
 
 ```python
+# require: pip install flaml[blendsearch]
 from flaml import BlendSearch
 tune.run(...
     search_alg = BlendSearch(points_to_evaluate=[init_config]),
@@ -156,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.
@@ -169,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.3"
+__version__ = "0.2.9"

notebook/finetune_transformer_demo.ipynb

@@ -1,788 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This notebook uses the Huggingface transformers library to finetune a transformer model.\n",
-    "\n",
-    "**Requirements.** This notebook has additional requirements:\n",
-    "\n",
-    "```bash\n",
-    "pip install -r transformers_requirements.txt\n",
-    "```"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Tokenizer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import AutoTokenizer"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "MODEL_CHECKPOINT = \"distilbert-base-uncased\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, use_fast=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'input_ids': [101, 2023, 2003, 1037, 3231, 102], 'attention_mask': [1, 1, 1, 1, 1, 1]}"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "tokenizer(\"this is a test\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Data"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "TASK = \"cola\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import datasets"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Reusing dataset glue (/home/amin/.cache/huggingface/datasets/glue/cola/1.0.0/7c99657241149a24692c402a5c3f34d4c9f1df5ac2e4c3759fadea38f6cb29c4)\n",
-      "/home/amin/miniconda/lib/python3.7/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)\n",
-      "  return torch._C._cuda_getDeviceCount() > 0\n"
-     ]
-    }
-   ],
-   "source": [
-    "raw_dataset = datasets.load_dataset(\"glue\", TASK)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# define tokenization function used to process data\n",
-    "COLUMN_NAME = \"sentence\"\n",
-    "def tokenize(examples):\n",
-    "    return tokenizer(examples[COLUMN_NAME], truncation=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5bd7b23a478043eaaf6e14e119143fcd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, max=9.0), HTML(value='')))"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d7b648c2dbdc4fb9907e43da7db8af9a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "36a9d6e62dbe462d94b1769f36fbd0f3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, max=2.0), HTML(value='')))"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "encoded_dataset = raw_dataset.map(tokenize, batched=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
-       " 'idx': 0,\n",
-       " 'input_ids': [101,\n",
-       "  2256,\n",
-       "  2814,\n",
-       "  2180,\n",
-       "  1005,\n",
-       "  1056,\n",
-       "  4965,\n",
-       "  2023,\n",
-       "  4106,\n",
-       "  1010,\n",
-       "  2292,\n",
-       "  2894,\n",
-       "  1996,\n",
-       "  2279,\n",
-       "  2028,\n",
-       "  2057,\n",
-       "  16599,\n",
-       "  1012,\n",
-       "  102],\n",
-       " 'label': 1,\n",
-       " 'sentence': \"Our friends won't buy this analysis, let alone the next one we propose.\"}"
-      ]
-     },
-     "execution_count": 19,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "encoded_dataset[\"train\"][0]"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import AutoModelForSequenceClassification"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "35b76e51b5c8406fae416fcdc3dd885e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=267967963.0, style=ProgressStyle(descri…"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertForSequenceClassification: ['vocab_transform.weight', 'vocab_transform.bias', 'vocab_layer_norm.weight', 'vocab_layer_norm.bias', 'vocab_projector.weight', 'vocab_projector.bias']\n",
-      "- This IS expected if you are initializing DistilBertForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
-      "- This IS NOT expected if you are initializing DistilBertForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
-      "Some weights of DistilBertForSequenceClassification were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['pre_classifier.weight', 'pre_classifier.bias', 'classifier.weight', 'classifier.bias']\n",
-      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
-     ]
-    }
-   ],
-   "source": [
-    "NUM_LABELS = 2\n",
-    "model = AutoModelForSequenceClassification.from_pretrained(MODEL_CHECKPOINT, num_labels=NUM_LABELS)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 31,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "DistilBertForSequenceClassification(\n",
-       "  (distilbert): DistilBertModel(\n",
-       "    (embeddings): Embeddings(\n",
-       "      (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
-       "      (position_embeddings): Embedding(512, 768)\n",
-       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "      (dropout): Dropout(p=0.1, inplace=False)\n",
-       "    )\n",
-       "    (transformer): Transformer(\n",
-       "      (layer): ModuleList(\n",
-       "        (0): TransformerBlock(\n",
-       "          (attention): MultiHeadSelfAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "          (ffn): FFN(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
-       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "        )\n",
-       "        (1): TransformerBlock(\n",
-       "          (attention): MultiHeadSelfAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "          (ffn): FFN(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
-       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "        )\n",
-       "        (2): TransformerBlock(\n",
-       "          (attention): MultiHeadSelfAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "          (ffn): FFN(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
-       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "        )\n",
-       "        (3): TransformerBlock(\n",
-       "          (attention): MultiHeadSelfAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "          (ffn): FFN(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
-       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "        )\n",
-       "        (4): TransformerBlock(\n",
-       "          (attention): MultiHeadSelfAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "          (ffn): FFN(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
-       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "        )\n",
-       "        (5): TransformerBlock(\n",
-       "          (attention): MultiHeadSelfAttention(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (q_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (k_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (v_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "            (out_lin): Linear(in_features=768, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (sa_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "          (ffn): FFN(\n",
-       "            (dropout): Dropout(p=0.1, inplace=False)\n",
-       "            (lin1): Linear(in_features=768, out_features=3072, bias=True)\n",
-       "            (lin2): Linear(in_features=3072, out_features=768, bias=True)\n",
-       "          )\n",
-       "          (output_layer_norm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
-       "        )\n",
-       "      )\n",
-       "    )\n",
-       "  )\n",
-       "  (pre_classifier): Linear(in_features=768, out_features=768, bias=True)\n",
-       "  (classifier): Linear(in_features=768, out_features=2, bias=True)\n",
-       "  (dropout): Dropout(p=0.2, inplace=False)\n",
-       ")"
-      ]
-     },
-     "execution_count": 31,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Metric"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "metric = datasets.load_metric(\"glue\", TASK)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Metric(name: \"glue\", features: {'predictions': Value(dtype='int64', id=None), 'references': Value(dtype='int64', id=None)}, usage: \"\"\"\n",
-       "Compute GLUE evaluation metric associated to each GLUE dataset.\n",
-       "Args:\n",
-       "    predictions: list of translations to score.\n",
-       "        Each translation should be tokenized into a list of tokens.\n",
-       "    references: list of lists of references for each translation.\n",
-       "        Each reference should be tokenized into a list of tokens.\n",
-       "Returns: depending on the GLUE subset, one or several of:\n",
-       "    \"accuracy\": Accuracy\n",
-       "    \"f1\": F1\n",
-       "    \"pearson\": Pearson Correlation\n",
-       "    \"spearmanr\": Spearman Correlation\n",
-       "    \"matthews_correlation\": Matthew Correlation\n",
-       "\"\"\", stored examples: 0)"
-      ]
-     },
-     "execution_count": 25,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "metric"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 26,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def compute_metrics(eval_pred):\n",
-    "    predictions, labels = eval_pred\n",
-    "    predictions = np.argmax(predictions, axis=1)\n",
-    "    return metric.compute(predictions=predictions, references=labels)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Training (aka Finetuning)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import Trainer\n",
-    "from transformers import TrainingArguments"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 29,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "args = TrainingArguments(\n",
-    "    output_dir='output',\n",
-    "    do_eval=True,\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 30,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "trainer = Trainer(\n",
-    "    model=model,\n",
-    "    args=args,\n",
-    "    train_dataset=encoded_dataset[\"train\"],\n",
-    "    eval_dataset=encoded_dataset[\"validation\"],\n",
-    "    tokenizer=tokenizer,\n",
-    "    compute_metrics=compute_metrics,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "\n",
-       "    <div>\n",
-       "        <style>\n",
-       "            /* Turns off some styling */\n",
-       "            progress {\n",
-       "                /* gets rid of default border in Firefox and Opera. */\n",
-       "                border: none;\n",
-       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
-       "                background-size: auto;\n",
-       "            }\n",
-       "        </style>\n",
-       "      \n",
-       "      <progress value='322' max='3207' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
-       "      [ 322/3207 02:51 < 25:41, 1.87 it/s, Epoch 0.30/3]\n",
-       "    </div>\n",
-       "    <table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: left;\">\n",
-       "      <th>Step</th>\n",
-       "      <th>Training Loss</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "  </tbody>\n",
-       "</table><p>"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "trainer.train()"
-   ]
-  },
-  {
-   "source": [
-    "## Hyperparameter Optimization\n",
-    "\n",
-    "`flaml.tune` is a module for economical hyperparameter tuning. It frees users from manually tuning many hyperparameters for a software, such as machine learning training procedures. \n",
-    "The API is compatible with ray tune.\n",
-    "\n",
-    "### Step 1. Define training method\n",
-    "\n",
-    "We define a function `train_distilbert(config: dict)` that accepts a hyperparameter configuration dict `config`. The specific configs will be generated by flaml's search algorithm in a given search space.\n"
-   ],
-   "cell_type": "markdown",
-   "metadata": {}
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import flaml\n",
-    "\n",
-    "def train_distilbert(config: dict):\n",
-    "\n",
-    "    # Define tokenize method\n",
-    "    tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, use_fast=True)\n",
-    "    def tokenize(examples):\n",
-    "        return tokenizer(examples[COLUMN_NAME], truncation=True)\n",
-    "    # Load CoLA dataset and apply tokenizer\n",
-    "    cola_raw = load_dataset(\"glue\", TASK)\n",
-    "    cola_encoded = cola_raw.map(tokenize, batched=True)\n",
-    "    # QUESTION: Write processed data to disk?\n",
-    "    train_dataset, eval_dataset = cola_encoded[\"train\"], cola_encoded[\"validation\"]\n",
-    "\n",
-    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
-    "        MODEL_CHECKPOINT, num_labels=NUM_LABELS\n",
-    "    )\n",
-    "\n",
-    "    metric = load_metric(\"glue\", TASK)\n",
-    "\n",
-    "    training_args = TrainingArguments(\n",
-    "        output_dir='.',\n",
-    "        do_eval=False,\n",
-    "        **config,\n",
-    "    )\n",
-    "\n",
-    "    trainer = Trainer(\n",
-    "        model,\n",
-    "        training_args,\n",
-    "        train_dataset=train_dataset,\n",
-    "        eval_dataset=eval_dataset,\n",
-    "        tokenizer=tokenizer,\n",
-    "        compute_metrics=compute_metrics,\n",
-    "    )\n",
-    "\n",
-    "    # train model\n",
-    "    trainer.train()\n",
-    "\n",
-    "    # evaluate model\n",
-    "    eval_output = trainer.evaluate()\n",
-    "\n",
-    "    # report the metric to optimize\n",
-    "    flaml.tune.report(\n",
-    "        loss=eval_output[\"eval_loss\"],\n",
-    "        matthews_correlation=eval_output[\"eval_matthews_correlation\"],\n",
-    "        )"
-   ]
-  },
-  {
-   "source": [
-    "### Step 2. Define the search\n",
-    "\n",
-    "We are now ready to define our search. This includes:\n",
-    "\n",
-    "- The `search_space` for our hyperparameters\n",
-    "- The metric and the mode ('max' or 'min') for optimization\n",
-    "- The constraints (`n_cpus`, `n_gpus`, `num_samples`, and `time_budget_s`)"
-   ],
-   "cell_type": "markdown",
-   "metadata": {}
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "max_num_epoch = 4\n",
-    "search_space = {\n",
-    "        # You can mix constants with search space objects.\n",
-    "        \"num_train_epochs\": flaml.tune.loguniform(1, max_num_epoch),\n",
-    "        \"learning_rate\": flaml.tune.loguniform(1e-6, 1e-4),\n",
-    "        \"adam_epsilon\": flaml.tune.loguniform(1e-9, 1e-7),\n",
-    "    }"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# optimization objective\n",
-    "HP_METRIC, MODE = \"matthews_correlation\", \"max\"\n",
-    "\n",
-    "# resources\n",
-    "num_cpus = 2\n",
-    "num_gpus = 2\n",
-    "\n",
-    "# constraints\n",
-    "num_samples = -1    # number of trials, -1 means unlimited\n",
-    "time_budget_s = 3600    # time budget in seconds"
-   ]
-  },
-  {
-   "source": [
-    "### Step 3. Launch with `flaml.tune.run`\n",
-    "\n",
-    "We are now ready to laungh the tuning using `flaml.tune.run`:"
-   ],
-   "cell_type": "markdown",
-   "metadata": {}
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import time\n",
-    "import ray\n",
-    "start_time = time.time()\n",
-    "ray.init(num_cpus=num_cpus, num_gpus=num_gpus)\n",
-    "\n",
-    "print(\"Tuning started...\")\n",
-    "analysis = flaml.tune.run(\n",
-    "    train_distilbert,\n",
-    "    config=search_space,\n",
-    "    init_config={\n",
-    "        \"num_train_epochs\": 1,\n",
-    "    },\n",
-    "    metric=HP_METRIC,\n",
-    "    mode=MODE,\n",
-    "    report_intermediate_result=False,\n",
-    "    # uncomment the following if report_intermediate_result = True\n",
-    "    # max_resource=max_num_epoch, min_resource=1,\n",
-    "    resources_per_trial={\"gpu\": 1},\n",
-    "    local_dir='logs/',\n",
-    "    num_samples=num_samples,\n",
-    "    time_budget_s=time_budget_s,\n",
-    "    use_ray=True,\n",
-    ")\n",
-    "\n",
-    "ray.shutdown()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "best_trial = analysis.get_best_trial(HP_METRIC, MODE, \"all\")\n",
-    "metric = best_trial.metric_analysis[HP_METRIC][MODE]\n",
-    "print(f\"n_trials={len(analysis.trials)}\")\n",
-    "print(f\"time={time.time()-start_time}\")\n",
-    "print(f\"Best model eval {HP_METRIC}: {metric:.4f}\")\n",
-    "print(f\"Best model parameters: {best_trial.config}\")\n"
-   ]
-  },
-  {
-   "source": [
-    "## Next Steps\n",
-    "\n",
-    "Notice that we only reported the metric with `flaml.tune.report` at the end of full training loop. It is possible to enable reporting of intermediate performance - allowing early stopping - as follows:\n",
-    "\n",
-    "- Huggingface provides _Callbacks_ which can be used to insert the `flaml.tune.report` call inside the training loop\n",
-    "- Make sure to set `do_eval=True` in the `TrainingArguments` provided to `Trainer` and adjust theevaluation frequency accordingly"
-   ],
-   "cell_type": "markdown",
-   "metadata": {}
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "flaml",
-   "language": "python",
-   "name": "flaml"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.7.6"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
-}

notebook/transformers_requirements.txt

@@ -1,4 +0,0 @@
-torch
-transformers
-datasets
-ipywidgets

setup.py

@@ -20,7 +20,6 @@ install_requires = [
     "scipy>=1.4.1",
     "catboost>=0.23",
     "scikit-learn>=0.23.2",
-    "optuna==2.3.0"
 ],
 
 
@@ -48,14 +47,21 @@ setuptools.setup(
             "coverage>=5.3",
             "xgboost<1.3",
             "rgf-python",
+            "optuna==2.3.0",
+        ],
+        "blendsearch": [
+            "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

@@ -0,0 +1,219 @@
+'''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 = "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"
+
+    # Define tokenize method
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT, use_fast=True)
+    def tokenize(examples):
+        return tokenizer(examples[COLUMN_NAME], 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_distilbert.log'))
+logger.setLevel(logging.INFO)
+
+import flaml
+
+def train_distilbert(config: dict):
+
+    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)
+
+    # 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,
+        disable_tqdm=True,
+        logging_steps=20000,
+        save_total_limit=0,
+        **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"],
+        matthews_correlation=eval_output["eval_matthews_correlation"],
+        )
+
+
+def _test_distillbert(method='BlendSearch'):
+ 
+    max_num_epoch = 64
+    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-6, 1e-4),
+        "adam_beta1": flaml.tune.uniform(0.8, 0.99),
+        "adam_beta2": flaml.tune.loguniform(98e-2, 9999e-4),
+        "adam_epsilon": flaml.tune.loguniform(1e-9, 1e-7),
+    }
+
+    start_time = time.time()
+    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,
+        }])
+    elif 'BlendSearch' == method:
+        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()
+    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_distilbert,
+        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_distillbert_cfo():
+    _test_distillbert('CFO')
+
+
+def _test_distillbert_dragonfly():
+    _test_distillbert('Dragonfly')
+
+
+def _test_distillbert_skopt():
+    _test_distillbert('SkOpt')
+
+
+def _test_distillbert_nevergrad():
+    _test_distillbert('Nevergrad')
+
+
+def _test_distillbert_zoopt():
+    _test_distillbert('ZOOpt')
+
+
+def _test_distillbert_ax():
+    _test_distillbert('Ax')
+
+
+def __test_distillbert_hyperopt():
+    _test_distillbert('HyperOpt')
+
+
+def _test_distillbert_optuna():
+    _test_distillbert('Optuna')
+
+
+def _test_distillbert_asha():
+    _test_distillbert('ASHA')
+
+
+def _test_distillbert_bohb():
+    _test_distillbert('BOHB')
+
+
+if __name__ == "__main__":
+    _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)
@@ -274,7 +281,7 @@ class TestAutoML(unittest.TestCase):
 
         automl_experiment = AutoML()
         automl_settings = {
-            "time_budget": 2,
+            "time_budget": 3,
             "metric": 'ap',
             "task": 'classification',
             "log_file_name": "test/sparse_classification.log",
@@ -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_pytorch_cifar10.py

@@ -1,3 +1,5 @@
+'''Require: pip install torchvision ray flaml[blendsearch]
+'''
 import unittest
 import os
 import time
@@ -24,7 +26,6 @@ def load_data(data_dir="./data"):
 # __load_data_end__
 
 
-import numpy as np
 try:
     import torch
     import torch.nn as nn

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

@@ -1,5 +1,6 @@
+'''Require: pip install flaml[test,ray]
+'''
 import unittest
-import os
 import time
 from sklearn.model_selection import train_test_split
 import sklearn.metrics
@@ -48,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]),
@@ -56,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()
@@ -138,6 +138,7 @@ def _test_xgboost(method='BlendSearch'):
                     scheduler=scheduler, search_alg=algo)
             ray.shutdown()
             # # Load the best model checkpoint
+            # import os
             # best_bst = xgb.Booster()
             # best_bst.load_model(os.path.join(analysis.best_checkpoint,
             #  "model.xgb"))
@@ -152,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()
 

test/test_xgboost2d.py

@@ -8,11 +8,7 @@ from flaml.model import XGBoostSklearnEstimator
 from flaml import tune
 
 
-# dataset = "blood-transfusion-service-center"
-# dataset = "Australian"
 dataset = "credit-g"
-# dataset = "phoneme"
-# dataset = "kc1"
 
 
 class XGBoost2D(XGBoostSklearnEstimator):
@@ -50,8 +46,11 @@ def test_simple(method=None):
         "log_type": "all",
         "time_budget": 3#6000,
     }
-
-    X, y = fetch_openml(name=dataset, return_X_y=True)
+    try:
+        X, y = fetch_openml(name=dataset, return_X_y=True)
+    except:
+        from sklearn.datasets import load_wine
+        X, y = load_wine(return_X_y=True)
     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33,
      random_state=42)
     automl.fit(X_train=X_train, y_train=y_train, **automl_settings)