mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-24 01:11:45 +00:00
2143 lines
76 KiB
Python
2143 lines
76 KiB
Python
# !
|
|
# * Copyright (c) FLAML authors. All rights reserved.
|
|
# * Licensed under the MIT License. See LICENSE file in the
|
|
# * project root for license information.
|
|
from contextlib import contextmanager
|
|
from functools import partial
|
|
import signal
|
|
import os
|
|
from typing import Callable, List
|
|
import numpy as np
|
|
import time
|
|
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
|
|
from sklearn.ensemble import ExtraTreesRegressor, ExtraTreesClassifier
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.dummy import DummyClassifier, DummyRegressor
|
|
from scipy.sparse import issparse
|
|
import logging
|
|
import shutil
|
|
from pandas import DataFrame, Series, to_datetime
|
|
import sys
|
|
import math
|
|
from . import tune
|
|
from .data import (
|
|
group_counts,
|
|
CLASSIFICATION,
|
|
TS_FORECASTREGRESSION,
|
|
TS_TIMESTAMP_COL,
|
|
TS_VALUE_COL,
|
|
SEQCLASSIFICATION,
|
|
SEQREGRESSION,
|
|
TOKENCLASSIFICATION,
|
|
SUMMARIZATION,
|
|
NLG_TASKS,
|
|
MULTICHOICECLASSIFICATION,
|
|
)
|
|
|
|
try:
|
|
import psutil
|
|
except ImportError:
|
|
psutil = None
|
|
try:
|
|
import resource
|
|
except ImportError:
|
|
resource = None
|
|
|
|
logger = logging.getLogger("flaml.automl")
|
|
FREE_MEM_RATIO = 0.2
|
|
|
|
|
|
def TimeoutHandler(sig, frame):
|
|
raise TimeoutError(sig, frame)
|
|
|
|
|
|
@contextmanager
|
|
def limit_resource(memory_limit, time_limit):
|
|
if memory_limit > 0:
|
|
soft, hard = resource.getrlimit(resource.RLIMIT_AS)
|
|
if soft < 0 and (hard < 0 or memory_limit <= hard) or memory_limit < soft:
|
|
resource.setrlimit(resource.RLIMIT_AS, (memory_limit, hard))
|
|
main_thread = False
|
|
if time_limit is not None:
|
|
try:
|
|
signal.signal(signal.SIGALRM, TimeoutHandler)
|
|
signal.alarm(int(time_limit) or 1)
|
|
main_thread = True
|
|
except ValueError:
|
|
pass
|
|
try:
|
|
yield
|
|
finally:
|
|
if main_thread:
|
|
signal.alarm(0)
|
|
if memory_limit > 0:
|
|
resource.setrlimit(resource.RLIMIT_AS, (soft, hard))
|
|
|
|
|
|
class BaseEstimator:
|
|
"""The abstract class for all learners.
|
|
|
|
Typical examples:
|
|
* XGBoostEstimator: for regression.
|
|
* XGBoostSklearnEstimator: for classification.
|
|
* LGBMEstimator, RandomForestEstimator, LRL1Classifier, LRL2Classifier:
|
|
for both regression and classification.
|
|
"""
|
|
|
|
def __init__(self, task="binary", **config):
|
|
"""Constructor.
|
|
|
|
Args:
|
|
task: A string of the task type, one of
|
|
'binary', 'multiclass', 'regression', 'rank', 'seq-classification',
|
|
'seq-regression', 'token-classification', 'multichoice-classification',
|
|
'summarization', 'ts_forecast', 'ts_forecast_classification'.
|
|
config: A dictionary containing the hyperparameter names, 'n_jobs' as keys.
|
|
n_jobs is the number of parallel threads.
|
|
"""
|
|
self._task = task
|
|
self.params = self.config2params(config)
|
|
self.estimator_class = self._model = None
|
|
if "_estimator_type" in config:
|
|
self._estimator_type = self.params.pop("_estimator_type")
|
|
else:
|
|
self._estimator_type = (
|
|
"classifier" if task in CLASSIFICATION else "regressor"
|
|
)
|
|
|
|
def get_params(self, deep=False):
|
|
params = self.params.copy()
|
|
params["task"] = self._task
|
|
if hasattr(self, "_estimator_type"):
|
|
params["_estimator_type"] = self._estimator_type
|
|
return params
|
|
|
|
@property
|
|
def classes_(self):
|
|
return self._model.classes_
|
|
|
|
@property
|
|
def n_features_in_(self):
|
|
return self._model.n_features_in_
|
|
|
|
@property
|
|
def model(self):
|
|
"""Trained model after fit() is called, or None before fit() is called."""
|
|
return self._model
|
|
|
|
@property
|
|
def estimator(self):
|
|
"""Trained model after fit() is called, or None before fit() is called."""
|
|
return self._model
|
|
|
|
def _preprocess(self, X):
|
|
return X
|
|
|
|
def _fit(self, X_train, y_train, **kwargs):
|
|
|
|
current_time = time.time()
|
|
if "groups" in kwargs:
|
|
kwargs = kwargs.copy()
|
|
groups = kwargs.pop("groups")
|
|
if self._task == "rank":
|
|
kwargs["group"] = group_counts(groups)
|
|
# groups_val = kwargs.get('groups_val')
|
|
# if groups_val is not None:
|
|
# kwargs['eval_group'] = [group_counts(groups_val)]
|
|
# kwargs['eval_set'] = [
|
|
# (kwargs['X_val'], kwargs['y_val'])]
|
|
# kwargs['verbose'] = False
|
|
# del kwargs['groups_val'], kwargs['X_val'], kwargs['y_val']
|
|
X_train = self._preprocess(X_train)
|
|
model = self.estimator_class(**self.params)
|
|
if logger.level == logging.DEBUG:
|
|
logger.debug(f"flaml.model - {model} fit started")
|
|
model.fit(X_train, y_train, **kwargs)
|
|
if logger.level == logging.DEBUG:
|
|
logger.debug(f"flaml.model - {model} fit finished")
|
|
train_time = time.time() - current_time
|
|
self._model = model
|
|
return train_time
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
"""Train the model from given training data.
|
|
|
|
Args:
|
|
X_train: A numpy array or a dataframe of training data in shape n*m.
|
|
y_train: A numpy array or a series of labels in shape n*1.
|
|
budget: A float of the time budget in seconds.
|
|
|
|
Returns:
|
|
train_time: A float of the training time in seconds.
|
|
"""
|
|
if (
|
|
getattr(self, "limit_resource", None)
|
|
and resource is not None
|
|
and (budget is not None or psutil is not None)
|
|
):
|
|
start_time = time.time()
|
|
mem = psutil.virtual_memory() if psutil is not None else None
|
|
try:
|
|
with limit_resource(
|
|
mem.available * (1 - FREE_MEM_RATIO)
|
|
+ psutil.Process(os.getpid()).memory_info().rss
|
|
if mem is not None
|
|
else -1,
|
|
budget,
|
|
):
|
|
train_time = self._fit(X_train, y_train, **kwargs)
|
|
except (MemoryError, TimeoutError) as e:
|
|
logger.warning(f"{e.__class__} {e}")
|
|
if self._task in CLASSIFICATION:
|
|
model = DummyClassifier()
|
|
else:
|
|
model = DummyRegressor()
|
|
X_train = self._preprocess(X_train)
|
|
model.fit(X_train, y_train)
|
|
self._model = model
|
|
train_time = time.time() - start_time
|
|
else:
|
|
train_time = self._fit(X_train, y_train, **kwargs)
|
|
return train_time
|
|
|
|
def predict(self, X, **kwargs):
|
|
"""Predict label from features.
|
|
|
|
Args:
|
|
X: A numpy array or a dataframe of featurized instances, shape n*m.
|
|
|
|
Returns:
|
|
A numpy array of shape n*1.
|
|
Each element is the label for a instance.
|
|
"""
|
|
if self._model is not None:
|
|
X = self._preprocess(X)
|
|
return self._model.predict(X)
|
|
else:
|
|
logger.warning(
|
|
"Estimator is not fit yet. Please run fit() before predict()."
|
|
)
|
|
return np.ones(X.shape[0])
|
|
|
|
def predict_proba(self, X, **kwargs):
|
|
"""Predict the probability of each class from features.
|
|
|
|
Only works for classification problems
|
|
|
|
Args:
|
|
X: A numpy array of featurized instances, shape n*m.
|
|
|
|
Returns:
|
|
A numpy array of shape n*c. c is the # classes.
|
|
Each element at (i,j) is the probability for instance i to be in
|
|
class j.
|
|
"""
|
|
assert self._task in CLASSIFICATION, "predict_proba() only for classification."
|
|
|
|
X = self._preprocess(X)
|
|
return self._model.predict_proba(X)
|
|
|
|
def score(self, X_val: DataFrame, y_val: Series, **kwargs):
|
|
"""Report the evaluation score of a trained estimator.
|
|
|
|
|
|
Args:
|
|
X_val: A pandas dataframe of the validation input data.
|
|
y_val: A pandas series of the validation label.
|
|
kwargs: keyword argument of the evaluation function, for example:
|
|
- metric: A string of the metric name or a function
|
|
e.g., 'accuracy', 'roc_auc', 'roc_auc_ovr', 'roc_auc_ovo',
|
|
'f1', 'micro_f1', 'macro_f1', 'log_loss', 'mae', 'mse', 'r2',
|
|
'mape'. Default is 'auto'.
|
|
If metric is given, the score will report the user specified metric.
|
|
If metric is not given, the metric is set to accuracy for classification and r2
|
|
for regression.
|
|
You can also pass a customized metric function, for examples on how to pass a
|
|
customized metric function, please check
|
|
[test/nlp/test_autohf_custom_metric.py](https://github.com/microsoft/FLAML/blob/main/test/nlp/test_autohf_custom_metric.py) and
|
|
[test/automl/test_multiclass.py](https://github.com/microsoft/FLAML/blob/main/test/automl/test_multiclass.py).
|
|
|
|
Returns:
|
|
The evaluation score on the validation dataset.
|
|
"""
|
|
from .ml import metric_loss_score
|
|
from .ml import is_min_metric
|
|
|
|
if self._model is not None:
|
|
if self._task == "rank":
|
|
raise NotImplementedError(
|
|
"AutoML.score() is not implemented for ranking"
|
|
)
|
|
else:
|
|
X_val = self._preprocess(X_val)
|
|
metric = kwargs.get("metric", None)
|
|
if metric:
|
|
y_pred = self.predict(X_val, **kwargs)
|
|
if is_min_metric(metric):
|
|
return metric_loss_score(metric, y_pred, y_val)
|
|
else:
|
|
return 1.0 - metric_loss_score(metric, y_pred, y_val)
|
|
else:
|
|
return self._model.score(X_val, y_val, **kwargs)
|
|
else:
|
|
logger.warning(
|
|
"Estimator is not fit yet. Please run fit() before predict()."
|
|
)
|
|
return 0.0
|
|
|
|
def cleanup(self):
|
|
del self._model
|
|
self._model = None
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, task, **params):
|
|
"""[required method] search space.
|
|
|
|
Args:
|
|
data_size: A tuple of two integers, number of rows and columns.
|
|
task: A str of the task type, e.g., "binary", "multiclass", "regression".
|
|
|
|
Returns:
|
|
A dictionary of the search space.
|
|
Each key is the name of a hyperparameter, and value is a dict with
|
|
its domain (required) and low_cost_init_value, init_value,
|
|
cat_hp_cost (if applicable).
|
|
e.g., ```{'domain': tune.randint(lower=1, upper=10), 'init_value': 1}```.
|
|
"""
|
|
return {}
|
|
|
|
@classmethod
|
|
def size(cls, config: dict) -> float:
|
|
"""[optional method] memory size of the estimator in bytes.
|
|
|
|
Args:
|
|
config: A dict of the hyperparameter config.
|
|
|
|
Returns:
|
|
A float of the memory size required by the estimator to train the
|
|
given config.
|
|
"""
|
|
return 1.0
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls) -> float:
|
|
"""[optional method] relative cost compared to lightgbm."""
|
|
return 1.0
|
|
|
|
@classmethod
|
|
def init(cls):
|
|
"""[optional method] initialize the class."""
|
|
pass
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
"""[optional method] config dict to params dict
|
|
|
|
Args:
|
|
config: A dict of the hyperparameter config.
|
|
|
|
Returns:
|
|
A dict that will be passed to self.estimator_class's constructor.
|
|
"""
|
|
params = config.copy()
|
|
if "FLAML_sample_size" in params:
|
|
params.pop("FLAML_sample_size")
|
|
return params
|
|
|
|
|
|
class TransformersEstimator(BaseEstimator):
|
|
"""The class for fine-tuning language models, using huggingface transformers API."""
|
|
|
|
ITER_HP = "global_max_steps"
|
|
|
|
def __init__(self, task="seq-classification", **config):
|
|
super().__init__(task, **config)
|
|
import uuid
|
|
|
|
self.trial_id = str(uuid.uuid1().hex)[:8]
|
|
if task not in NLG_TASKS: # TODO: not in NLG_TASKS
|
|
from .nlp.huggingface.training_args import (
|
|
TrainingArgumentsForAuto as TrainingArguments,
|
|
)
|
|
else:
|
|
from .nlp.huggingface.training_args import (
|
|
Seq2SeqTrainingArgumentsForAuto as TrainingArguments,
|
|
)
|
|
self._TrainingArguments = TrainingArguments
|
|
|
|
@staticmethod
|
|
def _join(X_train, y_train, task):
|
|
y_train = DataFrame(y_train, index=X_train.index)
|
|
y_train.columns = ["label"] if task != TOKENCLASSIFICATION else ["labels"]
|
|
train_df = X_train.join(y_train)
|
|
return train_df
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, task, **params):
|
|
search_space_dict = {
|
|
"learning_rate": {
|
|
"domain": tune.loguniform(lower=1e-6, upper=1e-3),
|
|
"init_value": 1e-5,
|
|
},
|
|
"num_train_epochs": {
|
|
"domain": tune.loguniform(lower=0.1, upper=10.0),
|
|
"init_value": 3.0, # to be consistent with roberta
|
|
},
|
|
"per_device_train_batch_size": {
|
|
"domain": tune.choice([4, 8, 16, 32]),
|
|
"init_value": 32,
|
|
},
|
|
"warmup_ratio": {
|
|
"domain": tune.uniform(lower=0.0, upper=0.3),
|
|
"init_value": 0.0,
|
|
},
|
|
"weight_decay": {
|
|
"domain": tune.uniform(lower=0.0, upper=0.3),
|
|
"init_value": 0.0,
|
|
},
|
|
"adam_epsilon": {
|
|
"domain": tune.loguniform(lower=1e-8, upper=1e-6),
|
|
"init_value": 1e-6,
|
|
},
|
|
"seed": {"domain": tune.choice(list(range(40, 45))), "init_value": 42},
|
|
"global_max_steps": {
|
|
"domain": sys.maxsize,
|
|
"init_value": sys.maxsize,
|
|
},
|
|
}
|
|
|
|
return search_space_dict
|
|
|
|
@property
|
|
def checkpoint_freq(self):
|
|
return (
|
|
int(
|
|
min(self._training_args.num_train_epochs, 1)
|
|
* len(self._X_train)
|
|
/ self._training_args.per_device_train_batch_size
|
|
/ self._training_args.ckpt_per_epoch
|
|
)
|
|
+ 1
|
|
)
|
|
|
|
@property
|
|
def fp16(self):
|
|
return self._kwargs.get("gpu_per_trial") and self._training_args.fp16
|
|
|
|
@property
|
|
def no_cuda(self):
|
|
return not self._kwargs.get("gpu_per_trial")
|
|
|
|
def _set_training_args(self, **kwargs):
|
|
from .nlp.utils import date_str, Counter
|
|
|
|
for (key, val) in kwargs.items():
|
|
assert key not in self.params, (
|
|
"Since {} is in the search space, it cannot exist in 'custom_fit_kwargs' at the same time."
|
|
"If you need to fix the value of {} to {}, the only way is to add a single-value domain in the search "
|
|
"space by adding:\n '{}': {{ 'domain': {} }} to 'custom_hp'. For example:"
|
|
'automl_settings["custom_hp"] = {{ "transformer": {{ "model_path": {{ "domain" : '
|
|
'"google/electra-small-discriminator" }} }} }}'.format(
|
|
key, key, val, key, val
|
|
)
|
|
)
|
|
|
|
"""
|
|
If use has specified any custom args for TrainingArguments, update these arguments
|
|
"""
|
|
self._training_args = self._TrainingArguments(**kwargs)
|
|
|
|
"""
|
|
Update the attributes in TrainingArguments with self.params values
|
|
"""
|
|
for key, val in self.params.items():
|
|
if hasattr(self._training_args, key):
|
|
setattr(self._training_args, key, val)
|
|
|
|
"""
|
|
Update the attributes in TrainingArguments that depends on the values of self.params
|
|
"""
|
|
local_dir = os.path.join(
|
|
self._training_args.output_dir, "train_{}".format(date_str())
|
|
)
|
|
if self._use_ray is True:
|
|
import ray
|
|
|
|
self._training_args.output_dir = ray.tune.get_trial_dir()
|
|
else:
|
|
self._training_args.output_dir = Counter.get_trial_fold_name(
|
|
local_dir, self.params, self.trial_id
|
|
)
|
|
|
|
self._training_args.eval_steps = (
|
|
self._training_args.logging_steps
|
|
) = self._training_args.saving_steps = self.checkpoint_freq
|
|
self._training_args.fp16 = self.fp16
|
|
self._training_args.no_cuda = self.no_cuda
|
|
|
|
def _preprocess(self, X, y=None, **kwargs):
|
|
from .nlp.utils import tokenize_text, is_a_list_of_str
|
|
|
|
is_str = str(X.dtypes[0]) in ("string", "str")
|
|
is_list_of_str = is_a_list_of_str(X[list(X.keys())[0]].to_list()[0])
|
|
|
|
if is_str or is_list_of_str:
|
|
return tokenize_text(
|
|
X=X,
|
|
Y=y,
|
|
task=self._task,
|
|
hf_args=self._training_args,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
else:
|
|
return X, None
|
|
|
|
def _model_init(self):
|
|
from .nlp.utils import load_model
|
|
|
|
this_model = load_model(
|
|
checkpoint_path=self._training_args.model_path,
|
|
task=self._task,
|
|
num_labels=self.num_labels,
|
|
)
|
|
return this_model
|
|
|
|
def preprocess_data(self, X, y):
|
|
from datasets import Dataset
|
|
|
|
if (self._task not in NLG_TASKS) and (self._task != TOKENCLASSIFICATION):
|
|
processed_X, _ = self._preprocess(X=X, **self._kwargs)
|
|
processed_y = y
|
|
else:
|
|
processed_X, processed_y = self._preprocess(X=X, y=y, **self._kwargs)
|
|
|
|
processed_dataset = Dataset.from_pandas(
|
|
TransformersEstimator._join(processed_X, processed_y, self._task)
|
|
)
|
|
return processed_dataset, processed_X, processed_y
|
|
|
|
@property
|
|
def num_labels(self):
|
|
from .data import SEQCLASSIFICATION, SEQREGRESSION, TOKENCLASSIFICATION
|
|
|
|
if self._task == SEQREGRESSION:
|
|
return 1
|
|
elif self._task == SEQCLASSIFICATION:
|
|
return len(set(self._y_train))
|
|
elif self._task == TOKENCLASSIFICATION:
|
|
return len(set([a for b in self._y_train.tolist() for a in b]))
|
|
else:
|
|
return None
|
|
|
|
@property
|
|
def tokenizer(self):
|
|
from transformers import AutoTokenizer
|
|
|
|
if self._task == SUMMARIZATION:
|
|
return AutoTokenizer.from_pretrained(
|
|
pretrained_model_name_or_path=self._training_args.model_path,
|
|
cache_dir=None,
|
|
use_fast=True,
|
|
revision="main",
|
|
use_auth_token=None,
|
|
)
|
|
else:
|
|
return AutoTokenizer.from_pretrained(
|
|
self._training_args.model_path, use_fast=True
|
|
)
|
|
|
|
@property
|
|
def data_collator(self):
|
|
from .nlp.huggingface.data_collator import task_to_datacollator_class
|
|
|
|
return (
|
|
task_to_datacollator_class[self._task](
|
|
tokenizer=self.tokenizer,
|
|
pad_to_multiple_of=8, # if self._training_args.fp16 else None,
|
|
)
|
|
if self._task in (MULTICHOICECLASSIFICATION, TOKENCLASSIFICATION)
|
|
else None
|
|
)
|
|
|
|
def fit(
|
|
self,
|
|
X_train: DataFrame,
|
|
y_train: Series,
|
|
budget=None,
|
|
X_val=None,
|
|
y_val=None,
|
|
gpu_per_trial=None,
|
|
metric=None,
|
|
**kwargs,
|
|
):
|
|
import transformers
|
|
|
|
transformers.logging.set_verbosity_error()
|
|
|
|
from transformers import TrainerCallback
|
|
from transformers.trainer_utils import set_seed
|
|
from .nlp.huggingface.trainer import TrainerForAuto
|
|
|
|
try:
|
|
from ray.tune import is_session_enabled
|
|
|
|
self._use_ray = is_session_enabled()
|
|
except ImportError:
|
|
self._use_ray = False
|
|
|
|
this_params = self.params
|
|
self._kwargs = kwargs
|
|
|
|
self._X_train, self._y_train = X_train, y_train
|
|
self._set_training_args(**kwargs)
|
|
|
|
train_dataset, self._X_train, self._y_train = self.preprocess_data(
|
|
X_train, y_train
|
|
)
|
|
if X_val is not None:
|
|
eval_dataset, self._X_val, self._y_val = self.preprocess_data(X_val, y_val)
|
|
else:
|
|
eval_dataset, self._X_val, self._y_val = None, None, None
|
|
|
|
set_seed(self.params.get("seed", self._training_args.seed))
|
|
self._metric = metric
|
|
|
|
class EarlyStoppingCallbackForAuto(TrainerCallback):
|
|
def on_train_begin(self, args, state, control, **callback_kwargs):
|
|
self.train_begin_time = time.time()
|
|
|
|
def on_step_begin(self, args, state, control, **callback_kwargs):
|
|
self.step_begin_time = time.time()
|
|
|
|
def on_step_end(self, args, state, control, **callback_kwargs):
|
|
if state.global_step == 1:
|
|
self.time_per_iter = time.time() - self.step_begin_time
|
|
if (
|
|
budget
|
|
and (
|
|
time.time() + self.time_per_iter
|
|
> self.train_begin_time + budget
|
|
)
|
|
or state.global_step >= this_params[TransformersEstimator.ITER_HP]
|
|
):
|
|
control.should_training_stop = True
|
|
control.should_save = True
|
|
control.should_evaluate = True
|
|
return control
|
|
|
|
def on_epoch_end(self, args, state, control, **callback_kwargs):
|
|
if (
|
|
control.should_training_stop
|
|
or state.epoch + 1 >= args.num_train_epochs
|
|
):
|
|
control.should_save = True
|
|
control.should_evaluate = True
|
|
|
|
self._trainer = TrainerForAuto(
|
|
args=self._training_args,
|
|
model_init=self._model_init,
|
|
train_dataset=train_dataset,
|
|
eval_dataset=eval_dataset,
|
|
tokenizer=self.tokenizer,
|
|
data_collator=self.data_collator,
|
|
compute_metrics=self._compute_metrics_by_dataset_name,
|
|
callbacks=[EarlyStoppingCallbackForAuto],
|
|
)
|
|
|
|
if self._task in NLG_TASKS:
|
|
setattr(self._trainer, "_is_seq2seq", True)
|
|
|
|
"""
|
|
When not using ray for tuning, set the limit of CUDA_VISIBLE_DEVICES to math.ceil(gpu_per_trial),
|
|
so each estimator does not see all the GPUs
|
|
"""
|
|
if gpu_per_trial is not None:
|
|
tmp_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "")
|
|
self._trainer.args._n_gpu = gpu_per_trial
|
|
|
|
# if gpu_per_trial == 0:
|
|
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
|
if tmp_cuda_visible_devices.count(",") != math.ceil(gpu_per_trial) - 1:
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
|
|
[str(x) for x in range(math.ceil(gpu_per_trial))]
|
|
)
|
|
|
|
import time
|
|
|
|
start_time = time.time()
|
|
self._trainer.train()
|
|
|
|
if gpu_per_trial is not None:
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = tmp_cuda_visible_devices
|
|
|
|
self.params[self.ITER_HP] = self._trainer.state.global_step
|
|
|
|
self._checkpoint_path = self._select_checkpoint(self._trainer)
|
|
self._ckpt_remains = list(self._trainer.ckpt_to_metric.keys())
|
|
|
|
if hasattr(self._trainer, "intermediate_results"):
|
|
self.intermediate_results = [
|
|
x[1]
|
|
for x in sorted(
|
|
self._trainer.intermediate_results.items(), key=lambda x: x[0]
|
|
)
|
|
]
|
|
self._trainer = None
|
|
|
|
return time.time() - start_time
|
|
|
|
def _delete_one_ckpt(self, ckpt_location):
|
|
if self._use_ray is False:
|
|
try:
|
|
shutil.rmtree(ckpt_location)
|
|
except FileNotFoundError:
|
|
logger.warning("checkpoint {} not found".format(ckpt_location))
|
|
|
|
def cleanup(self):
|
|
super().cleanup()
|
|
if hasattr(self, "_ckpt_remains"):
|
|
for each_ckpt in self._ckpt_remains:
|
|
self._delete_one_ckpt(each_ckpt)
|
|
|
|
def _select_checkpoint(self, trainer):
|
|
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
|
|
|
|
if trainer.ckpt_to_metric:
|
|
best_ckpt, _ = min(
|
|
trainer.ckpt_to_metric.items(), key=lambda x: x[1]["eval_loss"]
|
|
)
|
|
best_ckpt_global_step = trainer.ckpt_to_global_step[best_ckpt]
|
|
for each_ckpt in list(trainer.ckpt_to_metric):
|
|
if each_ckpt != best_ckpt:
|
|
del trainer.ckpt_to_metric[each_ckpt]
|
|
del trainer.ckpt_to_global_step[each_ckpt]
|
|
self._delete_one_ckpt(each_ckpt)
|
|
else:
|
|
best_ckpt_global_step = trainer.state.global_step
|
|
best_ckpt = os.path.join(
|
|
trainer.args.output_dir,
|
|
f"{PREFIX_CHECKPOINT_DIR}-{best_ckpt_global_step}",
|
|
)
|
|
self.params[self.ITER_HP] = best_ckpt_global_step
|
|
logger.debug(trainer.state.global_step)
|
|
logger.debug(trainer.ckpt_to_global_step)
|
|
return best_ckpt
|
|
|
|
def _compute_metrics_by_dataset_name(self, eval_pred):
|
|
if isinstance(self._metric, str):
|
|
from .ml import metric_loss_score
|
|
from .nlp.utils import postprocess_text
|
|
|
|
predictions, labels = eval_pred
|
|
if self._task in NLG_TASKS:
|
|
if isinstance(predictions, tuple):
|
|
predictions = np.argmax(predictions[0], axis=2)
|
|
decoded_preds = self.tokenizer.batch_decode(
|
|
predictions, skip_special_tokens=True
|
|
)
|
|
labels = np.where(labels != -100, labels, self.tokenizer.pad_token_id)
|
|
decoded_labels = self.tokenizer.batch_decode(
|
|
labels, skip_special_tokens=True
|
|
)
|
|
predictions, labels = postprocess_text(decoded_preds, decoded_labels)
|
|
else:
|
|
predictions = (
|
|
np.squeeze(predictions)
|
|
if self._task == SEQREGRESSION
|
|
else np.argmax(predictions, axis=2)
|
|
if self._task == TOKENCLASSIFICATION
|
|
else np.argmax(predictions, axis=1)
|
|
)
|
|
metric_dict = {
|
|
"automl_metric": metric_loss_score(
|
|
metric_name=self._metric,
|
|
y_predict=predictions,
|
|
y_true=labels,
|
|
labels=self._training_args.label_list,
|
|
)
|
|
}
|
|
else:
|
|
loss, metric_dict = self._metric(
|
|
X_test=self._X_val,
|
|
y_test=self._y_val,
|
|
estimator=self,
|
|
labels=None,
|
|
X_train=self._X_train,
|
|
y_train=self._y_train,
|
|
)
|
|
metric_dict["automl_metric"] = loss
|
|
|
|
return metric_dict
|
|
|
|
def _init_model_for_predict(self):
|
|
from .nlp.huggingface.trainer import TrainerForAuto
|
|
|
|
"""
|
|
Need to reinit training_args because of a bug in deepspeed: if not reinit, the deepspeed config will be inconsistent
|
|
with HF config https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py#L947
|
|
"""
|
|
training_args = self._TrainingArguments(
|
|
local_rank=-1, model_path=self._checkpoint_path, fp16=self.fp16
|
|
)
|
|
for key, val in self._training_args.__dict__.items():
|
|
if key not in ("local_rank", "model_path", "fp16"):
|
|
setattr(training_args, key, val)
|
|
self._training_args = training_args
|
|
|
|
new_trainer = TrainerForAuto(
|
|
model=self._model_init(),
|
|
args=self._training_args,
|
|
data_collator=self.data_collator,
|
|
compute_metrics=self._compute_metrics_by_dataset_name,
|
|
)
|
|
if self._task in NLG_TASKS:
|
|
setattr(new_trainer, "_is_seq2seq", True)
|
|
return new_trainer
|
|
|
|
def predict_proba(self, X, **pred_kwargs):
|
|
from datasets import Dataset
|
|
|
|
if pred_kwargs:
|
|
for key, val in pred_kwargs.items():
|
|
setattr(self._training_args, key, val)
|
|
|
|
assert (
|
|
self._task in CLASSIFICATION
|
|
), "predict_proba() only for classification tasks."
|
|
|
|
X_test, _ = self._preprocess(X, **self._kwargs)
|
|
test_dataset = Dataset.from_pandas(X_test)
|
|
|
|
new_trainer = self._init_model_for_predict()
|
|
predictions = new_trainer.predict(test_dataset)
|
|
return predictions.predictions
|
|
|
|
def score(self, X_val: DataFrame, y_val: Series, **kwargs):
|
|
import transformers
|
|
|
|
transformers.logging.set_verbosity_error()
|
|
|
|
self._metric = kwargs["metric"]
|
|
|
|
eval_dataset, X_val, y_val = self.preprocess_data(X_val, y_val)
|
|
|
|
new_trainer = self._init_model_for_predict()
|
|
return new_trainer.evaluate(eval_dataset)
|
|
|
|
def predict(self, X, **pred_kwargs):
|
|
import transformers
|
|
from datasets import Dataset
|
|
|
|
transformers.logging.set_verbosity_error()
|
|
|
|
if pred_kwargs:
|
|
for key, val in pred_kwargs.items():
|
|
setattr(self._training_args, key, val)
|
|
|
|
X_test, _ = self._preprocess(X, **self._kwargs)
|
|
test_dataset = Dataset.from_pandas(X_test)
|
|
|
|
new_trainer = self._init_model_for_predict()
|
|
|
|
if self._task not in NLG_TASKS:
|
|
predictions = new_trainer.predict(test_dataset)
|
|
else:
|
|
predictions = new_trainer.predict(
|
|
test_dataset,
|
|
metric_key_prefix="predict",
|
|
)
|
|
|
|
if self._task == SEQCLASSIFICATION:
|
|
return np.argmax(predictions.predictions, axis=1)
|
|
elif self._task == SEQREGRESSION:
|
|
return predictions.predictions.reshape((len(predictions.predictions),))
|
|
elif self._task == TOKENCLASSIFICATION:
|
|
return np.argmax(predictions.predictions, axis=2)
|
|
elif self._task == SUMMARIZATION:
|
|
decoded_preds = self.tokenizer.batch_decode(
|
|
predictions.predictions, skip_special_tokens=True
|
|
)
|
|
return decoded_preds
|
|
elif self._task == MULTICHOICECLASSIFICATION:
|
|
return np.argmax(predictions.predictions, axis=1)
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
params[TransformersEstimator.ITER_HP] = params.get(
|
|
TransformersEstimator.ITER_HP, sys.maxsize
|
|
)
|
|
return params
|
|
|
|
|
|
class TransformersEstimatorModelSelection(TransformersEstimator):
|
|
def __init__(self, task="seq-classification", **config):
|
|
super().__init__(task, **config)
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, task, **params):
|
|
search_space_dict = TransformersEstimator.search_space(
|
|
data_size, task, **params
|
|
)
|
|
|
|
"""
|
|
For model selection, use the same search space regardless of memory constraint
|
|
If OOM, user should change the search space themselves
|
|
"""
|
|
|
|
search_space_dict["model_path"] = {
|
|
"domain": tune.choice(
|
|
[
|
|
"google/electra-base-discriminator",
|
|
"bert-base-uncased",
|
|
"roberta-base",
|
|
"facebook/muppet-roberta-base",
|
|
"google/electra-small-discriminator",
|
|
]
|
|
),
|
|
"init_value": "facebook/muppet-roberta-base",
|
|
}
|
|
return search_space_dict
|
|
|
|
|
|
class SKLearnEstimator(BaseEstimator):
|
|
"""The base class for tuning scikit-learn estimators."""
|
|
|
|
def __init__(self, task="binary", **config):
|
|
super().__init__(task, **config)
|
|
|
|
def _preprocess(self, X):
|
|
if isinstance(X, DataFrame):
|
|
cat_columns = X.select_dtypes(include=["category"]).columns
|
|
if not cat_columns.empty:
|
|
X = X.copy()
|
|
X[cat_columns] = X[cat_columns].apply(lambda x: x.cat.codes)
|
|
elif isinstance(X, np.ndarray) and X.dtype.kind not in "buif":
|
|
# numpy array is not of numeric dtype
|
|
X = DataFrame(X)
|
|
for col in X.columns:
|
|
if isinstance(X[col][0], str):
|
|
X[col] = X[col].astype("category").cat.codes
|
|
X = X.to_numpy()
|
|
return X
|
|
|
|
|
|
class LGBMEstimator(BaseEstimator):
|
|
"""The class for tuning LGBM, using sklearn API."""
|
|
|
|
ITER_HP = "n_estimators"
|
|
HAS_CALLBACK = True
|
|
DEFAULT_ITER = 100
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, **params):
|
|
upper = max(5, min(32768, int(data_size[0]))) # upper must be larger than lower
|
|
return {
|
|
"n_estimators": {
|
|
"domain": tune.lograndint(lower=4, upper=upper),
|
|
"init_value": 4,
|
|
"low_cost_init_value": 4,
|
|
},
|
|
"num_leaves": {
|
|
"domain": tune.lograndint(lower=4, upper=upper),
|
|
"init_value": 4,
|
|
"low_cost_init_value": 4,
|
|
},
|
|
"min_child_samples": {
|
|
"domain": tune.lograndint(lower=2, upper=2 ** 7 + 1),
|
|
"init_value": 20,
|
|
},
|
|
"learning_rate": {
|
|
"domain": tune.loguniform(lower=1 / 1024, upper=1.0),
|
|
"init_value": 0.1,
|
|
},
|
|
# 'subsample': {
|
|
# 'domain': tune.uniform(lower=0.1, upper=1.0),
|
|
# 'init_value': 1.0,
|
|
# },
|
|
"log_max_bin": { # log transformed with base 2
|
|
"domain": tune.lograndint(lower=3, upper=11),
|
|
"init_value": 8,
|
|
},
|
|
"colsample_bytree": {
|
|
"domain": tune.uniform(lower=0.01, upper=1.0),
|
|
"init_value": 1.0,
|
|
},
|
|
"reg_alpha": {
|
|
"domain": tune.loguniform(lower=1 / 1024, upper=1024),
|
|
"init_value": 1 / 1024,
|
|
},
|
|
"reg_lambda": {
|
|
"domain": tune.loguniform(lower=1 / 1024, upper=1024),
|
|
"init_value": 1.0,
|
|
},
|
|
}
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
if "log_max_bin" in params:
|
|
params["max_bin"] = (1 << params.pop("log_max_bin")) - 1
|
|
return params
|
|
|
|
@classmethod
|
|
def size(cls, config):
|
|
num_leaves = int(
|
|
round(
|
|
config.get("num_leaves")
|
|
or config.get("max_leaves")
|
|
or 1 << config.get("max_depth", 16)
|
|
)
|
|
)
|
|
n_estimators = int(round(config["n_estimators"]))
|
|
return (num_leaves * 3 + (num_leaves - 1) * 4 + 1.0) * n_estimators * 8
|
|
|
|
def __init__(self, task="binary", **config):
|
|
super().__init__(task, **config)
|
|
if "verbose" not in self.params:
|
|
self.params["verbose"] = -1
|
|
if "regression" == task:
|
|
from lightgbm import LGBMRegressor
|
|
|
|
self.estimator_class = LGBMRegressor
|
|
elif "rank" == task:
|
|
from lightgbm import LGBMRanker
|
|
|
|
self.estimator_class = LGBMRanker
|
|
else:
|
|
from lightgbm import LGBMClassifier
|
|
|
|
self.estimator_class = LGBMClassifier
|
|
self._time_per_iter = None
|
|
self._train_size = 0
|
|
self._mem_per_iter = -1
|
|
self.HAS_CALLBACK = self.HAS_CALLBACK and self._callbacks(0, 0) is not None
|
|
|
|
def _preprocess(self, X):
|
|
if (
|
|
not isinstance(X, DataFrame)
|
|
and issparse(X)
|
|
and np.issubdtype(X.dtype, np.integer)
|
|
):
|
|
X = X.astype(float)
|
|
elif isinstance(X, np.ndarray) and X.dtype.kind not in "buif":
|
|
# numpy array is not of numeric dtype
|
|
X = DataFrame(X)
|
|
for col in X.columns:
|
|
if isinstance(X[col][0], str):
|
|
X[col] = X[col].astype("category").cat.codes
|
|
X = X.to_numpy()
|
|
return X
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
start_time = time.time()
|
|
deadline = start_time + budget if budget else np.inf
|
|
n_iter = self.params.get(self.ITER_HP, self.DEFAULT_ITER)
|
|
trained = False
|
|
if not self.HAS_CALLBACK:
|
|
mem0 = psutil.virtual_memory().available if psutil is not None else 1
|
|
if (
|
|
(
|
|
not self._time_per_iter
|
|
or abs(self._train_size - X_train.shape[0]) > 4
|
|
)
|
|
and budget is not None
|
|
or self._mem_per_iter < 0
|
|
and psutil is not None
|
|
) and n_iter > 1:
|
|
self.params[self.ITER_HP] = 1
|
|
self._t1 = self._fit(X_train, y_train, **kwargs)
|
|
if budget is not None and self._t1 >= budget or n_iter == 1:
|
|
# self.params[self.ITER_HP] = n_iter
|
|
return self._t1
|
|
mem1 = psutil.virtual_memory().available if psutil is not None else 1
|
|
self._mem1 = mem0 - mem1
|
|
self.params[self.ITER_HP] = min(n_iter, 4)
|
|
self._t2 = self._fit(X_train, y_train, **kwargs)
|
|
mem2 = psutil.virtual_memory().available if psutil is not None else 1
|
|
self._mem2 = max(mem0 - mem2, self._mem1)
|
|
# if self._mem1 <= 0:
|
|
# self._mem_per_iter = self._mem2 / (self.params[self.ITER_HP] + 1)
|
|
# elif self._mem2 <= 0:
|
|
# self._mem_per_iter = self._mem1
|
|
# else:
|
|
self._mem_per_iter = min(
|
|
self._mem1, self._mem2 / self.params[self.ITER_HP]
|
|
)
|
|
# if self._mem_per_iter <= 1 and psutil is not None:
|
|
# n_iter = self.params[self.ITER_HP]
|
|
self._time_per_iter = (
|
|
(self._t2 - self._t1) / (self.params[self.ITER_HP] - 1)
|
|
if self._t2 > self._t1
|
|
else self._t1
|
|
if self._t1
|
|
else 0.001
|
|
)
|
|
self._train_size = X_train.shape[0]
|
|
if (
|
|
budget is not None
|
|
and self._t1 + self._t2 >= budget
|
|
or n_iter == self.params[self.ITER_HP]
|
|
):
|
|
# self.params[self.ITER_HP] = n_iter
|
|
return time.time() - start_time
|
|
trained = True
|
|
# logger.debug(mem0)
|
|
# logger.debug(self._mem_per_iter)
|
|
if n_iter > 1:
|
|
max_iter = min(
|
|
n_iter,
|
|
int(
|
|
(budget - time.time() + start_time - self._t1)
|
|
/ self._time_per_iter
|
|
+ 1
|
|
)
|
|
if budget is not None
|
|
else n_iter,
|
|
int((1 - FREE_MEM_RATIO) * mem0 / self._mem_per_iter)
|
|
if psutil is not None and self._mem_per_iter > 0
|
|
else n_iter,
|
|
)
|
|
if trained and max_iter <= self.params[self.ITER_HP]:
|
|
return time.time() - start_time
|
|
# when not trained, train at least one iter
|
|
self.params[self.ITER_HP] = max(max_iter, 1)
|
|
if self.HAS_CALLBACK:
|
|
kwargs_callbacks = kwargs.get("callbacks")
|
|
if kwargs_callbacks:
|
|
callbacks = kwargs_callbacks + self._callbacks(start_time, deadline)
|
|
kwargs.pop("callbacks")
|
|
else:
|
|
callbacks = self._callbacks(start_time, deadline)
|
|
self._fit(
|
|
X_train,
|
|
y_train,
|
|
callbacks=callbacks,
|
|
**kwargs,
|
|
)
|
|
best_iteration = (
|
|
self._model.get_booster().best_iteration
|
|
if isinstance(self, XGBoostSklearnEstimator)
|
|
else self._model.best_iteration_
|
|
)
|
|
if best_iteration is not None:
|
|
self._model.set_params(n_estimators=best_iteration + 1)
|
|
else:
|
|
self._fit(X_train, y_train, **kwargs)
|
|
train_time = time.time() - start_time
|
|
return train_time
|
|
|
|
def _callbacks(self, start_time, deadline) -> List[Callable]:
|
|
return [partial(self._callback, start_time, deadline)]
|
|
|
|
def _callback(self, start_time, deadline, env) -> None:
|
|
from lightgbm.callback import EarlyStopException
|
|
|
|
now = time.time()
|
|
if env.iteration == 0:
|
|
self._time_per_iter = now - start_time
|
|
if now + self._time_per_iter > deadline:
|
|
raise EarlyStopException(env.iteration, env.evaluation_result_list)
|
|
if psutil is not None:
|
|
mem = psutil.virtual_memory()
|
|
if mem.available / mem.total < FREE_MEM_RATIO:
|
|
raise EarlyStopException(env.iteration, env.evaluation_result_list)
|
|
|
|
|
|
class XGBoostEstimator(SKLearnEstimator):
|
|
"""The class for tuning XGBoost regressor, not using sklearn API."""
|
|
|
|
DEFAULT_ITER = 10
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, **params):
|
|
upper = max(5, min(32768, int(data_size[0]))) # upper must be larger than lower
|
|
return {
|
|
"n_estimators": {
|
|
"domain": tune.lograndint(lower=4, upper=upper),
|
|
"init_value": 4,
|
|
"low_cost_init_value": 4,
|
|
},
|
|
"max_leaves": {
|
|
"domain": tune.lograndint(lower=4, upper=upper),
|
|
"init_value": 4,
|
|
"low_cost_init_value": 4,
|
|
},
|
|
"max_depth": {
|
|
"domain": tune.choice([0, 6, 12]),
|
|
"init_value": 0,
|
|
},
|
|
"min_child_weight": {
|
|
"domain": tune.loguniform(lower=0.001, upper=128),
|
|
"init_value": 1,
|
|
},
|
|
"learning_rate": {
|
|
"domain": tune.loguniform(lower=1 / 1024, upper=1.0),
|
|
"init_value": 0.1,
|
|
},
|
|
"subsample": {
|
|
"domain": tune.uniform(lower=0.1, upper=1.0),
|
|
"init_value": 1.0,
|
|
},
|
|
"colsample_bylevel": {
|
|
"domain": tune.uniform(lower=0.01, upper=1.0),
|
|
"init_value": 1.0,
|
|
},
|
|
"colsample_bytree": {
|
|
"domain": tune.uniform(lower=0.01, upper=1.0),
|
|
"init_value": 1.0,
|
|
},
|
|
"reg_alpha": {
|
|
"domain": tune.loguniform(lower=1 / 1024, upper=1024),
|
|
"init_value": 1 / 1024,
|
|
},
|
|
"reg_lambda": {
|
|
"domain": tune.loguniform(lower=1 / 1024, upper=1024),
|
|
"init_value": 1.0,
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def size(cls, config):
|
|
return LGBMEstimator.size(config)
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 1.6
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
max_depth = params["max_depth"] = params.get("max_depth", 0)
|
|
if max_depth == 0:
|
|
params["grow_policy"] = params.get("grow_policy", "lossguide")
|
|
params["tree_method"] = params.get("tree_method", "hist")
|
|
# params["booster"] = params.get("booster", "gbtree")
|
|
params["use_label_encoder"] = params.get("use_label_encoder", False)
|
|
if "n_jobs" in config:
|
|
params["nthread"] = params.pop("n_jobs")
|
|
return params
|
|
|
|
def __init__(
|
|
self,
|
|
task="regression",
|
|
**config,
|
|
):
|
|
super().__init__(task, **config)
|
|
self.params["verbosity"] = 0
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
import xgboost as xgb
|
|
|
|
start_time = time.time()
|
|
deadline = start_time + budget if budget else np.inf
|
|
if issparse(X_train):
|
|
self.params["tree_method"] = "auto"
|
|
else:
|
|
X_train = self._preprocess(X_train)
|
|
if "sample_weight" in kwargs:
|
|
dtrain = xgb.DMatrix(X_train, label=y_train, weight=kwargs["sample_weight"])
|
|
else:
|
|
dtrain = xgb.DMatrix(X_train, label=y_train)
|
|
|
|
objective = self.params.get("objective")
|
|
if isinstance(objective, str):
|
|
obj = None
|
|
else:
|
|
obj = objective
|
|
if "objective" in self.params:
|
|
del self.params["objective"]
|
|
_n_estimators = self.params.pop("n_estimators")
|
|
callbacks = XGBoostEstimator._callbacks(start_time, deadline)
|
|
if callbacks:
|
|
self._model = xgb.train(
|
|
self.params,
|
|
dtrain,
|
|
_n_estimators,
|
|
obj=obj,
|
|
callbacks=callbacks,
|
|
)
|
|
self.params["n_estimators"] = self._model.best_iteration + 1
|
|
else:
|
|
self._model = xgb.train(self.params, dtrain, _n_estimators, obj=obj)
|
|
self.params["n_estimators"] = _n_estimators
|
|
self.params["objective"] = objective
|
|
del dtrain
|
|
train_time = time.time() - start_time
|
|
return train_time
|
|
|
|
def predict(self, X, **kwargs):
|
|
import xgboost as xgb
|
|
|
|
if not issparse(X):
|
|
X = self._preprocess(X)
|
|
dtest = xgb.DMatrix(X)
|
|
return super().predict(dtest)
|
|
|
|
@classmethod
|
|
def _callbacks(cls, start_time, deadline):
|
|
try:
|
|
from xgboost.callback import TrainingCallback
|
|
except ImportError: # for xgboost<1.3
|
|
return None
|
|
|
|
class ResourceLimit(TrainingCallback):
|
|
def after_iteration(self, model, epoch, evals_log) -> bool:
|
|
now = time.time()
|
|
if epoch == 0:
|
|
self._time_per_iter = now - start_time
|
|
if now + self._time_per_iter > deadline:
|
|
return True
|
|
if psutil is not None:
|
|
mem = psutil.virtual_memory()
|
|
if mem.available / mem.total < FREE_MEM_RATIO:
|
|
return True
|
|
return False
|
|
|
|
return [ResourceLimit()]
|
|
|
|
|
|
class XGBoostSklearnEstimator(SKLearnEstimator, LGBMEstimator):
|
|
"""The class for tuning XGBoost with unlimited depth, using sklearn API."""
|
|
|
|
DEFAULT_ITER = 10
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, **params):
|
|
space = XGBoostEstimator.search_space(data_size)
|
|
space.pop("max_depth")
|
|
return space
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return XGBoostEstimator.cost_relative2lgbm()
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
max_depth = params["max_depth"] = params.get("max_depth", 0)
|
|
if max_depth == 0:
|
|
params["grow_policy"] = params.get("grow_policy", "lossguide")
|
|
params["tree_method"] = params.get("tree_method", "hist")
|
|
params["use_label_encoder"] = params.get("use_label_encoder", False)
|
|
return params
|
|
|
|
def __init__(
|
|
self,
|
|
task="binary",
|
|
**config,
|
|
):
|
|
super().__init__(task, **config)
|
|
del self.params["verbose"]
|
|
self.params["verbosity"] = 0
|
|
import xgboost as xgb
|
|
|
|
self.estimator_class = xgb.XGBRegressor
|
|
if "rank" == task:
|
|
self.estimator_class = xgb.XGBRanker
|
|
elif task in CLASSIFICATION:
|
|
self.estimator_class = xgb.XGBClassifier
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
if issparse(X_train):
|
|
self.params["tree_method"] = "auto"
|
|
if kwargs.get("gpu_per_trial"):
|
|
self.params["tree_method"] = "gpu_hist"
|
|
kwargs.pop("gpu_per_trial")
|
|
return super().fit(X_train, y_train, budget, **kwargs)
|
|
|
|
def _callbacks(self, start_time, deadline) -> List[Callable]:
|
|
return XGBoostEstimator._callbacks(start_time, deadline)
|
|
|
|
|
|
class XGBoostLimitDepthEstimator(XGBoostSklearnEstimator):
|
|
"""The class for tuning XGBoost with limited depth, using sklearn API."""
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, **params):
|
|
space = XGBoostEstimator.search_space(data_size)
|
|
space.pop("max_leaves")
|
|
upper = max(6, int(np.log2(data_size[0])))
|
|
space["max_depth"] = {
|
|
"domain": tune.randint(lower=1, upper=min(upper, 16)),
|
|
"init_value": 6,
|
|
"low_cost_init_value": 1,
|
|
}
|
|
space["learning_rate"]["init_value"] = 0.3
|
|
space["n_estimators"]["init_value"] = 10
|
|
return space
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 64
|
|
|
|
|
|
class RandomForestEstimator(SKLearnEstimator, LGBMEstimator):
|
|
"""The class for tuning Random Forest."""
|
|
|
|
HAS_CALLBACK = False
|
|
nrows = 101
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, task, **params):
|
|
RandomForestEstimator.nrows = int(data_size[0])
|
|
upper = min(2048, RandomForestEstimator.nrows)
|
|
init = 1 / np.sqrt(data_size[1]) if task in CLASSIFICATION else 1
|
|
lower = min(0.1, init)
|
|
space = {
|
|
"n_estimators": {
|
|
"domain": tune.lograndint(lower=4, upper=max(5, upper)),
|
|
"init_value": 4,
|
|
"low_cost_init_value": 4,
|
|
},
|
|
"max_features": {
|
|
"domain": tune.loguniform(lower=lower, upper=1.0),
|
|
"init_value": init,
|
|
},
|
|
"max_leaves": {
|
|
"domain": tune.lograndint(
|
|
lower=4,
|
|
upper=max(5, min(32768, RandomForestEstimator.nrows >> 1)), #
|
|
),
|
|
"init_value": 4,
|
|
"low_cost_init_value": 4,
|
|
},
|
|
}
|
|
if task in CLASSIFICATION:
|
|
space["criterion"] = {
|
|
"domain": tune.choice(["gini", "entropy"]),
|
|
# "init_value": "gini",
|
|
}
|
|
return space
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 2
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
if "max_leaves" in params:
|
|
params["max_leaf_nodes"] = params.get(
|
|
"max_leaf_nodes", params.pop("max_leaves")
|
|
)
|
|
if self._task not in CLASSIFICATION and "criterion" in config:
|
|
params.pop("criterion")
|
|
return params
|
|
|
|
def __init__(
|
|
self,
|
|
task="binary",
|
|
**params,
|
|
):
|
|
super().__init__(task, **params)
|
|
self.params["verbose"] = 0
|
|
self.estimator_class = RandomForestRegressor
|
|
if task in CLASSIFICATION:
|
|
self.estimator_class = RandomForestClassifier
|
|
|
|
|
|
class ExtraTreesEstimator(RandomForestEstimator):
|
|
"""The class for tuning Extra Trees."""
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 1.9
|
|
|
|
def __init__(self, task="binary", **params):
|
|
super().__init__(task, **params)
|
|
if "regression" in task:
|
|
self.estimator_class = ExtraTreesRegressor
|
|
else:
|
|
self.estimator_class = ExtraTreesClassifier
|
|
|
|
|
|
class LRL1Classifier(SKLearnEstimator):
|
|
"""The class for tuning Logistic Regression with L1 regularization."""
|
|
|
|
@classmethod
|
|
def search_space(cls, **params):
|
|
return {
|
|
"C": {
|
|
"domain": tune.loguniform(lower=0.03125, upper=32768.0),
|
|
"init_value": 1.0,
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 160
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
params["tol"] = params.get("tol", 0.0001)
|
|
params["solver"] = params.get("solver", "saga")
|
|
params["penalty"] = params.get("penalty", "l1")
|
|
return params
|
|
|
|
def __init__(self, task="binary", **config):
|
|
super().__init__(task, **config)
|
|
assert task in CLASSIFICATION, "LogisticRegression for classification task only"
|
|
self.estimator_class = LogisticRegression
|
|
|
|
|
|
class LRL2Classifier(SKLearnEstimator):
|
|
"""The class for tuning Logistic Regression with L2 regularization."""
|
|
|
|
limit_resource = True
|
|
|
|
@classmethod
|
|
def search_space(cls, **params):
|
|
return LRL1Classifier.search_space(**params)
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 25
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
params["tol"] = params.get("tol", 0.0001)
|
|
params["solver"] = params.get("solver", "lbfgs")
|
|
params["penalty"] = params.get("penalty", "l2")
|
|
return params
|
|
|
|
def __init__(self, task="binary", **config):
|
|
super().__init__(task, **config)
|
|
assert task in CLASSIFICATION, "LogisticRegression for classification task only"
|
|
self.estimator_class = LogisticRegression
|
|
|
|
|
|
class CatBoostEstimator(BaseEstimator):
|
|
"""The class for tuning CatBoost."""
|
|
|
|
ITER_HP = "n_estimators"
|
|
DEFAULT_ITER = 1000
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, **params):
|
|
upper = max(min(round(1500000 / data_size[0]), 150), 12)
|
|
return {
|
|
"early_stopping_rounds": {
|
|
"domain": tune.lograndint(lower=10, upper=upper),
|
|
"init_value": 10,
|
|
"low_cost_init_value": 10,
|
|
},
|
|
"learning_rate": {
|
|
"domain": tune.loguniform(lower=0.005, upper=0.2),
|
|
"init_value": 0.1,
|
|
},
|
|
"n_estimators": {
|
|
"domain": 8192,
|
|
"init_value": 8192,
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def size(cls, config):
|
|
n_estimators = config.get("n_estimators", 8192)
|
|
max_leaves = 64
|
|
return (max_leaves * 3 + (max_leaves - 1) * 4 + 1.0) * n_estimators * 8
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 15
|
|
|
|
def _preprocess(self, X):
|
|
if isinstance(X, DataFrame):
|
|
cat_columns = X.select_dtypes(include=["category"]).columns
|
|
if not cat_columns.empty:
|
|
X = X.copy()
|
|
X[cat_columns] = X[cat_columns].apply(
|
|
lambda x: x.cat.rename_categories(
|
|
[
|
|
str(c) if isinstance(c, float) else c
|
|
for c in x.cat.categories
|
|
]
|
|
)
|
|
)
|
|
elif isinstance(X, np.ndarray) and X.dtype.kind not in "buif":
|
|
# numpy array is not of numeric dtype
|
|
X = DataFrame(X)
|
|
for col in X.columns:
|
|
if isinstance(X[col][0], str):
|
|
X[col] = X[col].astype("category").cat.codes
|
|
X = X.to_numpy()
|
|
return X
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
params["n_estimators"] = params.get("n_estimators", 8192)
|
|
if "n_jobs" in params:
|
|
params["thread_count"] = params.pop("n_jobs")
|
|
return params
|
|
|
|
def __init__(
|
|
self,
|
|
task="binary",
|
|
**config,
|
|
):
|
|
super().__init__(task, **config)
|
|
self.params.update(
|
|
{
|
|
"verbose": config.get("verbose", False),
|
|
"random_seed": config.get("random_seed", 10242048),
|
|
}
|
|
)
|
|
from catboost import CatBoostRegressor
|
|
|
|
self.estimator_class = CatBoostRegressor
|
|
if task in CLASSIFICATION:
|
|
from catboost import CatBoostClassifier
|
|
|
|
self.estimator_class = CatBoostClassifier
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
start_time = time.time()
|
|
deadline = start_time + budget if budget else np.inf
|
|
train_dir = f"catboost_{str(start_time)}"
|
|
X_train = self._preprocess(X_train)
|
|
if isinstance(X_train, DataFrame):
|
|
cat_features = list(X_train.select_dtypes(include="category").columns)
|
|
else:
|
|
cat_features = []
|
|
n = max(int(len(y_train) * 0.9), len(y_train) - 1000)
|
|
X_tr, y_tr = X_train[:n], y_train[:n]
|
|
if "sample_weight" in kwargs:
|
|
weight = kwargs["sample_weight"]
|
|
if weight is not None:
|
|
kwargs["sample_weight"] = weight[:n]
|
|
else:
|
|
weight = None
|
|
from catboost import Pool, __version__
|
|
|
|
model = self.estimator_class(train_dir=train_dir, **self.params)
|
|
if __version__ >= "0.26":
|
|
model.fit(
|
|
X_tr,
|
|
y_tr,
|
|
cat_features=cat_features,
|
|
eval_set=Pool(
|
|
data=X_train[n:], label=y_train[n:], cat_features=cat_features
|
|
),
|
|
callbacks=CatBoostEstimator._callbacks(start_time, deadline),
|
|
**kwargs,
|
|
)
|
|
else:
|
|
model.fit(
|
|
X_tr,
|
|
y_tr,
|
|
cat_features=cat_features,
|
|
eval_set=Pool(
|
|
data=X_train[n:], label=y_train[n:], cat_features=cat_features
|
|
),
|
|
**kwargs,
|
|
)
|
|
shutil.rmtree(train_dir, ignore_errors=True)
|
|
if weight is not None:
|
|
kwargs["sample_weight"] = weight
|
|
self._model = model
|
|
self.params[self.ITER_HP] = self._model.tree_count_
|
|
train_time = time.time() - start_time
|
|
return train_time
|
|
|
|
@classmethod
|
|
def _callbacks(cls, start_time, deadline):
|
|
class ResourceLimit:
|
|
def after_iteration(self, info) -> bool:
|
|
now = time.time()
|
|
if info.iteration == 1:
|
|
self._time_per_iter = now - start_time
|
|
if now + self._time_per_iter > deadline:
|
|
return False
|
|
if psutil is not None:
|
|
mem = psutil.virtual_memory()
|
|
if mem.available / mem.total < FREE_MEM_RATIO:
|
|
return False
|
|
return True # can continue
|
|
|
|
return [ResourceLimit()]
|
|
|
|
|
|
class KNeighborsEstimator(BaseEstimator):
|
|
@classmethod
|
|
def search_space(cls, data_size, **params):
|
|
upper = min(512, int(data_size[0] / 2))
|
|
return {
|
|
"n_neighbors": {
|
|
"domain": tune.lograndint(lower=1, upper=max(2, upper)),
|
|
"init_value": 5,
|
|
"low_cost_init_value": 1,
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def cost_relative2lgbm(cls):
|
|
return 30
|
|
|
|
def config2params(self, config: dict) -> dict:
|
|
params = super().config2params(config)
|
|
params["weights"] = params.get("weights", "distance")
|
|
return params
|
|
|
|
def __init__(self, task="binary", **config):
|
|
super().__init__(task, **config)
|
|
if task in CLASSIFICATION:
|
|
from sklearn.neighbors import KNeighborsClassifier
|
|
|
|
self.estimator_class = KNeighborsClassifier
|
|
else:
|
|
from sklearn.neighbors import KNeighborsRegressor
|
|
|
|
self.estimator_class = KNeighborsRegressor
|
|
|
|
def _preprocess(self, X):
|
|
if isinstance(X, DataFrame):
|
|
cat_columns = X.select_dtypes(["category"]).columns
|
|
if X.shape[1] == len(cat_columns):
|
|
raise ValueError("kneighbor requires at least one numeric feature")
|
|
X = X.drop(cat_columns, axis=1)
|
|
elif isinstance(X, np.ndarray) and X.dtype.kind not in "buif":
|
|
# drop categocial columns if any
|
|
X = DataFrame(X)
|
|
cat_columns = []
|
|
for col in X.columns:
|
|
if isinstance(X[col][0], str):
|
|
cat_columns.append(col)
|
|
X = X.drop(cat_columns, axis=1)
|
|
X = X.to_numpy()
|
|
return X
|
|
|
|
|
|
class Prophet(SKLearnEstimator):
|
|
"""The class for tuning Prophet."""
|
|
|
|
@classmethod
|
|
def search_space(cls, **params):
|
|
space = {
|
|
"changepoint_prior_scale": {
|
|
"domain": tune.loguniform(lower=0.001, upper=0.05),
|
|
"init_value": 0.05,
|
|
"low_cost_init_value": 0.001,
|
|
},
|
|
"seasonality_prior_scale": {
|
|
"domain": tune.loguniform(lower=0.01, upper=10),
|
|
"init_value": 10,
|
|
},
|
|
"holidays_prior_scale": {
|
|
"domain": tune.loguniform(lower=0.01, upper=10),
|
|
"init_value": 10,
|
|
},
|
|
"seasonality_mode": {
|
|
"domain": tune.choice(["additive", "multiplicative"]),
|
|
"init_value": "multiplicative",
|
|
},
|
|
}
|
|
return space
|
|
|
|
def __init__(self, task="ts_forecast", n_jobs=1, **params):
|
|
super().__init__(task, **params)
|
|
|
|
def _join(self, X_train, y_train):
|
|
assert TS_TIMESTAMP_COL in X_train, (
|
|
"Dataframe for training ts_forecast model must have column"
|
|
f' "{TS_TIMESTAMP_COL}" with the dates in X_train.'
|
|
)
|
|
y_train = DataFrame(y_train, columns=[TS_VALUE_COL])
|
|
train_df = X_train.join(y_train)
|
|
return train_df
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
from prophet import Prophet
|
|
|
|
current_time = time.time()
|
|
train_df = self._join(X_train, y_train)
|
|
train_df = self._preprocess(train_df)
|
|
cols = list(train_df)
|
|
cols.remove(TS_TIMESTAMP_COL)
|
|
cols.remove(TS_VALUE_COL)
|
|
logging.getLogger("prophet").setLevel(logging.WARNING)
|
|
model = Prophet(**self.params)
|
|
for regressor in cols:
|
|
model.add_regressor(regressor)
|
|
with suppress_stdout_stderr():
|
|
model.fit(train_df)
|
|
train_time = time.time() - current_time
|
|
self._model = model
|
|
return train_time
|
|
|
|
def predict(self, X, **kwargs):
|
|
if isinstance(X, int):
|
|
raise ValueError(
|
|
"predict() with steps is only supported for arima/sarimax."
|
|
" For Prophet, pass a dataframe with the first column containing"
|
|
" the timestamp values."
|
|
)
|
|
if self._model is not None:
|
|
X = self._preprocess(X)
|
|
forecast = self._model.predict(X)
|
|
return forecast["yhat"]
|
|
else:
|
|
logger.warning(
|
|
"Estimator is not fit yet. Please run fit() before predict()."
|
|
)
|
|
return np.ones(X.shape[0])
|
|
|
|
def score(self, X_val: DataFrame, y_val: Series, **kwargs):
|
|
from sklearn.metrics import r2_score
|
|
from .ml import metric_loss_score
|
|
|
|
y_pred = self.predict(X_val)
|
|
self._metric = kwargs.get("metric", None)
|
|
if self._metric:
|
|
return metric_loss_score(self._metric, y_pred, y_val)
|
|
else:
|
|
return r2_score(y_pred, y_val)
|
|
|
|
|
|
class ARIMA(Prophet):
|
|
"""The class for tuning ARIMA."""
|
|
|
|
@classmethod
|
|
def search_space(cls, **params):
|
|
space = {
|
|
"p": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 2,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"d": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 2,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"q": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 1,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
}
|
|
return space
|
|
|
|
def _join(self, X_train, y_train):
|
|
train_df = super()._join(X_train, y_train)
|
|
train_df.index = to_datetime(train_df[TS_TIMESTAMP_COL])
|
|
train_df = train_df.drop(TS_TIMESTAMP_COL, axis=1)
|
|
return train_df
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
import warnings
|
|
|
|
warnings.filterwarnings("ignore")
|
|
from statsmodels.tsa.arima.model import ARIMA as ARIMA_estimator
|
|
|
|
current_time = time.time()
|
|
train_df = self._join(X_train, y_train)
|
|
train_df = self._preprocess(train_df)
|
|
regressors = list(train_df)
|
|
regressors.remove(TS_VALUE_COL)
|
|
if regressors:
|
|
model = ARIMA_estimator(
|
|
train_df[[TS_VALUE_COL]],
|
|
exog=train_df[regressors],
|
|
order=(self.params["p"], self.params["d"], self.params["q"]),
|
|
enforce_stationarity=False,
|
|
enforce_invertibility=False,
|
|
)
|
|
else:
|
|
model = ARIMA_estimator(
|
|
train_df,
|
|
order=(self.params["p"], self.params["d"], self.params["q"]),
|
|
enforce_stationarity=False,
|
|
enforce_invertibility=False,
|
|
)
|
|
with suppress_stdout_stderr():
|
|
model = model.fit()
|
|
train_time = time.time() - current_time
|
|
self._model = model
|
|
return train_time
|
|
|
|
def predict(self, X, **kwargs):
|
|
if self._model is not None:
|
|
if isinstance(X, int):
|
|
forecast = self._model.forecast(steps=X)
|
|
elif isinstance(X, DataFrame):
|
|
start = X[TS_TIMESTAMP_COL].iloc[0]
|
|
end = X[TS_TIMESTAMP_COL].iloc[-1]
|
|
if len(X.columns) > 1:
|
|
X = self._preprocess(X.drop(columns=TS_TIMESTAMP_COL))
|
|
regressors = list(X)
|
|
forecast = self._model.predict(
|
|
start=start, end=end, exog=X[regressors]
|
|
)
|
|
else:
|
|
forecast = self._model.predict(start=start, end=end)
|
|
else:
|
|
raise ValueError(
|
|
"X needs to be either a pandas Dataframe with dates as the first column"
|
|
" or an int number of periods for predict()."
|
|
)
|
|
return forecast
|
|
else:
|
|
return np.ones(X if isinstance(X, int) else X.shape[0])
|
|
|
|
|
|
class SARIMAX(ARIMA):
|
|
"""The class for tuning SARIMA."""
|
|
|
|
@classmethod
|
|
def search_space(cls, **params):
|
|
space = {
|
|
"p": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 2,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"d": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 2,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"q": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 1,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"P": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 1,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"D": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 1,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"Q": {
|
|
"domain": tune.quniform(lower=0, upper=10, q=1),
|
|
"init_value": 1,
|
|
"low_cost_init_value": 0,
|
|
},
|
|
"s": {
|
|
"domain": tune.choice([1, 4, 6, 12]),
|
|
"init_value": 12,
|
|
},
|
|
}
|
|
return space
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
import warnings
|
|
|
|
warnings.filterwarnings("ignore")
|
|
from statsmodels.tsa.statespace.sarimax import SARIMAX as SARIMAX_estimator
|
|
|
|
current_time = time.time()
|
|
train_df = self._join(X_train, y_train)
|
|
train_df = self._preprocess(train_df)
|
|
regressors = list(train_df)
|
|
regressors.remove(TS_VALUE_COL)
|
|
if regressors:
|
|
model = SARIMAX_estimator(
|
|
train_df[[TS_VALUE_COL]],
|
|
exog=train_df[regressors],
|
|
order=(self.params["p"], self.params["d"], self.params["q"]),
|
|
seasonality_order=(
|
|
self.params["P"],
|
|
self.params["D"],
|
|
self.params["Q"],
|
|
self.params["s"],
|
|
),
|
|
enforce_stationarity=False,
|
|
enforce_invertibility=False,
|
|
)
|
|
else:
|
|
model = SARIMAX_estimator(
|
|
train_df,
|
|
order=(self.params["p"], self.params["d"], self.params["q"]),
|
|
seasonality_order=(
|
|
self.params["P"],
|
|
self.params["D"],
|
|
self.params["Q"],
|
|
self.params["s"],
|
|
),
|
|
enforce_stationarity=False,
|
|
enforce_invertibility=False,
|
|
)
|
|
with suppress_stdout_stderr():
|
|
model = model.fit()
|
|
train_time = time.time() - current_time
|
|
self._model = model
|
|
return train_time
|
|
|
|
|
|
class TS_SKLearn(SKLearnEstimator):
|
|
"""The class for tuning SKLearn Regressors for time-series forecasting, using hcrystalball"""
|
|
|
|
base_class = SKLearnEstimator
|
|
|
|
@classmethod
|
|
def search_space(cls, data_size, pred_horizon, **params):
|
|
space = cls.base_class.search_space(data_size, **params)
|
|
space.update(
|
|
{
|
|
"optimize_for_horizon": {
|
|
"domain": tune.choice([True, False]),
|
|
"init_value": False,
|
|
"low_cost_init_value": False,
|
|
},
|
|
"lags": {
|
|
"domain": tune.randint(
|
|
lower=1, upper=max(2, int(np.sqrt(data_size[0])))
|
|
),
|
|
"init_value": 3,
|
|
},
|
|
}
|
|
)
|
|
return space
|
|
|
|
def __init__(self, task="ts_forecast", **params):
|
|
super().__init__(task, **params)
|
|
self.hcrystaball_model = None
|
|
self.ts_task = (
|
|
"regression" if task in TS_FORECASTREGRESSION else "classification"
|
|
)
|
|
|
|
def transform_X(self, X):
|
|
cols = list(X)
|
|
if len(cols) == 1:
|
|
ds_col = cols[0]
|
|
X = DataFrame(index=X[ds_col])
|
|
elif len(cols) > 1:
|
|
ds_col = cols[0]
|
|
exog_cols = cols[1:]
|
|
X = X[exog_cols].set_index(X[ds_col])
|
|
return X
|
|
|
|
def _fit(self, X_train, y_train, budget=None, **kwargs):
|
|
from hcrystalball.wrappers import get_sklearn_wrapper
|
|
|
|
X_train = self.transform_X(X_train)
|
|
X_train = self._preprocess(X_train)
|
|
params = self.params.copy()
|
|
lags = params.pop("lags")
|
|
optimize_for_horizon = params.pop("optimize_for_horizon")
|
|
estimator = self.base_class(task=self.ts_task, **params)
|
|
self.hcrystaball_model = get_sklearn_wrapper(estimator.estimator_class)
|
|
self.hcrystaball_model.lags = int(lags)
|
|
self.hcrystaball_model.fit(X_train, y_train)
|
|
if optimize_for_horizon:
|
|
# Direct Multi-step Forecast Strategy - fit a seperate model for each horizon
|
|
model_list = []
|
|
for i in range(1, kwargs["period"] + 1):
|
|
(
|
|
X_fit,
|
|
y_fit,
|
|
) = self.hcrystaball_model._transform_data_to_tsmodel_input_format(
|
|
X_train, y_train, i
|
|
)
|
|
self.hcrystaball_model.model.set_params(**estimator.params)
|
|
model = self.hcrystaball_model.model.fit(X_fit, y_fit)
|
|
model_list.append(model)
|
|
self._model = model_list
|
|
else:
|
|
(
|
|
X_fit,
|
|
y_fit,
|
|
) = self.hcrystaball_model._transform_data_to_tsmodel_input_format(
|
|
X_train, y_train, kwargs["period"]
|
|
)
|
|
self.hcrystaball_model.model.set_params(**estimator.params)
|
|
model = self.hcrystaball_model.model.fit(X_fit, y_fit)
|
|
self._model = model
|
|
|
|
def fit(self, X_train, y_train, budget=None, **kwargs):
|
|
current_time = time.time()
|
|
self._fit(X_train, y_train, budget=budget, **kwargs)
|
|
train_time = time.time() - current_time
|
|
return train_time
|
|
|
|
def predict(self, X, **kwargs):
|
|
if self._model is not None:
|
|
X = self.transform_X(X)
|
|
X = self._preprocess(X)
|
|
if isinstance(self._model, list):
|
|
assert len(self._model) == len(
|
|
X
|
|
), "Model is optimized for horizon, length of X must be equal to `period`."
|
|
preds = []
|
|
for i in range(1, len(self._model) + 1):
|
|
(
|
|
X_pred,
|
|
_,
|
|
) = self.hcrystaball_model._transform_data_to_tsmodel_input_format(
|
|
X.iloc[:i, :]
|
|
)
|
|
preds.append(self._model[i - 1].predict(X_pred)[-1])
|
|
forecast = DataFrame(
|
|
data=np.asarray(preds).reshape(-1, 1),
|
|
columns=[self.hcrystaball_model.name],
|
|
index=X.index,
|
|
)
|
|
else:
|
|
(
|
|
X_pred,
|
|
_,
|
|
) = self.hcrystaball_model._transform_data_to_tsmodel_input_format(X)
|
|
forecast = self._model.predict(X_pred)
|
|
return forecast
|
|
else:
|
|
logger.warning(
|
|
"Estimator is not fit yet. Please run fit() before predict()."
|
|
)
|
|
return np.ones(X.shape[0])
|
|
|
|
|
|
class LGBM_TS(TS_SKLearn):
|
|
"""The class for tuning LGBM Regressor for time-series forecasting"""
|
|
|
|
base_class = LGBMEstimator
|
|
|
|
|
|
class XGBoost_TS(TS_SKLearn):
|
|
"""The class for tuning XGBoost Regressor for time-series forecasting"""
|
|
|
|
base_class = XGBoostSklearnEstimator
|
|
|
|
|
|
# catboost regressor is invalid because it has a `name` parameter, making it incompatible with hcrystalball
|
|
# class CatBoost_TS_Regressor(TS_Regressor):
|
|
# base_class = CatBoostEstimator
|
|
|
|
|
|
class RF_TS(TS_SKLearn):
|
|
"""The class for tuning Random Forest Regressor for time-series forecasting"""
|
|
|
|
base_class = RandomForestEstimator
|
|
|
|
|
|
class ExtraTrees_TS(TS_SKLearn):
|
|
"""The class for tuning Extra Trees Regressor for time-series forecasting"""
|
|
|
|
base_class = ExtraTreesEstimator
|
|
|
|
|
|
class XGBoostLimitDepth_TS(TS_SKLearn):
|
|
"""The class for tuning XGBoost Regressor with unlimited depth for time-series forecasting"""
|
|
|
|
base_class = XGBoostLimitDepthEstimator
|
|
|
|
|
|
class suppress_stdout_stderr(object):
|
|
def __init__(self):
|
|
# Open a pair of null files
|
|
self.null_fds = [os.open(os.devnull, os.O_RDWR) for x in range(2)]
|
|
# Save the actual stdout (1) and stderr (2) file descriptors.
|
|
self.save_fds = (os.dup(1), os.dup(2))
|
|
|
|
def __enter__(self):
|
|
# Assign the null pointers to stdout and stderr.
|
|
os.dup2(self.null_fds[0], 1)
|
|
os.dup2(self.null_fds[1], 2)
|
|
|
|
def __exit__(self, *_):
|
|
# Re-assign the real stdout/stderr back to (1) and (2)
|
|
os.dup2(self.save_fds[0], 1)
|
|
os.dup2(self.save_fds[1], 2)
|
|
# Close the null files
|
|
os.close(self.null_fds[0])
|
|
os.close(self.null_fds[1])
|